[autofix.ci] apply automated fixes

api/.importlinter

@@ -171,6 +171,48 @@ forbidden_modules =
     sqlalchemy
     werkzeug
 
+[importlinter:contract:dataset-application-boundary]
+name = Dataset application services and contracts are framework and persistence neutral
+type = forbidden
+source_modules =
+    services.data_source.binding_application_service
+    services.data_source.notion_import_application_service
+    services.entities.data_source.credential
+    services.entities.data_source.notion_import
+    services.entities.knowledge_entities.indexing_estimate
+    services.knowledge.application
+forbidden_modules =
+    configs
+    controllers
+    extensions
+    flask
+    models
+    repositories
+    sqlalchemy
+    tasks
+    werkzeug
+    services.dataset_service
+    services.datasource_provider_service
+    core.datasource.datasource_manager
+    core.indexing_runner
+    core.rag.extractor.notion_extractor
+
+[importlinter:contract:dataset-reference-boundary]
+name = Dataset owner references do not depend on web application layers
+type = forbidden
+source_modules =
+    core.rag.entities.dataset_reference
+forbidden_modules =
+    controllers
+    extensions
+    flask
+    models
+    repositories
+    services
+    sqlalchemy
+    tasks
+    werkzeug
+
 [importlinter:contract:init-validation-service-boundary]
 name = Initialization validation application service is framework and persistence neutral
 type = forbidden
@@ -310,7 +352,7 @@ forbidden_modules =
 name = Data source OAuth application service is framework and persistence neutral
 type = forbidden
 source_modules =
-    services.data_source_oauth_service
+    services.data_source.oauth_service
 forbidden_modules =
     configs
     controllers

api/controllers/console/auth/data_source_bearer_auth.py

@@ -26,7 +26,7 @@ from services.auth.errors import (
     InvalidDataSourceApiKeyAuthCredentialsError,
     UnsupportedDataSourceApiKeyAuthProviderError,
 )
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthCredentials,
 )
@@ -79,7 +79,7 @@ class ApiKeyAuthDataSource(Resource):
     @console_ns.response(HTTPStatus.OK, "Success", console_ns.models[ApiKeyAuthDataSourceListResponse.__name__])
     @console_account_admission()
     def get(self, request_context: RequestContext):
-        bindings = application_services().data_source_api_key_auth.list_bindings(request_context)
+        bindings = application_services().data_sources.api_key_auth.list_bindings(request_context)
         return dump_response(
             ApiKeyAuthDataSourceListResponse,
             {
@@ -121,7 +121,7 @@ class ApiKeyAuthDataSourceBinding(Resource):
             ),
         )
         try:
-            application_services().data_source_api_key_auth.create_binding(request_context, command)
+            application_services().data_sources.api_key_auth.create_binding(request_context, command)
         except UnsupportedDataSourceApiKeyAuthProviderError as exc:
             raise DataSourceApiKeyAuthProviderNotSupportedError() from exc
         except InvalidDataSourceApiKeyAuthCredentialsError as exc:
@@ -143,5 +143,5 @@ class ApiKeyAuthDataSourceBindingDelete(Resource):
         rbac_resource_required=False,
     )
     def delete(self, request_context: RequestContext, binding_id: UUID):
-        application_services().data_source_api_key_auth.delete_binding(request_context, str(binding_id))
+        application_services().data_sources.api_key_auth.delete_binding(request_context, str(binding_id))
         return "", HTTPStatus.NO_CONTENT

api/controllers/console/auth/data_source_oauth.py

@@ -23,13 +23,13 @@ from fields.base import ResponseModel
 from libs.helper import dump_response
 from machinery.context import RequestContext
 from models.account import TenantAccountRole
-from services.data_source_oauth_service import (
+from services.data_source.oauth_service import (
     DataSourceOAuthConfigurationError,
     DataSourceOAuthError,
     InvalidDataSourceOAuthCodeError,
     InvalidDataSourceOAuthProviderError,
 )
-from services.entities.data_source_oauth_entities import DataSourceOAuthCallback
+from services.entities.data_source.oauth import DataSourceOAuthCallback
 
 from .. import console_ns
 
@@ -117,7 +117,7 @@ class OAuthDataSource(Resource):
     )
     def get(self, request_context: RequestContext, provider: str):
         try:
-            service = application_services().resolve_data_source_oauth(provider)
+            service = application_services().data_sources.resolve_oauth(provider)
             authorization = service.start_authorization(request_context)
         except InvalidDataSourceOAuthProviderError:
             return _invalid_provider_response()
@@ -142,7 +142,7 @@ class OAuthDataSourceCallback(Resource):
     def get(self, provider: str):
         query = query_params_from_request(OAuthDataSourceCallbackQuery)
         try:
-            service = application_services().resolve_data_source_oauth(provider)
+            service = application_services().data_sources.resolve_oauth(provider)
             callback = service.complete_callback(
                 code=query.code,
                 error=query.error,
@@ -173,7 +173,7 @@ class OAuthDataSourceBinding(Resource):
         if not query.code:
             return {"error": "Invalid code"}, HTTPStatus.BAD_REQUEST
         try:
-            service = application_services().resolve_data_source_oauth(provider)
+            service = application_services().data_sources.resolve_oauth(provider)
             service.bind(request_context, code=query.code)
         except InvalidDataSourceOAuthProviderError:
             return _invalid_provider_response()
@@ -199,7 +199,7 @@ class OAuthDataSourceSync(Resource):
     @console_account_admission()
     def get(self, request_context: RequestContext, provider: str, binding_id: UUID):
         try:
-            service = application_services().resolve_data_source_oauth(provider)
+            service = application_services().data_sources.resolve_oauth(provider)
             service.sync(
                 request_context,
                 binding_id=str(binding_id),

api/controllers/console/datasets/datasets_segments.py

@@ -44,6 +44,7 @@ from controllers.console.wraps import (
 )
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
+from core.rag.entities.dataset_reference import SegmentRef
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from extensions.ext_redis import redis_client
 from fields.base import ResponseModel
@@ -63,7 +64,7 @@ from libs.pagination import paginate_query
 from models import Account
 from models.dataset import Dataset, Document, DocumentSegment
 from models.model import UploadFile
-from services.dataset_ref_service import DatasetRefService, SegmentRef
+from services.dataset_ref_service import DatasetRefService
 from services.dataset_service import DatasetService, DocumentService, SegmentService
 from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs
 from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError

api/controllers/service_api/dataset/segment.py

@@ -26,6 +26,7 @@ from controllers.service_api.wraps import (
 )
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
+from core.rag.entities.dataset_reference import SegmentRef
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from fields.base import ResponseModel
 from fields.segment_fields import (
@@ -40,7 +41,7 @@ from graphon.model_runtime.entities.model_entities import ModelType
 from libs.helper import dump_response
 from libs.login import current_account_with_tenant
 from models.dataset import Dataset, Document, DocumentSegment
-from services.dataset_ref_service import DatasetRefService, SegmentRef
+from services.dataset_ref_service import DatasetRefService
 from services.dataset_service import DatasetService, DocumentService, SegmentService
 from services.entities.knowledge_entities.knowledge_entities import SegmentUpdateArgs
 from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError

api/core/provider_manager.py

@@ -1157,6 +1157,7 @@ class ProviderManager:
             if user is not None:
                 stmt = CredentialPermissionService.apply_visibility_filter(
                     stmt,
+                    tenant_id=tenant_id,
                     model_id_column=ProviderCredential.id,
                     model_user_id_column=ProviderCredential.user_id,
                     model_visibility_column=ProviderCredential.visibility,

api/core/rag/entities/__init__.py

@@ -3,7 +3,13 @@ from core.rag.entities.context_entities import DocumentContext
 from core.rag.entities.event import DatasourceCompletedEvent, DatasourceErrorEvent, DatasourceProcessingEvent
 from core.rag.entities.index_entities import EconomySetting, EmbeddingSetting, IndexMethod
 from core.rag.entities.metadata_entities import Condition, MetadataFilteringCondition, SupportedComparisonOperator
-from core.rag.entities.processing_entities import ParentMode, PreProcessingRule, Rule, Segmentation
+from core.rag.entities.processing_entities import (
+    ParentMode,
+    PreProcessingRule,
+    PreProcessingRuleKey,
+    Rule,
+    Segmentation,
+)
 from core.rag.entities.retrieval_settings import (
     KeywordSetting,
     RerankingModelConfig,
@@ -24,6 +30,7 @@ __all__ = [
     "MetadataFilteringCondition",
     "ParentMode",
     "PreProcessingRule",
+    "PreProcessingRuleKey",
     "RerankingModelConfig",
     "RetrievalSourceMetadata",
     "Rule",

api/core/rag/entities/dataset_reference.py

@@ -0,0 +1,57 @@
+"""Framework-neutral owner chains for knowledge resources."""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+
+class DatasetRef(NamedTuple):
+    """Identify a dataset inside its owning workspace."""
+
+    tenant_id: str
+    dataset_id: str
+
+    def document(self, document_id: str) -> DocumentRef:
+        """Bind a document identifier to this trusted dataset owner."""
+
+        return DocumentRef(dataset=self, document_id=document_id)
+
+
+class DocumentRef(NamedTuple):
+    """Identify a document through its complete dataset owner chain."""
+
+    dataset: DatasetRef
+    document_id: str
+
+    def segment(self, segment_id: str) -> SegmentRef:
+        """Bind a segment identifier to this trusted document owner."""
+
+        return SegmentRef(document=self, segment_id=segment_id)
+
+    def matches(self, *, tenant_id: str, dataset_id: str, document_id: str) -> bool:
+        """Return whether raw resource identifiers belong to this reference."""
+
+        return (
+            tenant_id == self.dataset.tenant_id
+            and dataset_id == self.dataset.dataset_id
+            and document_id == self.document_id
+        )
+
+
+class SegmentRef(NamedTuple):
+    """Identify a segment through its complete document owner chain."""
+
+    document: DocumentRef
+    segment_id: str
+
+    def matches(self, *, tenant_id: str, dataset_id: str, document_id: str, segment_id: str) -> bool:
+        """Return whether raw resource identifiers belong to this reference."""
+
+        return (
+            self.document.matches(
+                tenant_id=tenant_id,
+                dataset_id=dataset_id,
+                document_id=document_id,
+            )
+            and segment_id == self.segment_id
+        )

api/core/rag/entities/processing_entities.py

@@ -1,7 +1,7 @@
 from enum import StrEnum
-from typing import Annotated, Literal
+from typing import Literal
 
-from pydantic import AliasChoices, BaseModel, Field, WithJsonSchema
+from pydantic import AliasChoices, BaseModel, Field
 
 
 class ParentMode(StrEnum):
@@ -9,19 +9,11 @@ class ParentMode(StrEnum):
     PARAGRAPH = "paragraph"
 
 
-PreProcessingRuleID = Annotated[
-    str,
-    WithJsonSchema(
-        {
-            "enum": ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"],
-            "type": "string",
-        }
-    ),
-]
+type PreProcessingRuleKey = Literal["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
 
 
 class PreProcessingRule(BaseModel):
-    id: PreProcessingRuleID = Field(description="Rule identifier.")
+    id: PreProcessingRuleKey = Field(description="Rule identifier.")
     enabled: bool = Field(description="Whether this preprocessing rule is enabled.")
 
 

api/core/rag/extractor/entity/extract_setting.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
 from models.dataset import Document
 from models.model import UploadFile
@@ -15,6 +15,7 @@ class NotionInfo(BaseModel):
     notion_page_type: str
     document: Document | None = None
     tenant_id: str
+    notion_access_token: str | None = Field(default=None, exclude=True, repr=False)
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 

api/core/rag/extractor/extract_processor.py

@@ -224,6 +224,7 @@ class ExtractProcessor:
                 notion_page_type=extract_setting.notion_info.notion_page_type,
                 document_model=extract_setting.notion_info.document,
                 tenant_id=extract_setting.notion_info.tenant_id,
+                notion_access_token=extract_setting.notion_info.notion_access_token,
                 credential_id=extract_setting.notion_info.credential_id,
             )
             return extractor.extract()

api/extensions/ext_application_services.py

@@ -16,6 +16,7 @@ from configs import dify_config
 from constants.dsl_version import CURRENT_APP_DSL_VERSION
 from core.db.session_factory import get_session_maker
 from core.helper.ssrf_proxy import ssrf_proxy
+from core.plugin.impl.datasource import PluginDatasourceManager
 from core.schemas.schema_manager import SchemaManager
 from enums import DeploymentEdition, WebAppAccessMode
 from extensions.ext_redis import RedisClientWrapper, redis_client
@@ -27,11 +28,14 @@ from repositories.account_integration_repository import SQLAlchemyAccountIntegra
 from repositories.account_repository import SQLAlchemyAccountRepository
 from repositories.app_definition_query_repository import AppDefinitionQueryRepository
 from repositories.app_site_command_repository import AppSiteCommandRepository
-from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
-from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
+from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
+from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository
+from repositories.data_source.oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
 from repositories.explore_banner_query_repository import ExploreBannerQueryRepository
 from repositories.factory import DifyAPIRepositoryFactory
 from repositories.installation_state_repository import InstallationStateRepository
+from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
+from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
 from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository
 from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository
 from repositories.step_by_step_tour_repository import SQLAlchemyStepByStepTourStateRepository
@@ -109,7 +113,16 @@ from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthS
 from services.billing_portal_service import BillingPortalService
 from services.billing_service import BillingService
 from services.compliance_download_service import ComplianceDownloadService
-from services.data_source_oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError
+from services.data_source.binding_application_service import DataSourceBindingApplicationService
+from services.data_source.credential_gateway import (
+    ActorAwareDatasourceCredentialGateway,
+    OAuthDatasourceCredentialRefresher,
+    PluginDatasourceCredentialCodec,
+)
+from services.data_source.notion_import_adapters import PluginNotionSourceGateway
+from services.data_source.notion_import_application_service import NotionImportApplicationService
+from services.data_source.notion_oauth_gateway import NotionDataSourceGateway
+from services.data_source.oauth_service import DataSourceOAuthService, InvalidDataSourceOAuthProviderError
 from services.enterprise.enterprise_service import EnterpriseService
 from services.errors.enterprise import EnterpriseServiceError
 from services.explore_banner_query_service import ExploreBannerQueryService
@@ -118,9 +131,14 @@ from services.feature_service_gateway import FeatureServiceGateway
 from services.file_service import FileService
 from services.init_validation_service import InitValidationService
 from services.inner_mail_service import InnerMailService
+from services.knowledge.adapters import CeleryDocumentSyncDispatcher, IndexingRunnerEstimateGateway
+from services.knowledge.application import (
+    DatasetAccessService,
+    DocumentSyncApplicationService,
+    IndexingEstimateApplicationService,
+)
 from services.notification_gateway import BillingNotificationGateway
 from services.notification_service import NotificationService
-from services.notion_data_source_gateway import NotionDataSourceGateway
 from services.oauth_server_service import OAUTH_ACCESS_TOKEN_EXPIRES_IN, OAuthServerService
 from services.partner_tenant_binding_service import PartnerTenantBindingService
 from services.recommended_app_catalog_gateway import (
@@ -195,6 +213,26 @@ class AccountServices:
     profile: AccountProfileService
 
 
+@dataclass(frozen=True, slots=True)
+class DataSourceServices:
+    api_key_auth: DataSourceApiKeyAuthService
+    oauth: Mapping[str, DataSourceOAuthService]
+    bindings: DataSourceBindingApplicationService
+    notion_imports: NotionImportApplicationService
+
+    def resolve_oauth(self, provider: str) -> DataSourceOAuthService:
+        service = self.oauth.get(provider)
+        if service is None:
+            raise InvalidDataSourceOAuthProviderError("Invalid provider")
+        return service
+
+
+@dataclass(frozen=True, slots=True)
+class KnowledgeServices:
+    document_sync: DocumentSyncApplicationService
+    indexing_estimates: IndexingEstimateApplicationService
+
+
 @dataclass(frozen=True, slots=True)
 class ApplicationServices:
     accounts: AccountServices
@@ -203,8 +241,8 @@ class ApplicationServices:
     app_sites: AppSiteService
     billing_portal: BillingPortalService
     compliance_downloads: ComplianceDownloadService
-    data_source_api_key_auth: DataSourceApiKeyAuthService
-    data_source_oauth: Mapping[str, DataSourceOAuthService]
+    data_sources: DataSourceServices
+    knowledge: KnowledgeServices
     webapp_access: WebAppAccessQueryService
     web_app_runtime: WebAppRuntimeQueryService
     explore_banner_queries: ExploreBannerQueryService
@@ -227,16 +265,10 @@ class ApplicationServices:
     tags: TagApplicationService
     workflow_statistics: WorkflowStatisticQueryService
 
-    def resolve_data_source_oauth(self, provider: str) -> DataSourceOAuthService:
-        service = self.data_source_oauth.get(provider)
-        if service is None:
-            raise InvalidDataSourceOAuthProviderError("Invalid provider")
-        return service
-
 
 def _build_data_source_oauth_services(
     *,
-    database_client: sessionmaker[Session],
+    bindings: SQLAlchemyDataSourceOAuthBindingRepository,
 ) -> Mapping[str, DataSourceOAuthService]:
     notion_data_source = NotionDataSourceGateway(
         client_id=dify_config.NOTION_CLIENT_ID or "",
@@ -244,7 +276,6 @@ def _build_data_source_oauth_services(
         redirect_uri=dify_config.CONSOLE_API_URL + "/console/api/oauth/data-source/callback/notion",
         http_client=ssrf_proxy,
     )
-    bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client)
     return {
         "notion": DataSourceOAuthService(
             provider_name="notion",
@@ -277,6 +308,7 @@ def build_application_services(
 ) -> ApplicationServices:
     installation_state = InstallationStateRepository(session_factory=database_client)
     data_source_api_key_auth_bindings = SQLAlchemyDataSourceApiKeyAuthBindingRepository(session_factory=database_client)
+    data_source_oauth_bindings = SQLAlchemyDataSourceOAuthBindingRepository(session_factory=database_client)
     app_definition_repository = AppDefinitionQueryRepository(session_factory=database_client)
     feature_gateway = FeatureServiceGateway()
     accounts = SQLAlchemyAccountRepository(session_factory=database_client)
@@ -291,6 +323,29 @@ def build_application_services(
         builtin=builtin_catalog,
     )
     workspace_query_repository = WorkspaceQueryRepository(session_factory=database_client)
+    workspace_member_repository = WorkspaceMemberQueryRepository(session_factory=database_client)
+    datasets = SQLAlchemyDatasetRepository(session_factory=database_client)
+    documents = SQLAlchemyDocumentRepository(session_factory=database_client)
+    dataset_access = DatasetAccessService(
+        datasets=datasets,
+        workspace_roles=workspace_member_repository,
+        legacy_permissions_enabled=not dify_config.RBAC_ENABLED,
+    )
+    datasource_credentials = SQLAlchemyDatasourceCredentialRepository(session_factory=database_client)
+    datasource_provider_manager = PluginDatasourceManager()
+    credential_gateway = ActorAwareDatasourceCredentialGateway(
+        credentials=datasource_credentials,
+        codec=PluginDatasourceCredentialCodec(provider_manager=datasource_provider_manager),
+        refresher=OAuthDatasourceCredentialRefresher(
+            session_factory=database_client,
+            provider_manager=datasource_provider_manager,
+        ),
+    )
+    notion_imports = NotionImportApplicationService(
+        dataset_access=dataset_access,
+        documents=documents,
+        source=PluginNotionSourceGateway(credentials=credential_gateway),
+    )
     file_service = FileService(session_factory=database_client)
     passwords = DefaultAccountPasswordHasher()
     return ApplicationServices(
@@ -444,12 +499,31 @@ def build_application_services(
                 redis_client=redis,
             ),
         ),
-        data_source_api_key_auth=DataSourceApiKeyAuthService(
-            bindings=data_source_api_key_auth_bindings,
-            validator=ProviderApiKeyAuthCredentialValidator(),
-            encryptor=TenantApiKeyAuthCredentialEncryptor(),
+        data_sources=DataSourceServices(
+            api_key_auth=DataSourceApiKeyAuthService(
+                bindings=data_source_api_key_auth_bindings,
+                validator=ProviderApiKeyAuthCredentialValidator(),
+                encryptor=TenantApiKeyAuthCredentialEncryptor(),
+            ),
+            oauth=_build_data_source_oauth_services(bindings=data_source_oauth_bindings),
+            bindings=DataSourceBindingApplicationService(bindings=data_source_oauth_bindings),
+            notion_imports=notion_imports,
+        ),
+        knowledge=KnowledgeServices(
+            document_sync=DocumentSyncApplicationService(
+                dataset_access=dataset_access,
+                documents=documents,
+                dispatcher=CeleryDocumentSyncDispatcher(),
+            ),
+            indexing_estimates=IndexingEstimateApplicationService(
+                dataset_access=dataset_access,
+                documents=documents,
+                gateway=IndexingRunnerEstimateGateway(
+                    session_factory=database_client,
+                    credentials=credential_gateway,
+                ),
+            ),
         ),
-        data_source_oauth=_build_data_source_oauth_services(database_client=database_client),
         webapp_access=WebAppAccessQueryService(
             access=WebAppAccessQueryRepository(session_factory=database_client),
             webapp_auth_enabled=SystemFeatureService.is_webapp_auth_enabled(deployment_edition=deployment_edition),
@@ -514,9 +588,7 @@ def build_application_services(
             plans=DeploymentWorkspacePlanGateway(),
         ),
         workspace_member_queries=WorkspaceMemberQueryService(
-            members=WorkspaceMemberQueryRepository(
-                session_factory=database_client,
-            ),
+            members=workspace_member_repository,
             roles=DeploymentWorkspaceMemberRoleResolver(),
         ),
         inner_mail=InnerMailService(dispatch=enqueue_inner_mail),

api/models/dataset.py

@@ -18,7 +18,7 @@ from sqlalchemy import DateTime, String, func, select
 from sqlalchemy.orm import Mapped, Session, mapped_column, scoped_session
 
 from configs import dify_config
-from core.rag.entities import ParentMode, Rule
+from core.rag.entities import ParentMode, PreProcessingRuleKey, Rule
 from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.index_processor.constant.query_type import QueryType
@@ -53,7 +53,7 @@ logger = logging.getLogger(__name__)
 
 
 class PreProcessingRuleItem(TypedDict):
-    id: str
+    id: PreProcessingRuleKey
     enabled: bool
 
 
@@ -490,7 +490,6 @@ class DatasetProcessRule(TypeBase):
     )
 
     MODES = ["automatic", "custom", "hierarchical"]
-    PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
     AUTOMATIC_RULES: ClassVar[AutomaticRulesConfig] = {
         "pre_processing_rules": [
             {"id": "remove_extra_spaces", "enabled": True},

api/openapi/markdown/console-openapi.md

@@ -20849,7 +20849,13 @@ Verification of the plugin.
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
 | enabled | boolean | Whether this preprocessing rule is enabled. | Yes |
-| id | string, <br>**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.<br>*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes |
+| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes |
+
+#### PreProcessingRuleKey
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| PreProcessingRuleKey | string |  |  |
 
 #### PreviewDetail
 

api/openapi/markdown/service-openapi.md

@@ -3848,7 +3848,13 @@ JSON result for published runs and draft runs using `response_mode: blocking`.
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
 | enabled | boolean | Whether this preprocessing rule is enabled. | Yes |
-| id | string, <br>**Available values:** "remove_extra_spaces", "remove_stopwords", "remove_urls_emails" | Rule identifier.<br>*Enum:* `"remove_extra_spaces"`, `"remove_stopwords"`, `"remove_urls_emails"` | Yes |
+| id | [PreProcessingRuleKey](#preprocessingrulekey) | Rule identifier. | Yes |
+
+#### PreProcessingRuleKey
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| PreProcessingRuleKey | string |  |  |
 
 #### ProcessRule
 

api/repositories/data_source/__init__.py

@@ -0,0 +1 @@
+"""Persistence adapters for the data-source domain."""

api/repositories/data_source/api_key_auth_repository.py

@@ -9,7 +9,7 @@ from sqlalchemy.orm import Session, sessionmaker
 
 from models.source import DataSourceApiKeyAuthBinding
 from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthBindingRepository
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingRecord,
     DataSourceApiKeyAuthCredentials,
 )

api/repositories/data_source/credential_repository.py

@@ -0,0 +1,92 @@
+"""Actor-aware SQLAlchemy repository for datasource credentials."""
+
+from collections.abc import Mapping
+from typing import cast
+
+from pydantic import TypeAdapter
+from sqlalchemy import select, update
+from sqlalchemy.engine import CursorResult
+from sqlalchemy.orm import Session, sessionmaker
+
+from libs.datetime_utils import naive_utc_now
+from models.credential_permission import CredentialType
+from models.oauth import DatasourceProvider
+from services.credential_permission_service import CredentialPermissionService
+from services.entities.data_source.credential import DatasourceCredentialRecord
+
+_CREDENTIALS_ADAPTER = TypeAdapter(dict[str, object])
+
+
+def _record(provider: DatasourceProvider) -> DatasourceCredentialRecord:
+    return DatasourceCredentialRecord(
+        id=provider.id,
+        workspace_id=provider.tenant_id,
+        name=provider.name,
+        provider=provider.provider,
+        plugin_id=provider.plugin_id,
+        auth_type=provider.auth_type,
+        encrypted_credentials=_CREDENTIALS_ADAPTER.validate_python(provider.encrypted_credentials),
+        expires_at=provider.expires_at,
+        updated_at=provider.updated_at,
+    )
+
+
+class SQLAlchemyDatasourceCredentialRepository:
+    def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
+        self._session_factory = session_factory
+
+    def get_visible(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> DatasourceCredentialRecord | None:
+        query = select(DatasourceProvider).where(
+            DatasourceProvider.tenant_id == workspace_id,
+            DatasourceProvider.id == credential_id,
+            DatasourceProvider.provider == provider,
+            DatasourceProvider.plugin_id == plugin_id,
+        )
+        query = CredentialPermissionService.apply_visibility_filter_for_actor(
+            query,
+            tenant_id=workspace_id,
+            model_id_column=DatasourceProvider.id,
+            model_user_id_column=DatasourceProvider.user_id,
+            model_visibility_column=DatasourceProvider.visibility,
+            credential_type=CredentialType.DATASOURCE_PROVIDER,
+            actor_id=actor_id,
+        )
+        with self._session_factory() as session:
+            provider_record = session.scalar(query.limit(1))
+            return _record(provider_record) if provider_record is not None else None
+
+    def update_if_unchanged(
+        self,
+        *,
+        record: DatasourceCredentialRecord,
+        encrypted_credentials: Mapping[str, object],
+        expires_at: int,
+    ) -> bool:
+        credentials = _CREDENTIALS_ADAPTER.validate_python(encrypted_credentials)
+        with self._session_factory.begin() as session:
+            result = session.execute(
+                update(DatasourceProvider)
+                .where(
+                    DatasourceProvider.id == record.id,
+                    DatasourceProvider.tenant_id == record.workspace_id,
+                    DatasourceProvider.provider == record.provider,
+                    DatasourceProvider.plugin_id == record.plugin_id,
+                    DatasourceProvider.encrypted_credentials
+                    == _CREDENTIALS_ADAPTER.validate_python(record.encrypted_credentials),
+                    DatasourceProvider.expires_at == record.expires_at,
+                )
+                .values(
+                    encrypted_credentials=credentials,
+                    expires_at=expires_at,
+                    updated_at=naive_utc_now(),
+                )
+            )
+            return cast(CursorResult[object], result).rowcount == 1

api/repositories/data_source/oauth_binding_repository.py

@@ -1,16 +1,19 @@
 """SQLAlchemy repository for OAuth data-source bindings."""
 
 from collections.abc import Mapping
-from typing import override
+from typing import cast, override
 
 from pydantic import TypeAdapter
-from sqlalchemy import select
+from sqlalchemy import select, update
+from sqlalchemy.engine import CursorResult
 from sqlalchemy.orm import Session, sessionmaker
 
 from libs.datetime_utils import naive_utc_now
 from models.source import DataSourceOauthBinding
-from services.data_source_oauth_service import DataSourceOAuthBindingRepository
-from services.entities.data_source_oauth_entities import (
+from services.data_source.binding_application_service import BindingMutationResult, DataSourceBindingStore
+from services.data_source.oauth_service import DataSourceOAuthBindingRepository
+from services.entities.data_source.oauth import (
+    DataSourceBindingSummary,
     DataSourceOAuthAuthorization,
     DataSourceOAuthBindingRecord,
 )
@@ -18,7 +21,7 @@ from services.entities.data_source_oauth_entities import (
 _SOURCE_INFO_ADAPTER = TypeAdapter(dict[str, object])
 
 
-class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository):
+class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepository, DataSourceBindingStore):
     def __init__(self, session_factory: sessionmaker[Session]) -> None:
         self._session_factory = session_factory
 
@@ -68,7 +71,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
                     DataSourceOauthBinding.tenant_id == workspace_id,
                     DataSourceOauthBinding.provider == provider,
                     DataSourceOauthBinding.id == binding_id,
-                    DataSourceOauthBinding.disabled.is_(False),
+                    DataSourceOauthBinding.disabled.is_not(True),
                 )
             )
             if binding is None:
@@ -96,7 +99,7 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
                     DataSourceOauthBinding.tenant_id == workspace_id,
                     DataSourceOauthBinding.provider == provider,
                     DataSourceOauthBinding.id == binding_id,
-                    DataSourceOauthBinding.disabled.is_(False),
+                    DataSourceOauthBinding.disabled.is_not(True),
                 )
             )
             if binding is None:
@@ -105,3 +108,59 @@ class SQLAlchemyDataSourceOAuthBindingRepository(DataSourceOAuthBindingRepositor
             binding.source_info = validated_source_info
             binding.updated_at = naive_utc_now()
             return True
+
+    @override
+    def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]:
+        with self._session_factory() as session:
+            bindings = session.scalars(
+                select(DataSourceOauthBinding).where(
+                    DataSourceOauthBinding.tenant_id == workspace_id,
+                    DataSourceOauthBinding.disabled.is_not(True),
+                )
+            ).all()
+            return tuple(
+                DataSourceBindingSummary(
+                    id=binding.id,
+                    provider=binding.provider,
+                    created_at=binding.created_at,
+                    disabled=False,
+                    source_info=_SOURCE_INFO_ADAPTER.validate_python(binding.source_info),
+                )
+                for binding in bindings
+            )
+
+    @override
+    def change_disabled_state(
+        self,
+        *,
+        workspace_id: str,
+        binding_id: str,
+        disabled: bool,
+    ) -> BindingMutationResult:
+        state_predicate = (
+            DataSourceOauthBinding.disabled.is_not(True) if disabled else DataSourceOauthBinding.disabled.is_(True)
+        )
+        with self._session_factory.begin() as session:
+            result = session.execute(
+                update(DataSourceOauthBinding)
+                .where(
+                    DataSourceOauthBinding.tenant_id == workspace_id,
+                    DataSourceOauthBinding.id == binding_id,
+                    state_predicate,
+                )
+                .values(disabled=disabled, updated_at=naive_utc_now())
+            )
+            if cast(CursorResult[object], result).rowcount == 1:
+                return "updated"
+
+            binding_exists = session.scalar(
+                select(DataSourceOauthBinding.id)
+                .where(
+                    DataSourceOauthBinding.tenant_id == workspace_id,
+                    DataSourceOauthBinding.id == binding_id,
+                )
+                .limit(1)
+            )
+            if binding_exists is None:
+                return "not_found"
+            return "already_disabled" if disabled else "already_enabled"

api/repositories/knowledge/__init__.py

@@ -0,0 +1 @@
+"""Persistence adapters for the knowledge domain."""

api/repositories/knowledge/dataset_repository.py

@@ -0,0 +1,123 @@
+"""SQLAlchemy repository for tenant-owned dataset state."""
+
+from collections.abc import Sequence
+from typing import override
+
+from sqlalchemy import delete, exists, select
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.rag.entities.dataset_reference import DatasetRef
+from models.dataset import Dataset, DatasetPermission
+from services.entities.knowledge_entities.records import DatasetAccessSnapshot, DatasetRecord
+from services.knowledge.application import DatasetAccessReader, DatasetNotFoundError
+
+
+def _dataset_record(dataset: Dataset) -> DatasetRecord:
+    return DatasetRecord(
+        id=dataset.id,
+        workspace_id=dataset.tenant_id,
+        maintainer_id=dataset.maintainer,
+        permission=str(dataset.permission),
+        data_source_type=str(dataset.data_source_type) if dataset.data_source_type is not None else None,
+        indexing_technique=str(dataset.indexing_technique) if dataset.indexing_technique is not None else None,
+        embedding_model=dataset.embedding_model,
+        embedding_model_provider=dataset.embedding_model_provider,
+    )
+
+
+class SQLAlchemyDatasetRepository(DatasetAccessReader):
+    """Own SQL access to datasets and their partial-member grants."""
+
+    def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
+        self._session_factory = session_factory
+
+    def get_by_ref(self, dataset_ref: DatasetRef) -> DatasetRecord | None:
+        with self._session_factory() as session:
+            dataset = session.scalar(
+                select(Dataset)
+                .where(
+                    Dataset.id == dataset_ref.dataset_id,
+                    Dataset.tenant_id == dataset_ref.tenant_id,
+                )
+                .limit(1)
+            )
+            return _dataset_record(dataset) if dataset is not None else None
+
+    @override
+    def get_access_snapshot(
+        self,
+        *,
+        workspace_id: str,
+        dataset_id: str,
+        actor_id: str,
+    ) -> DatasetAccessSnapshot | None:
+        with self._session_factory() as session:
+            dataset = session.scalar(
+                select(Dataset)
+                .where(
+                    Dataset.id == dataset_id,
+                    Dataset.tenant_id == workspace_id,
+                )
+                .limit(1)
+            )
+            if dataset is None:
+                return None
+
+            actor_has_partial_access = bool(
+                session.scalar(
+                    select(
+                        exists().where(
+                            DatasetPermission.tenant_id == workspace_id,
+                            DatasetPermission.dataset_id == dataset_id,
+                            DatasetPermission.account_id == actor_id,
+                            DatasetPermission.has_permission.is_(True),
+                        )
+                    )
+                )
+            )
+            return DatasetAccessSnapshot(
+                dataset=_dataset_record(dataset),
+                actor_has_partial_access=actor_has_partial_access,
+            )
+
+    def list_partial_member_ids(self, dataset_ref: DatasetRef) -> tuple[str, ...]:
+        with self._session_factory() as session:
+            member_ids = session.scalars(
+                select(DatasetPermission.account_id).where(
+                    DatasetPermission.tenant_id == dataset_ref.tenant_id,
+                    DatasetPermission.dataset_id == dataset_ref.dataset_id,
+                    DatasetPermission.has_permission.is_(True),
+                )
+            ).all()
+            return tuple(member_ids)
+
+    def replace_partial_member_ids(self, dataset_ref: DatasetRef, member_ids: Sequence[str]) -> None:
+        unique_member_ids = tuple(dict.fromkeys(member_ids))
+        with self._session_factory.begin() as session:
+            # Serialize replacements on the aggregate root so concurrent writers cannot
+            # leave the union of two independently requested member sets.
+            locked_dataset_id = session.scalar(
+                select(Dataset.id)
+                .where(
+                    Dataset.id == dataset_ref.dataset_id,
+                    Dataset.tenant_id == dataset_ref.tenant_id,
+                )
+                .with_for_update()
+            )
+            if locked_dataset_id is None:
+                raise DatasetNotFoundError()
+
+            session.execute(
+                delete(DatasetPermission).where(
+                    DatasetPermission.tenant_id == dataset_ref.tenant_id,
+                    DatasetPermission.dataset_id == dataset_ref.dataset_id,
+                )
+            )
+            session.add_all(
+                DatasetPermission(
+                    tenant_id=dataset_ref.tenant_id,
+                    dataset_id=dataset_ref.dataset_id,
+                    account_id=member_id,
+                )
+                for member_id in unique_member_ids
+            )

api/repositories/knowledge/document_repository.py

@@ -0,0 +1,170 @@
+"""SQLAlchemy repository for tenant-owned document state."""
+
+import json
+from collections.abc import Mapping, Sequence
+
+from pydantic import TypeAdapter, ValidationError
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef
+from models.dataset import Dataset, Document
+from models.enums import IndexingStatus
+from services.entities.knowledge_entities.records import DocumentRecord
+
+_MAPPING_ADAPTER = TypeAdapter(dict[str, object])
+
+
+def _mapping(value: object) -> Mapping[str, object] | None:
+    if value is None:
+        return None
+    try:
+        if isinstance(value, str):
+            return _MAPPING_ADAPTER.validate_python(json.loads(value))
+        return _MAPPING_ADAPTER.validate_python(value)
+    except (TypeError, ValueError, ValidationError):
+        return {}
+
+
+def _document_record(document: Document) -> DocumentRecord:
+    return DocumentRecord(
+        id=document.id,
+        workspace_id=document.tenant_id,
+        dataset_id=document.dataset_id,
+        name=document.name,
+        data_source_type=str(document.data_source_type),
+        data_source_info=_mapping(document.data_source_info),
+        enabled=document.enabled,
+        archived=document.archived,
+        indexing_status=str(document.indexing_status),
+        batch=document.batch,
+        doc_form=str(document.doc_form),
+        doc_language=document.doc_language,
+        dataset_process_rule_id=document.dataset_process_rule_id,
+        need_summary=document.need_summary,
+        doc_metadata=_mapping(document.doc_metadata),
+    )
+
+
+class SQLAlchemyDocumentRepository:
+    """Own SQL access to documents through their complete owner chain."""
+
+    def __init__(self, *, session_factory: sessionmaker[Session]) -> None:
+        self._session_factory = session_factory
+
+    def get_by_ref(self, document_ref: DocumentRef) -> DocumentRecord | None:
+        with self._session_factory() as session:
+            document = session.scalar(
+                select(Document)
+                .where(
+                    Document.id == document_ref.document_id,
+                    Document.dataset_id == document_ref.dataset.dataset_id,
+                    Document.tenant_id == document_ref.dataset.tenant_id,
+                )
+                .limit(1)
+            )
+            return _document_record(document) if document is not None else None
+
+    def get_by_id_for_workspace(self, *, workspace_id: str, document_id: str) -> DocumentRecord | None:
+        with self._session_factory() as session:
+            document = session.scalar(
+                select(Document)
+                .join(Dataset, Dataset.id == Document.dataset_id)
+                .where(
+                    Document.id == document_id,
+                    Document.tenant_id == workspace_id,
+                    Dataset.tenant_id == workspace_id,
+                )
+                .limit(1)
+            )
+            return _document_record(document) if document is not None else None
+
+    def list_by_refs(self, dataset_ref: DatasetRef, document_ids: Sequence[str]) -> tuple[DocumentRecord, ...]:
+        return self._list_by_refs(dataset_ref, document_ids, available_only=False)
+
+    def list_available_by_refs(
+        self,
+        dataset_ref: DatasetRef,
+        document_ids: Sequence[str],
+    ) -> tuple[DocumentRecord, ...]:
+        return self._list_by_refs(dataset_ref, document_ids, available_only=True)
+
+    def _list_by_refs(
+        self,
+        dataset_ref: DatasetRef,
+        document_ids: Sequence[str],
+        *,
+        available_only: bool,
+    ) -> tuple[DocumentRecord, ...]:
+        unique_document_ids = tuple(dict.fromkeys(document_ids))
+        if not unique_document_ids:
+            return ()
+
+        stmt = select(Document).where(
+            Document.tenant_id == dataset_ref.tenant_id,
+            Document.dataset_id == dataset_ref.dataset_id,
+            Document.id.in_(unique_document_ids),
+        )
+        if available_only:
+            stmt = stmt.where(
+                Document.enabled.is_(True),
+                Document.indexing_status == IndexingStatus.COMPLETED,
+                Document.archived.is_(False),
+            )
+
+        with self._session_factory() as session:
+            documents = session.scalars(stmt).all()
+            records_by_id = {document.id: _document_record(document) for document in documents}
+            return tuple(
+                records_by_id[document_id] for document_id in unique_document_ids if document_id in records_by_id
+            )
+
+    def list_by_batch(self, dataset_ref: DatasetRef, batch: str) -> tuple[DocumentRecord, ...]:
+        with self._session_factory() as session:
+            documents = session.scalars(
+                select(Document).where(
+                    Document.tenant_id == dataset_ref.tenant_id,
+                    Document.dataset_id == dataset_ref.dataset_id,
+                    Document.batch == batch,
+                )
+            ).all()
+            return tuple(_document_record(document) for document in documents)
+
+    def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]:
+        return tuple(dataset_ref.document(document.id) for document in self._list_active_notion(dataset_ref))
+
+    def _list_active_notion(self, dataset_ref: DatasetRef) -> tuple[DocumentRecord, ...]:
+        with self._session_factory() as session:
+            documents = session.scalars(
+                select(Document).where(
+                    Document.tenant_id == dataset_ref.tenant_id,
+                    Document.dataset_id == dataset_ref.dataset_id,
+                    Document.data_source_type == "notion_import",
+                    Document.enabled.is_(True),
+                    Document.archived.is_(False),
+                )
+            ).all()
+            return tuple(_document_record(document) for document in documents)
+
+    def list_working_by_dataset(self, dataset_ref: DatasetRef) -> tuple[DocumentRecord, ...]:
+        with self._session_factory() as session:
+            documents = session.scalars(
+                select(Document).where(
+                    Document.tenant_id == dataset_ref.tenant_id,
+                    Document.dataset_id == dataset_ref.dataset_id,
+                    Document.enabled.is_(True),
+                    Document.indexing_status == IndexingStatus.COMPLETED,
+                    Document.archived.is_(False),
+                )
+            ).all()
+            return tuple(_document_record(document) for document in documents)
+
+    def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]:
+        page_ids: set[str] = set()
+        for document in self._list_active_notion(dataset_ref):
+            if document.data_source_info is None:
+                continue
+            page_id = document.data_source_info.get("notion_page_id")
+            if isinstance(page_id, str) and page_id:
+                page_ids.add(page_id)
+        return frozenset(page_ids)

api/repositories/workspace_member_query_repository.py

@@ -60,3 +60,17 @@ class WorkspaceMemberQueryRepository(WorkspaceMemberQuery):
             )
 
         return records
+
+    def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None:
+        """Return one member's role without loading the whole workspace roster."""
+
+        with self._session_factory() as session:
+            role = session.scalar(
+                select(TenantAccountJoin.role)
+                .where(
+                    TenantAccountJoin.tenant_id == workspace_id,
+                    TenantAccountJoin.account_id == account_id,
+                )
+                .limit(1)
+            )
+            return role.value if role is not None else None

api/services/auth/data_source_api_key_auth_gateways.py

@@ -11,7 +11,7 @@ from services.auth.errors import (
     DataSourceApiKeyAuthProviderUnavailableError,
     UnsupportedDataSourceApiKeyAuthProviderError,
 )
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 class _ProviderApiKeyAuthValidator(Protocol):

api/services/auth/data_source_api_key_auth_service.py

@@ -5,7 +5,7 @@ from dataclasses import replace
 from typing import Protocol
 
 from machinery.context import RequestContext
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthBindingRecord,
     DataSourceApiKeyAuthCredentials,

api/services/auth/firecrawl/firecrawl.py

@@ -9,7 +9,7 @@ from services.auth.errors import (
     DataSourceApiKeyAuthProviderUnavailableError,
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 # Explicit bounded timeout for credential-validation requests so a slow or
 # hanging Firecrawl endpoint cannot block the worker indefinitely.

api/services/auth/jina/jina.py

@@ -10,7 +10,7 @@ from services.auth.errors import (
     DataSourceApiKeyAuthProviderUnavailableError,
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 _http_client: httpx.Client = get_pooled_http_client(
     "auth:jina",

api/services/auth/watercrawl/watercrawl.py

@@ -10,7 +10,7 @@ from services.auth.errors import (
     DataSourceApiKeyAuthProviderUnavailableError,
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 # Explicit bounded timeout for credential-validation requests so a slow or
 # hanging WaterCrawl endpoint cannot block the worker indefinitely.

api/services/credential_permission_service.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 
-from sqlalchemy import or_, select
+from sqlalchemy import and_, or_, select
 from sqlalchemy.orm import InstrumentedAttribute, Session
 
 from models.account import Account
@@ -30,28 +30,52 @@ class CredentialPermissionService:
         cls,
         query,
         *,
+        tenant_id: str,
         model_id_column: InstrumentedAttribute,
         model_user_id_column: InstrumentedAttribute,
         model_visibility_column: InstrumentedAttribute,
         credential_type: str,
         user: Account,
+    ):
+        return cls.apply_visibility_filter_for_actor(
+            query,
+            tenant_id=tenant_id,
+            model_id_column=model_id_column,
+            model_user_id_column=model_user_id_column,
+            model_visibility_column=model_visibility_column,
+            credential_type=credential_type,
+            actor_id=user.id,
+        )
+
+    @classmethod
+    def apply_visibility_filter_for_actor(
+        cls,
+        query,
+        *,
+        tenant_id: str,
+        model_id_column: InstrumentedAttribute,
+        model_user_id_column: InstrumentedAttribute,
+        model_visibility_column: InstrumentedAttribute,
+        credential_type: str,
+        actor_id: str,
     ):
         """
-        Add WHERE clauses to a SQLAlchemy query so it only returns credentials
-        visible to the given user.
+        Add WHERE clauses so a query only returns credentials visible to an actor.
 
         - all_team_members: always visible
-        - only_me: visible only to the creator (user.id matches)
+        - only_me: visible only to the creator
         - partial_members: visible to the creator OR users in credential_permissions
         - Legacy rows with NULL user_id are treated as all_team_members
         - No admin bypass: personal credentials are private regardless of role
         """
-        # Subquery: credential_ids where user has partial-member permission
+        # Subquery: credential_ids where the actor has partial-member permission.
         partial_subquery = (
             select(CredentialPermission.credential_id)
             .where(
+                CredentialPermission.tenant_id == tenant_id,
                 CredentialPermission.credential_type == credential_type,
-                CredentialPermission.account_id == user.id,
+                CredentialPermission.account_id == actor_id,
+                CredentialPermission.has_permission.is_(True),
             )
             .correlate_except(CredentialPermission)
         )
@@ -63,8 +87,11 @@ class CredentialPermissionService:
                 # legacy rows with NULL user_id treated as all_team
                 model_user_id_column.is_(None),
                 # only_me: creator sees their own
-                (model_user_id_column == user.id),
+                (model_user_id_column == actor_id),
                 # partial_members: user is in the permission table
-                model_id_column.in_(partial_subquery),
+                and_(
+                    model_visibility_column == PermissionEnum.PARTIAL_TEAM,
+                    model_id_column.in_(partial_subquery),
+                ),
             )
         )

api/services/data_source/__init__.py

@@ -0,0 +1 @@
+"""Application services and adapters for data-source use cases."""

api/services/data_source/binding_application_service.py

@@ -0,0 +1,68 @@
+"""Application service for managing data-source OAuth bindings."""
+
+from collections.abc import Sequence
+from typing import Literal, Protocol
+
+from machinery.context import RequestContext
+from services.entities.data_source.oauth import DataSourceBindingSummary
+
+type BindingMutationResult = Literal["updated", "not_found", "already_enabled", "already_disabled"]
+
+_SUPPORTED_PROVIDERS = frozenset({"notion"})
+
+
+class DataSourceBindingStore(Protocol):
+    def list_enabled_bindings(self, *, workspace_id: str) -> Sequence[DataSourceBindingSummary]: ...
+
+    def change_disabled_state(
+        self,
+        *,
+        workspace_id: str,
+        binding_id: str,
+        disabled: bool,
+    ) -> BindingMutationResult: ...
+
+
+class DataSourceBindingError(Exception):
+    """Base class for framework-neutral binding management failures."""
+
+
+class DataSourceBindingNotFoundError(DataSourceBindingError):
+    def __init__(self) -> None:
+        super().__init__("Data source binding not found")
+
+
+class DataSourceBindingStateError(DataSourceBindingError):
+    def __init__(self, *, disabled: bool) -> None:
+        state = "disabled" if disabled else "enabled"
+        super().__init__(f"Data source is already {state}")
+
+
+class DataSourceBindingApplicationService:
+    def __init__(self, *, bindings: DataSourceBindingStore) -> None:
+        self._bindings = bindings
+
+    def list_integrations(self, context: RequestContext) -> tuple[DataSourceBindingSummary, ...]:
+        return tuple(
+            binding
+            for binding in self._bindings.list_enabled_bindings(workspace_id=context.active_workspace_id)
+            if binding.provider in _SUPPORTED_PROVIDERS
+        )
+
+    def enable(self, context: RequestContext, binding_id: str) -> None:
+        self._change_state(context, binding_id, disabled=False)
+
+    def disable(self, context: RequestContext, binding_id: str) -> None:
+        self._change_state(context, binding_id, disabled=True)
+
+    def _change_state(self, context: RequestContext, binding_id: str, *, disabled: bool) -> None:
+        result = self._bindings.change_disabled_state(
+            workspace_id=context.active_workspace_id,
+            binding_id=binding_id,
+            disabled=disabled,
+        )
+        if result == "updated":
+            return
+        if result == "not_found":
+            raise DataSourceBindingNotFoundError()
+        raise DataSourceBindingStateError(disabled=disabled)

api/services/data_source/credential_gateway.py

@@ -0,0 +1,297 @@
+"""Explicit-actor credential resolution and OAuth refresh infrastructure."""
+
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Protocol
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from configs import dify_config
+from core.helper import encrypter
+from core.helper.provider_cache import NoOpProviderCredentialCache
+from core.plugin.entities.plugin_daemon import CredentialType
+from core.plugin.impl.datasource import PluginDatasourceManager
+from core.plugin.impl.oauth import OAuthHandler
+from core.plugin.plugin_service import PluginService
+from core.tools.utils.encryption import create_provider_encrypter
+from graphon.model_runtime.entities.provider_entities import FormType
+from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig
+from models.provider_ids import DatasourceProviderID
+from services.entities.data_source.credential import DatasourceCredentialRecord
+
+
+class DatasourceCredentialStore(Protocol):
+    def get_visible(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> DatasourceCredentialRecord | None: ...
+
+    def update_if_unchanged(
+        self,
+        *,
+        record: DatasourceCredentialRecord,
+        encrypted_credentials: Mapping[str, object],
+        expires_at: int,
+    ) -> bool: ...
+
+
+class DatasourceCredentialCodec(Protocol):
+    def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]: ...
+
+    def encrypt(
+        self,
+        record: DatasourceCredentialRecord,
+        credentials: Mapping[str, object],
+    ) -> dict[str, object]: ...
+
+
+@dataclass(frozen=True, slots=True)
+class RefreshedDatasourceCredential:
+    credentials: Mapping[str, object]
+    expires_at: int
+
+
+class DatasourceCredentialRefresher(Protocol):
+    def refresh(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        record: DatasourceCredentialRecord,
+        credentials: Mapping[str, object],
+    ) -> RefreshedDatasourceCredential: ...
+
+
+class DatasourceCredentialError(Exception):
+    """Base class for credential resolution failures."""
+
+
+class DatasourceCredentialNotFoundError(DatasourceCredentialError):
+    def __init__(self) -> None:
+        super().__init__("Credential not found")
+
+
+class DatasourceCredentialConcurrentUpdateError(DatasourceCredentialError):
+    def __init__(self) -> None:
+        super().__init__("Credential changed while it was being refreshed")
+
+
+class DatasourceCredentialRefreshError(DatasourceCredentialError):
+    def __init__(self, credential_id: str) -> None:
+        super().__init__(f"Failed to refresh datasource credential: {credential_id}")
+
+
+class PluginDatasourceCredentialCodec:
+    def __init__(self, provider_manager: PluginDatasourceManager | None = None) -> None:
+        self._provider_manager = provider_manager or PluginDatasourceManager()
+
+    def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]:
+        result = dict(record.encrypted_credentials)
+        for name in self._secret_variables(record):
+            value = result.get(name)
+            if isinstance(value, str):
+                result[name] = encrypter.decrypt_token(record.workspace_id, value)
+        return result
+
+    def encrypt(
+        self,
+        record: DatasourceCredentialRecord,
+        credentials: Mapping[str, object],
+    ) -> dict[str, object]:
+        result = dict(credentials)
+        for name in self._secret_variables(record):
+            value = result.get(name)
+            if isinstance(value, str):
+                result[name] = encrypter.encrypt_token(record.workspace_id, value)
+        return result
+
+    def _secret_variables(self, record: DatasourceCredentialRecord) -> tuple[str, ...]:
+        provider = self._provider_manager.fetch_datasource_provider(
+            tenant_id=record.workspace_id,
+            provider_id=f"{record.plugin_id}/{record.provider}",
+        )
+        credential_type = CredentialType.of(record.auth_type)
+        if credential_type == CredentialType.API_KEY:
+            schemas = provider.declaration.credentials_schema
+        elif credential_type == CredentialType.OAUTH2 and provider.declaration.oauth_schema:
+            schemas = provider.declaration.oauth_schema.credentials_schema
+        else:
+            raise DatasourceCredentialError(f"Unsupported credential type: {record.auth_type}")
+        return tuple(schema.name for schema in schemas if schema.type.value == FormType.SECRET_INPUT)
+
+
+class OAuthDatasourceCredentialRefresher:
+    def __init__(
+        self,
+        *,
+        session_factory: sessionmaker[Session],
+        provider_manager: PluginDatasourceManager | None = None,
+        oauth_handler: OAuthHandler | None = None,
+    ) -> None:
+        self._session_factory = session_factory
+        self._provider_manager = provider_manager or PluginDatasourceManager()
+        self._oauth_handler = oauth_handler or OAuthHandler()
+
+    def refresh(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        record: DatasourceCredentialRecord,
+        credentials: Mapping[str, object],
+    ) -> RefreshedDatasourceCredential:
+        provider_id = DatasourceProviderID(f"{record.plugin_id}/{record.provider}")
+        try:
+            refreshed = self._oauth_handler.refresh_credentials(
+                tenant_id=workspace_id,
+                user_id=actor_id,
+                plugin_id=provider_id.plugin_id,
+                provider=provider_id.provider_name,
+                redirect_uri=(
+                    f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/datasource/callback"
+                ),
+                system_credentials=self._oauth_client(workspace_id, provider_id),
+                credentials=dict(credentials),
+            )
+        except DatasourceCredentialError:
+            raise
+        except Exception as error:
+            raise DatasourceCredentialRefreshError(record.id) from error
+        return RefreshedDatasourceCredential(
+            credentials=dict(refreshed.credentials),
+            expires_at=refreshed.expires_at,
+        )
+
+    def _oauth_client(self, workspace_id: str, provider_id: DatasourceProviderID) -> dict[str, object]:
+        with self._session_factory() as session:
+            tenant_config = session.scalar(
+                select(DatasourceOauthTenantParamConfig)
+                .where(
+                    DatasourceOauthTenantParamConfig.tenant_id == workspace_id,
+                    DatasourceOauthTenantParamConfig.provider == provider_id.provider_name,
+                    DatasourceOauthTenantParamConfig.plugin_id == provider_id.plugin_id,
+                    DatasourceOauthTenantParamConfig.enabled.is_(True),
+                )
+                .limit(1)
+            )
+            encrypted_tenant_params = dict(tenant_config.client_params) if tenant_config is not None else None
+            system_config = session.scalar(
+                select(DatasourceOauthParamConfig)
+                .where(
+                    DatasourceOauthParamConfig.provider == provider_id.provider_name,
+                    DatasourceOauthParamConfig.plugin_id == provider_id.plugin_id,
+                )
+                .limit(1)
+            )
+            system_credentials = dict(system_config.system_credentials) if system_config is not None else None
+
+        provider = self._provider_manager.fetch_datasource_provider(
+            tenant_id=workspace_id,
+            provider_id=str(provider_id),
+        )
+        if encrypted_tenant_params is not None:
+            oauth_schema = provider.declaration.oauth_schema
+            if oauth_schema is None:
+                raise DatasourceCredentialError("Datasource provider oauth schema not found")
+            credential_encrypter, _ = create_provider_encrypter(
+                tenant_id=workspace_id,
+                config=[item.to_basic_provider_config() for item in oauth_schema.client_schema],
+                cache=NoOpProviderCredentialCache(),
+            )
+            return dict(credential_encrypter.decrypt(encrypted_tenant_params))
+
+        if system_credentials is not None and PluginService.is_plugin_verified(
+            workspace_id, provider.plugin_unique_identifier
+        ):
+            return system_credentials
+        raise DatasourceCredentialError(f"OAuth client is not configured for {provider_id}")
+
+
+class ActorAwareDatasourceCredentialGateway:
+    def __init__(
+        self,
+        *,
+        credentials: DatasourceCredentialStore,
+        codec: DatasourceCredentialCodec,
+        refresher: DatasourceCredentialRefresher,
+        now: Callable[[], float] = time.time,
+    ) -> None:
+        self._credentials = credentials
+        self._codec = codec
+        self._refresher = refresher
+        self._now = now
+
+    def resolve(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> dict[str, object]:
+        record = self._get_visible(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            credential_id=credential_id,
+            provider=provider,
+            plugin_id=plugin_id,
+        )
+        decrypted = self._codec.decrypt(record)
+        if not self._should_refresh(record):
+            return decrypted
+
+        refreshed = self._refresher.refresh(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            record=record,
+            credentials=decrypted,
+        )
+        encrypted = self._codec.encrypt(record, refreshed.credentials)
+        if self._credentials.update_if_unchanged(
+            record=record,
+            encrypted_credentials=encrypted,
+            expires_at=refreshed.expires_at,
+        ):
+            return dict(refreshed.credentials)
+
+        latest = self._get_visible(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            credential_id=credential_id,
+            provider=provider,
+            plugin_id=plugin_id,
+        )
+        if self._should_refresh(latest):
+            raise DatasourceCredentialConcurrentUpdateError()
+        return self._codec.decrypt(latest)
+
+    def _get_visible(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> DatasourceCredentialRecord:
+        record = self._credentials.get_visible(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            credential_id=credential_id,
+            provider=provider,
+            plugin_id=plugin_id,
+        )
+        if record is None:
+            raise DatasourceCredentialNotFoundError()
+        return record
+
+    def _should_refresh(self, record: DatasourceCredentialRecord) -> bool:
+        return record.expires_at != -1 and record.expires_at - 60 < int(self._now())

api/services/data_source/notion_import_adapters.py

@@ -0,0 +1,155 @@
+"""Notion plugin and extractor adapter for import use cases."""
+
+from collections.abc import Callable, Mapping
+from typing import Protocol, cast
+
+from core.datasource.entities.datasource_entities import DatasourceProviderType
+from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
+from core.rag.extractor.notion_extractor import NotionExtractor
+from core.rag.models.document import Document as ExtractedDocument
+from services.entities.data_source.notion_import import (
+    AuthorizedNotionPage,
+    NotionWorkspace,
+    notion_page_icon,
+)
+
+_NOTION_PROVIDER = "notion_datasource"
+_NOTION_PLUGIN = "langgenius/notion_datasource"
+_NOTION_RUNTIME_PROVIDER = f"{_NOTION_PLUGIN}/{_NOTION_PROVIDER}"
+
+
+class _CredentialResolver(Protocol):
+    def resolve(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> dict[str, object]: ...
+
+
+class _RuntimeLoader(Protocol):
+    def __call__(
+        self,
+        *,
+        provider_id: str,
+        datasource_name: str,
+        tenant_id: str,
+        datasource_type: DatasourceProviderType,
+    ) -> object: ...
+
+
+class _NotionPageExtractor(Protocol):
+    def extract(self) -> list[ExtractedDocument]: ...
+
+
+class NotionSourceAdapterError(Exception):
+    """Base class for Notion adapter failures."""
+
+
+class NotionCredentialDataError(NotionSourceAdapterError):
+    def __init__(self) -> None:
+        super().__init__("Notion credential does not contain an integration secret")
+
+
+class PluginNotionSourceGateway:
+    def __init__(
+        self,
+        *,
+        credentials: _CredentialResolver,
+        runtime_loader: _RuntimeLoader | None = None,
+        extractor_factory: Callable[..., _NotionPageExtractor] = NotionExtractor,
+    ) -> None:
+        if runtime_loader is None:
+            from core.datasource.datasource_manager import DatasourceManager
+
+            runtime_loader = DatasourceManager.get_datasource_runtime
+        self._credentials = credentials
+        self._runtime_loader = runtime_loader
+        self._extractor_factory = extractor_factory
+
+    def list_authorized_pages(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+    ) -> tuple[NotionWorkspace, ...]:
+        credentials = self._resolve_credentials(workspace_id, actor_id, credential_id)
+        runtime = cast(
+            OnlineDocumentDatasourcePlugin,
+            self._runtime_loader(
+                provider_id=_NOTION_RUNTIME_PROVIDER,
+                datasource_name=_NOTION_PROVIDER,
+                tenant_id=workspace_id,
+                datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
+            ),
+        )
+        runtime.runtime.credentials = dict(credentials)
+
+        grouped: dict[tuple[str | None, str | None, str | None], list[AuthorizedNotionPage]] = {}
+        for message in runtime.get_online_document_pages(
+            user_id=actor_id,
+            datasource_parameters={},
+            provider_type=runtime.datasource_provider_type(),
+        ):
+            for workspace in message.result:
+                key = (workspace.workspace_id, workspace.workspace_name, workspace.workspace_icon)
+                pages = grouped.setdefault(key, [])
+                pages.extend(
+                    AuthorizedNotionPage(
+                        page_id=page.page_id,
+                        page_name=page.page_name,
+                        page_icon=notion_page_icon(page.page_icon),
+                        parent_id=page.parent_id,
+                        page_type=page.type,
+                    )
+                    for page in workspace.pages
+                )
+        return tuple(
+            NotionWorkspace(
+                workspace_id=key[0],
+                workspace_name=key[1],
+                workspace_icon=key[2],
+                pages=tuple(pages),
+            )
+            for key, pages in grouped.items()
+        )
+
+    def preview_page(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        page_id: str,
+        page_type: str,
+    ) -> str:
+        credentials = self._resolve_credentials(workspace_id, actor_id, credential_id)
+        integration_secret = credentials.get("integration_secret")
+        if not isinstance(integration_secret, str) or not integration_secret:
+            raise NotionCredentialDataError()
+        extractor = self._extractor_factory(
+            notion_workspace_id="",
+            notion_obj_id=page_id,
+            notion_page_type=page_type,
+            notion_access_token=integration_secret,
+            tenant_id=workspace_id,
+        )
+        return "\n".join(document.page_content for document in extractor.extract())
+
+    def _resolve_credentials(
+        self,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+    ) -> Mapping[str, object]:
+        return self._credentials.resolve(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            credential_id=credential_id,
+            provider=_NOTION_PROVIDER,
+            plugin_id=_NOTION_PLUGIN,
+        )

api/services/data_source/notion_import_application_service.py

@@ -0,0 +1,122 @@
+"""Application service for listing and previewing importable Notion pages."""
+
+from typing import Protocol
+
+from core.rag.entities.dataset_reference import DatasetRef
+from machinery.context import RequestContext
+from services.entities.data_source.notion_import import (
+    NotionImportPage,
+    NotionImportResult,
+    NotionImportWorkspace,
+    NotionWorkspace,
+)
+from services.entities.knowledge_entities.records import DatasetRecord
+
+
+class NotionDatasetAccess(Protocol):
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord: ...
+
+
+class NotionDocumentBindingReader(Protocol):
+    def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]: ...
+
+
+class NotionSourceGateway(Protocol):
+    def list_authorized_pages(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+    ) -> tuple[NotionWorkspace, ...]: ...
+
+    def preview_page(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        page_id: str,
+        page_type: str,
+    ) -> str: ...
+
+
+class NotionImportError(Exception):
+    """Base class for framework-neutral Notion import failures."""
+
+
+class DatasetIsNotNotionSourceError(NotionImportError):
+    def __init__(self) -> None:
+        super().__init__("Dataset is not notion type")
+
+
+class NotionImportApplicationService:
+    def __init__(
+        self,
+        *,
+        dataset_access: NotionDatasetAccess,
+        documents: NotionDocumentBindingReader,
+        source: NotionSourceGateway,
+    ) -> None:
+        self._dataset_access = dataset_access
+        self._documents = documents
+        self._source = source
+
+    def list_pages(
+        self,
+        context: RequestContext,
+        *,
+        credential_id: str,
+        dataset_id: str | None = None,
+    ) -> NotionImportResult:
+        bound_page_ids: frozenset[str] = frozenset()
+        if dataset_id is not None:
+            dataset = self._dataset_access.require_accessible(context, dataset_id)
+            if dataset.data_source_type != "notion_import":
+                raise DatasetIsNotNotionSourceError()
+            bound_page_ids = self._documents.list_bound_notion_page_ids(
+                DatasetRef(context.active_workspace_id, dataset.id)
+            )
+
+        workspaces = self._source.list_authorized_pages(
+            workspace_id=context.active_workspace_id,
+            actor_id=context.account_id,
+            credential_id=credential_id,
+        )
+        return NotionImportResult(
+            workspaces=tuple(
+                NotionImportWorkspace(
+                    workspace_id=workspace.workspace_id,
+                    workspace_name=workspace.workspace_name,
+                    workspace_icon=workspace.workspace_icon,
+                    pages=tuple(
+                        NotionImportPage(
+                            page_id=page.page_id,
+                            page_name=page.page_name,
+                            page_icon=page.page_icon,
+                            parent_id=page.parent_id,
+                            page_type=page.page_type,
+                            is_bound=page.page_id in bound_page_ids,
+                        )
+                        for page in workspace.pages
+                    ),
+                )
+                for workspace in workspaces
+            )
+        )
+
+    def preview_page(
+        self,
+        context: RequestContext,
+        *,
+        credential_id: str,
+        page_id: str,
+        page_type: str,
+    ) -> str:
+        return self._source.preview_page(
+            workspace_id=context.active_workspace_id,
+            actor_id=context.account_id,
+            credential_id=credential_id,
+            page_id=page_id,
+            page_type=page_type,
+        )

api/services/data_source/notion_oauth_gateway.py

@@ -9,8 +9,8 @@ from pydantic import TypeAdapter, ValidationError
 
 from core.helper.ssrf_proxy import MaxRetriesExceededError, SSRFProxy
 from core.tools.errors import ToolSSRFError
-from services.data_source_oauth_service import DataSourceOAuthError, DataSourceProviderGateway
-from services.entities.data_source_oauth_entities import DataSourceOAuthAuthorization
+from services.data_source.oauth_service import DataSourceOAuthError, DataSourceProviderGateway
+from services.entities.data_source.oauth import DataSourceOAuthAuthorization
 
 
 class _NotionPageSummary(TypedDict):

api/services/data_source/oauth_service.py

@@ -4,7 +4,7 @@ from collections.abc import Mapping
 from typing import Protocol
 
 from machinery.context import RequestContext
-from services.entities.data_source_oauth_entities import (
+from services.entities.data_source.oauth import (
     DataSourceOAuthAuthorization,
     DataSourceOAuthBindingRecord,
     DataSourceOAuthCallback,

api/services/dataset_ref_service.py

@@ -1,34 +1,12 @@
-"""Typed resource references for dataset ownership chains."""
-
-from typing import NamedTuple
+"""ORM-backed compatibility helpers for dataset ownership chains."""
 
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef, SegmentRef
 from models.dataset import Dataset, Document
 
 
-class DatasetRef(NamedTuple):
-    """Dataset identifiers used to scope downstream resource lookups."""
-
-    tenant_id: str
-    dataset_id: str
-
-
-class DocumentRef(NamedTuple):
-    """Owner-bound lookup coordinates, not proof that a document is authorized or exists."""
-
-    dataset: DatasetRef
-    document_id: str
-
-
-class SegmentRef(NamedTuple):
-    """Segment identifiers used to scope downstream resource lookups."""
-
-    document: DocumentRef
-    segment_id: str
-
-
 class DatasetRefService:
     """Build child locators from validated dataset roots and resolve them with owner predicates."""
 
@@ -46,12 +24,12 @@ class DatasetRefService:
     @staticmethod
     def create_document_ref_from_id(dataset_ref: DatasetRef, document_id: str) -> DocumentRef:
         """Bind a candidate document ID; ownership is enforced when the ref is consumed."""
-        return DocumentRef(dataset=dataset_ref, document_id=document_id)
+        return dataset_ref.document(document_id)
 
     @staticmethod
     def create_segment_ref(document_ref: DocumentRef, segment_id: str) -> SegmentRef:
         """Bind a candidate segment ID; ownership is enforced when the ref is consumed."""
-        return SegmentRef(document=document_ref, segment_id=segment_id)
+        return document_ref.segment(segment_id)
 
     @staticmethod
     def get_document_by_ref(document_ref: DocumentRef, *, session: Session) -> Document | None:

api/services/dataset_service.py

@@ -7,10 +7,9 @@ import time
 import uuid
 from collections import Counter
 from collections.abc import Sequence
-from typing import Annotated, Any, Literal, TypedDict, cast
+from typing import Any, Literal, TypedDict, cast
 
 import sqlalchemy as sa
-from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
 from redis.exceptions import LockNotOwnedError
 from sqlalchemy import ColumnElement, delete, exists, func, select, update
 from sqlalchemy.orm import Session
@@ -20,6 +19,7 @@ from configs import dify_config
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.helper.name_generator import generate_incremental_name
 from core.model_manager import ModelManager
+from core.rag.entities.dataset_reference import DatasetRef, SegmentRef
 from core.rag.index_processor.constant.built_in_field import BuiltInField
 from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
@@ -65,11 +65,12 @@ from models.provider_ids import ModelProviderID
 from models.source import DataSourceOauthBinding
 from models.workflow import Workflow
 from services import dataset_api_key_service
-from services.dataset_ref_service import DatasetRef, DatasetRefService, SegmentRef
+from services.dataset_ref_service import DatasetRefService
 from services.document_indexing_proxy.document_indexing_task_proxy import DocumentIndexingTaskProxy
 from services.document_indexing_proxy.duplicate_document_indexing_task_proxy import DuplicateDocumentIndexingTaskProxy
 from services.enterprise import rbac_service as enterprise_rbac_service
 from services.entities.feature_entities import FeatureModel
+from services.entities.knowledge_entities.indexing_estimate import normalize_indexing_estimate_args
 from services.entities.knowledge_entities.knowledge_entities import (
     ChildChunkUpdateArgs,
     KnowledgeConfig,
@@ -121,121 +122,6 @@ class AutoDisableLogsDict(TypedDict):
     count: int
 
 
-class _EstimatePreProcessingRule(BaseModel):
-    id: str = Field(min_length=1)
-    enabled: bool
-
-    @field_validator("id")
-    @classmethod
-    def _validate_id(cls, v: str) -> str:
-        if v not in DatasetProcessRule.PRE_PROCESSING_RULES:
-            raise ValueError("Process rule pre_processing_rules id is invalid")
-        return v
-
-
-class _EstimateSegmentation(BaseModel):
-    separator: str = Field(min_length=1)
-    max_tokens: int = Field(gt=0)
-
-
-class _EstimateRules(BaseModel):
-    pre_processing_rules: list[_EstimatePreProcessingRule]
-    segmentation: _EstimateSegmentation
-
-    @field_validator("pre_processing_rules")
-    @classmethod
-    def _deduplicate(cls, v: list[_EstimatePreProcessingRule]) -> list[_EstimatePreProcessingRule]:
-        seen: dict[str, _EstimatePreProcessingRule] = {}
-        for rule in v:
-            seen[rule.id] = rule
-        return list(seen.values())
-
-
-class _EstimateHierarchicalRules(_EstimateRules):
-    parent_mode: Literal["full-doc", "paragraph"] | None = None
-    subchunk_segmentation: _EstimateSegmentation | None = None
-
-
-class _SummaryIndexSettingDisabled(BaseModel):
-    enable: Literal[False] = False
-
-
-class _SummaryIndexSettingEnabled(BaseModel):
-    enable: Literal[True]
-    model_name: str = Field(min_length=1)
-    model_provider_name: str = Field(min_length=1)
-
-
-_SummaryIndexSetting = Annotated[
-    _SummaryIndexSettingDisabled | _SummaryIndexSettingEnabled,
-    Field(discriminator="enable"),
-]
-
-
-class _AutomaticProcessRule(BaseModel):
-    model_config = ConfigDict(extra="allow")
-
-    mode: Literal[ProcessRuleMode.AUTOMATIC]
-    summary_index_setting: _SummaryIndexSetting | None = None
-
-    @field_validator("summary_index_setting", mode="before")
-    @classmethod
-    def _normalize_summary_index_setting(cls, v: Any) -> Any:
-        """Treat dicts with enable=None (or missing enable) as None (#36602)."""
-        if v is None:
-            return None
-        if isinstance(v, dict) and v.get("enable") is None:
-            return None
-        return v
-
-
-class _CustomProcessRule(BaseModel):
-    model_config = ConfigDict(extra="allow")
-
-    mode: Literal[ProcessRuleMode.CUSTOM]
-    rules: _EstimateRules
-    summary_index_setting: _SummaryIndexSetting | None = None
-
-    @field_validator("summary_index_setting", mode="before")
-    @classmethod
-    def _normalize_summary_index_setting(cls, v: Any) -> Any:
-        """Treat dicts with enable=None (or missing enable) as None (#36602)."""
-        if v is None:
-            return None
-        if isinstance(v, dict) and v.get("enable") is None:
-            return None
-        return v
-
-
-class _HierarchicalProcessRule(BaseModel):
-    model_config = ConfigDict(extra="allow")
-
-    mode: Literal[ProcessRuleMode.HIERARCHICAL]
-    rules: _EstimateHierarchicalRules
-    summary_index_setting: _SummaryIndexSetting | None = None
-
-    @field_validator("summary_index_setting", mode="before")
-    @classmethod
-    def _normalize_summary_index_setting(cls, v: Any) -> Any:
-        """Treat dicts with enable=None (or missing enable) as None (#36602)."""
-        if v is None:
-            return None
-        if isinstance(v, dict) and v.get("enable") is None:
-            return None
-        return v
-
-
-_EstimateProcessRule = Annotated[
-    _AutomaticProcessRule | _CustomProcessRule | _HierarchicalProcessRule,
-    Field(discriminator="mode"),
-]
-
-
-class _EstimateArgs(BaseModel):
-    info_list: dict[str, Any]
-    process_rule: _EstimateProcessRule
-
-
 class DatasetService:
     @staticmethod
     def _can_manage_all_datasets(tenant_id: str, account_id: str, *, session: Session) -> bool:
@@ -3176,20 +3062,7 @@ class DocumentService:
 
     @classmethod
     def estimate_args_validate(cls, args: dict[str, Any]):
-        try:
-            validated = _EstimateArgs.model_validate(args)
-        except ValidationError as e:
-            first = e.errors()[0]
-            original = first.get("ctx", {}).get("error")
-            raise ValueError(str(original) if isinstance(original, ValueError) else first["msg"]) from e
-        process_rule_dict = validated.process_rule.model_dump(exclude_none=True)
-        if validated.process_rule.mode == ProcessRuleMode.AUTOMATIC:
-            process_rule_dict["rules"] = {}
-        elif validated.process_rule.mode == ProcessRuleMode.HIERARCHICAL:
-            rules = process_rule_dict.get("rules")
-            if isinstance(rules, dict) and not rules.get("parent_mode"):
-                rules["parent_mode"] = "paragraph"
-        args["process_rule"] = process_rule_dict
+        args["process_rule"] = normalize_indexing_estimate_args(args)
 
     @staticmethod
     def batch_update_document_status(

api/services/datasource_provider_service.py

@@ -851,6 +851,7 @@ class DatasourceProviderService:
         if user is not None:
             query = CredentialPermissionService.apply_visibility_filter(
                 query,
+                tenant_id=tenant_id,
                 model_id_column=DatasourceProvider.id,
                 model_user_id_column=DatasourceProvider.user_id,
                 model_visibility_column=DatasourceProvider.visibility,

api/services/entities/data_source/__init__.py

@@ -0,0 +1 @@
+"""Shared data contracts for data-source use cases."""

api/services/entities/data_source/credential.py

@@ -0,0 +1,18 @@
+"""Detached data contracts for datasource credentials."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass(frozen=True, slots=True)
+class DatasourceCredentialRecord:
+    id: str
+    workspace_id: str
+    name: str
+    provider: str
+    plugin_id: str
+    auth_type: str
+    encrypted_credentials: Mapping[str, object]
+    expires_at: int
+    updated_at: datetime

api/services/entities/data_source/notion_import.py

@@ -0,0 +1,66 @@
+"""Framework-neutral values used by Notion import use cases."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class NotionPageIcon:
+    type: str | None = None
+    url: str | None = None
+    emoji: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class AuthorizedNotionPage:
+    page_id: str
+    page_name: str
+    page_icon: NotionPageIcon | None
+    parent_id: str | None
+    page_type: str
+
+
+@dataclass(frozen=True, slots=True)
+class NotionImportPage:
+    page_id: str
+    page_name: str
+    page_icon: NotionPageIcon | None
+    parent_id: str | None
+    page_type: str
+    is_bound: bool
+
+
+@dataclass(frozen=True, slots=True)
+class NotionWorkspace:
+    workspace_id: str | None
+    workspace_name: str | None
+    workspace_icon: str | None
+    pages: tuple[AuthorizedNotionPage, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class NotionImportWorkspace:
+    workspace_id: str | None
+    workspace_name: str | None
+    workspace_icon: str | None
+    pages: tuple[NotionImportPage, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class NotionImportResult:
+    workspaces: tuple[NotionImportWorkspace, ...]
+
+
+def notion_page_icon(value: Mapping[str, object] | None) -> NotionPageIcon | None:
+    """Map a plugin icon payload without leaking the plugin entity type."""
+
+    if value is None:
+        return None
+    icon_type = value.get("type")
+    url = value.get("url")
+    emoji = value.get("emoji")
+    return NotionPageIcon(
+        type=icon_type if isinstance(icon_type, str) else None,
+        url=url if isinstance(url, str) else None,
+        emoji=emoji if isinstance(emoji, str) else None,
+    )

api/services/entities/data_source/oauth.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Mapping
 from dataclasses import dataclass
+from datetime import datetime
 
 
 @dataclass(frozen=True, slots=True)
@@ -22,3 +23,14 @@ class DataSourceOAuthBindingRecord:
     id: str
     access_token: str
     source_info: Mapping[str, object]
+
+
+@dataclass(frozen=True, slots=True)
+class DataSourceBindingSummary:
+    """Secret-free management view of a data-source OAuth binding."""
+
+    id: str
+    provider: str
+    created_at: datetime
+    disabled: bool
+    source_info: Mapping[str, object]

api/services/entities/knowledge_entities/__init__.py

@@ -0,0 +1 @@
+"""Shared data contracts for knowledge use cases."""

api/services/entities/knowledge_entities/indexing_estimate.py

@@ -0,0 +1,183 @@
+"""Commands and validation shared by knowledge indexing-estimate use cases."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator
+
+from core.rag.entities.processing_entities import PreProcessingRuleKey
+from services.entities.knowledge_entities.records import DatasetRecord, DocumentRecord
+
+
+class EstimateValidationError(ValueError):
+    """Raised when an estimate process rule is invalid."""
+
+
+class _PreProcessingRule(BaseModel):
+    id: PreProcessingRuleKey
+    enabled: bool
+
+
+class _Segmentation(BaseModel):
+    separator: str = Field(min_length=1)
+    max_tokens: int = Field(gt=0)
+
+
+class _Rules(BaseModel):
+    pre_processing_rules: list[_PreProcessingRule]
+    segmentation: _Segmentation
+
+    @field_validator("pre_processing_rules")
+    @classmethod
+    def _deduplicate(cls, value: list[_PreProcessingRule]) -> list[_PreProcessingRule]:
+        by_id = {rule.id: rule for rule in value}
+        return list(by_id.values())
+
+
+class _HierarchicalRules(_Rules):
+    parent_mode: Literal["full-doc", "paragraph"] | None = None
+    subchunk_segmentation: _Segmentation | None = None
+
+
+class _SummaryIndexSettingDisabled(BaseModel):
+    enable: Literal[False] = False
+
+
+class _SummaryIndexSettingEnabled(BaseModel):
+    enable: Literal[True]
+    model_name: str = Field(min_length=1)
+    model_provider_name: str = Field(min_length=1)
+
+
+_SummaryIndexSetting = Annotated[
+    _SummaryIndexSettingDisabled | _SummaryIndexSettingEnabled,
+    Field(discriminator="enable"),
+]
+
+
+class _ProcessRuleBase(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    summary_index_setting: _SummaryIndexSetting | None = None
+
+    @field_validator("summary_index_setting", mode="before")
+    @classmethod
+    def _normalize_summary_index_setting(cls, value: object) -> object:
+        """Treat mappings with a missing/None enable flag as absent (#36602)."""
+
+        if value is None:
+            return None
+        if isinstance(value, Mapping) and value.get("enable") is None:
+            return None
+        return value
+
+
+class _AutomaticProcessRule(_ProcessRuleBase):
+    mode: Literal["automatic"]
+
+
+class _CustomProcessRule(_ProcessRuleBase):
+    mode: Literal["custom"]
+    rules: _Rules
+
+
+class _HierarchicalProcessRule(_ProcessRuleBase):
+    mode: Literal["hierarchical"]
+    rules: _HierarchicalRules
+
+
+_ProcessRule = Annotated[
+    _AutomaticProcessRule | _CustomProcessRule | _HierarchicalProcessRule,
+    Field(discriminator="mode"),
+]
+
+
+class _EstimateArgs(BaseModel):
+    info_list: dict[str, object]
+    process_rule: _ProcessRule
+
+
+class _ProcessRuleEnvelope(BaseModel):
+    process_rule: _ProcessRule
+
+
+def _validation_message(error: ValidationError) -> str:
+    first = error.errors()[0]
+    original = first.get("ctx", {}).get("error")
+    return str(original) if isinstance(original, ValueError) else str(first["msg"])
+
+
+def normalize_process_rule(process_rule: Mapping[str, object]) -> dict[str, object]:
+    """Validate and normalize a process rule without mutating caller-owned input."""
+
+    try:
+        validated = _ProcessRuleEnvelope.model_validate({"process_rule": process_rule})
+    except ValidationError as error:
+        raise EstimateValidationError(_validation_message(error)) from error
+
+    result = validated.process_rule.model_dump(exclude_none=True)
+    if validated.process_rule.mode == "automatic":
+        result["rules"] = {}
+    elif validated.process_rule.mode == "hierarchical":
+        rules = result.get("rules")
+        if isinstance(rules, dict) and not rules.get("parent_mode"):
+            rules["parent_mode"] = "paragraph"
+    return result
+
+
+def normalize_indexing_estimate_args(args: Mapping[str, object]) -> dict[str, object]:
+    """Compatibility validator for the existing generic estimate transport."""
+
+    try:
+        validated = _EstimateArgs.model_validate(args)
+    except ValidationError as error:
+        raise EstimateValidationError(_validation_message(error)) from error
+    return normalize_process_rule(validated.process_rule.model_dump(exclude_none=True))
+
+
+@dataclass(frozen=True, slots=True)
+class UploadFileEstimateSource:
+    file_id: str
+    source_type: Literal["upload_file"] = field(default="upload_file", init=False)
+
+
+@dataclass(frozen=True, slots=True)
+class NotionEstimateSource:
+    workspace_id: str
+    page_id: str
+    page_type: str
+    credential_id: str
+    source_type: Literal["notion_import"] = field(default="notion_import", init=False)
+
+
+@dataclass(frozen=True, slots=True)
+class WebsiteEstimateSource:
+    provider: str
+    job_id: str
+    url: str
+    mode: str = "crawl"
+    only_main_content: bool = False
+    source_type: Literal["website_crawl"] = field(default="website_crawl", init=False)
+
+
+type NewEstimateSource = UploadFileEstimateSource | NotionEstimateSource | WebsiteEstimateSource
+
+
+@dataclass(frozen=True, slots=True)
+class NewSourcesEstimateCommand:
+    sources: tuple[NewEstimateSource, ...]
+    process_rule: Mapping[str, object]
+    doc_form: str = "text_model"
+    doc_language: str = "English"
+    dataset_id: str | None = None
+    indexing_technique: str = "economy"
+
+
+@dataclass(frozen=True, slots=True)
+class ExistingDocumentsEstimateCommand:
+    dataset: DatasetRecord
+    documents: tuple[DocumentRecord, ...]
+
+
+type EstimateCommand = NewSourcesEstimateCommand | ExistingDocumentsEstimateCommand

api/services/entities/knowledge_entities/records.py

@@ -0,0 +1,44 @@
+"""Immutable records shared by knowledge application services."""
+
+from collections.abc import Mapping
+from typing import NamedTuple
+
+
+class DatasetRecord(NamedTuple):
+    """Framework- and ORM-independent dataset state used at application boundaries."""
+
+    id: str
+    workspace_id: str
+    maintainer_id: str | None
+    permission: str
+    data_source_type: str | None
+    indexing_technique: str | None
+    embedding_model: str | None
+    embedding_model_provider: str | None
+
+
+class DatasetAccessSnapshot(NamedTuple):
+    """Dataset state plus the requesting account's partial-member grant."""
+
+    dataset: DatasetRecord
+    actor_has_partial_access: bool
+
+
+class DocumentRecord(NamedTuple):
+    """Document state detached from its SQLAlchemy session."""
+
+    id: str
+    workspace_id: str
+    dataset_id: str
+    name: str
+    data_source_type: str
+    data_source_info: Mapping[str, object] | None
+    enabled: bool
+    archived: bool
+    indexing_status: str
+    batch: str
+    doc_form: str
+    doc_language: str | None
+    dataset_process_rule_id: str | None
+    need_summary: bool
+    doc_metadata: Mapping[str, object] | None

api/services/knowledge/__init__.py

@@ -0,0 +1 @@
+"""Application services and adapters for knowledge use cases."""

api/services/knowledge/adapters.py

@@ -0,0 +1,364 @@
+"""Infrastructure adapters for knowledge application services."""
+
+import json
+from collections.abc import Callable, Mapping
+from typing import Protocol
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.entities.knowledge_entities import IndexingEstimate
+from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
+from core.indexing_runner import IndexingRunner
+from core.plugin.impl.exc import PluginDaemonClientSideError
+from core.rag.extractor.entity.datasource_type import DatasourceType
+from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
+from models.dataset import DatasetProcessRule
+from models.model import UploadFile
+from services.entities.knowledge_entities.indexing_estimate import (
+    EstimateCommand,
+    ExistingDocumentsEstimateCommand,
+    NewSourcesEstimateCommand,
+    NotionEstimateSource,
+    UploadFileEstimateSource,
+    WebsiteEstimateSource,
+)
+from services.entities.knowledge_entities.records import DocumentRecord
+from tasks.document_indexing_sync_task import document_indexing_sync_task
+
+_NOTION_PROVIDER = "notion_datasource"
+_NOTION_PLUGIN = "langgenius/notion_datasource"
+
+
+class CeleryDocumentSyncDispatcher:
+    def __init__(self, delay: Callable[[str, str], object] | None = None) -> None:
+        self._delay = delay or document_indexing_sync_task.delay
+
+    def dispatch(self, *, dataset_id: str, document_id: str) -> None:
+        self._delay(dataset_id, document_id)
+
+
+class _EstimateRunner(Protocol):
+    def indexing_estimate(
+        self,
+        tenant_id: str,
+        extract_settings: list[ExtractSetting],
+        tmp_processing_rule: Mapping[str, object],
+        doc_form: str | None = None,
+        doc_language: str = "English",
+        dataset_id: str | None = None,
+        indexing_technique: str = "economy",
+        *,
+        session: Session,
+    ) -> IndexingEstimate: ...
+
+
+class _CredentialResolver(Protocol):
+    def resolve(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> dict[str, object]: ...
+
+
+class IndexingEstimateAdapterError(Exception):
+    """Base class for failures produced by the estimate infrastructure adapter."""
+
+
+class EstimateSourceNotFoundError(IndexingEstimateAdapterError):
+    def __init__(self, source_id: str) -> None:
+        super().__init__(f"Estimate source not found: {source_id}")
+
+
+class UnsupportedEstimateSourceError(IndexingEstimateAdapterError):
+    def __init__(self, source_type: str) -> None:
+        super().__init__(f"Data source type not supported: {source_type}")
+
+
+class IndexingEstimateProviderUnavailableError(IndexingEstimateAdapterError):
+    """Raised when a configured model or plugin credential is unavailable."""
+
+
+class IndexingEstimateExecutionError(IndexingEstimateAdapterError):
+    """Raised when the existing runner cannot complete an estimate."""
+
+
+class IndexingRunnerEstimateGateway:
+    """Keep the runner's caller-owned Session requirement inside one adapter."""
+
+    def __init__(
+        self,
+        *,
+        session_factory: sessionmaker[Session],
+        credentials: _CredentialResolver,
+        runner_factory: Callable[[], _EstimateRunner] = IndexingRunner,
+    ) -> None:
+        self._session_factory = session_factory
+        self._credentials = credentials
+        self._runner_factory = runner_factory
+
+    def estimate(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        command: EstimateCommand,
+    ) -> IndexingEstimate:
+        notion_tokens = self._resolve_notion_tokens(
+            workspace_id=workspace_id,
+            actor_id=actor_id,
+            command=command,
+        )
+        with self._session_factory() as session:
+            if isinstance(command, NewSourcesEstimateCommand):
+                extract_settings = self._new_source_settings(session, workspace_id, command, notion_tokens)
+                process_rule = command.process_rule
+                doc_form = command.doc_form
+                doc_language = command.doc_language
+                dataset_id = command.dataset_id
+                indexing_technique = command.indexing_technique
+            else:
+                extract_settings = self._document_settings(
+                    session,
+                    workspace_id,
+                    command.documents,
+                    notion_tokens,
+                )
+                process_rule = self._load_process_rule(session, command)
+                first_document = command.documents[0] if command.documents else None
+                doc_form = first_document.doc_form if first_document else "text_model"
+                doc_language = first_document.doc_language or "English" if first_document else "English"
+                dataset_id = command.dataset.id
+                indexing_technique = command.dataset.indexing_technique or "economy"
+
+            if not extract_settings:
+                return IndexingEstimate(total_segments=0, preview=[])
+
+            try:
+                return self._runner_factory().indexing_estimate(
+                    tenant_id=workspace_id,
+                    extract_settings=extract_settings,
+                    tmp_processing_rule=process_rule,
+                    doc_form=doc_form,
+                    doc_language=doc_language,
+                    dataset_id=dataset_id,
+                    indexing_technique=indexing_technique,
+                    session=session,
+                )
+            except (LLMBadRequestError, ProviderTokenNotInitError, PluginDaemonClientSideError) as error:
+                description = getattr(error, "description", str(error))
+                raise IndexingEstimateProviderUnavailableError(description) from error
+            except Exception as error:
+                raise IndexingEstimateExecutionError(str(error)) from error
+
+    @staticmethod
+    def _new_source_settings(
+        session: Session,
+        workspace_id: str,
+        command: NewSourcesEstimateCommand,
+        notion_tokens: Mapping[str, str],
+    ) -> list[ExtractSetting]:
+        file_ids = tuple(source.file_id for source in command.sources if isinstance(source, UploadFileEstimateSource))
+        files_by_id: dict[str, UploadFile] = {}
+        if file_ids:
+            files = session.scalars(
+                select(UploadFile).where(UploadFile.tenant_id == workspace_id, UploadFile.id.in_(file_ids))
+            ).all()
+            files_by_id = {file.id: file for file in files}
+
+        settings: list[ExtractSetting] = []
+        for source in command.sources:
+            if isinstance(source, UploadFileEstimateSource):
+                upload_file = files_by_id.get(source.file_id)
+                if upload_file is None:
+                    raise EstimateSourceNotFoundError(source.file_id)
+                settings.append(
+                    ExtractSetting(
+                        datasource_type=DatasourceType.FILE,
+                        upload_file=upload_file,
+                        document_model=command.doc_form,
+                    )
+                )
+            elif isinstance(source, NotionEstimateSource):
+                settings.append(
+                    ExtractSetting(
+                        datasource_type=DatasourceType.NOTION,
+                        notion_info=NotionInfo(
+                            credential_id=source.credential_id,
+                            notion_workspace_id=source.workspace_id,
+                            notion_obj_id=source.page_id,
+                            notion_page_type=source.page_type,
+                            tenant_id=workspace_id,
+                            notion_access_token=notion_tokens[source.credential_id],
+                        ),
+                        document_model=command.doc_form,
+                    )
+                )
+            elif isinstance(source, WebsiteEstimateSource):
+                settings.append(
+                    ExtractSetting(
+                        datasource_type=DatasourceType.WEBSITE,
+                        website_info=WebsiteInfo(
+                            provider=source.provider,
+                            job_id=source.job_id,
+                            url=source.url,
+                            mode=source.mode,
+                            tenant_id=workspace_id,
+                            only_main_content=source.only_main_content,
+                        ),
+                        document_model=command.doc_form,
+                    )
+                )
+        return settings
+
+    def _document_settings(
+        self,
+        session: Session,
+        workspace_id: str,
+        documents: tuple[DocumentRecord, ...],
+        notion_tokens: Mapping[str, str],
+    ) -> list[ExtractSetting]:
+        upload_ids = tuple(
+            source_id
+            for document in documents
+            if document.data_source_type == "upload_file"
+            if (source_id := self._string_value(document.data_source_info, "upload_file_id")) is not None
+        )
+        files_by_id: dict[str, UploadFile] = {}
+        if upload_ids:
+            files = session.scalars(
+                select(UploadFile).where(UploadFile.tenant_id == workspace_id, UploadFile.id.in_(upload_ids))
+            ).all()
+            files_by_id = {file.id: file for file in files}
+
+        settings: list[ExtractSetting] = []
+        for document in documents:
+            info = document.data_source_info
+            if not info:
+                continue
+            match document.data_source_type:
+                case "upload_file":
+                    file_id = self._required_string(info, "upload_file_id")
+                    upload_file = files_by_id.get(file_id)
+                    if upload_file is None:
+                        raise EstimateSourceNotFoundError(file_id)
+                    settings.append(
+                        ExtractSetting(
+                            datasource_type=DatasourceType.FILE,
+                            upload_file=upload_file,
+                            document_model=document.doc_form,
+                        )
+                    )
+                case "notion_import":
+                    credential_id = self._required_string(info, "credential_id")
+                    settings.append(
+                        ExtractSetting(
+                            datasource_type=DatasourceType.NOTION,
+                            notion_info=NotionInfo(
+                                credential_id=credential_id,
+                                notion_workspace_id=self._required_string(info, "notion_workspace_id"),
+                                notion_obj_id=self._required_string(info, "notion_page_id"),
+                                notion_page_type=self._required_string(info, "type"),
+                                tenant_id=workspace_id,
+                                notion_access_token=notion_tokens[credential_id],
+                            ),
+                            document_model=document.doc_form,
+                        )
+                    )
+                case "website_crawl":
+                    settings.append(
+                        ExtractSetting(
+                            datasource_type=DatasourceType.WEBSITE,
+                            website_info=WebsiteInfo(
+                                provider=self._required_string(info, "provider"),
+                                job_id=self._required_string(info, "job_id"),
+                                url=self._required_string(info, "url"),
+                                mode=self._required_string(info, "mode"),
+                                tenant_id=workspace_id,
+                                only_main_content=bool(info.get("only_main_content", False)),
+                            ),
+                            document_model=document.doc_form,
+                        )
+                    )
+                case _:
+                    raise UnsupportedEstimateSourceError(document.data_source_type)
+        return settings
+
+    def _resolve_notion_tokens(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        command: EstimateCommand,
+    ) -> dict[str, str]:
+        if isinstance(command, NewSourcesEstimateCommand):
+            credential_ids = tuple(
+                dict.fromkeys(
+                    source.credential_id for source in command.sources if isinstance(source, NotionEstimateSource)
+                )
+            )
+        else:
+            credential_ids = tuple(
+                dict.fromkeys(
+                    self._required_string(document.data_source_info or {}, "credential_id")
+                    for document in command.documents
+                    if document.data_source_type == "notion_import"
+                )
+            )
+
+        tokens: dict[str, str] = {}
+        for credential_id in credential_ids:
+            credentials = self._credentials.resolve(
+                workspace_id=workspace_id,
+                actor_id=actor_id,
+                credential_id=credential_id,
+                provider=_NOTION_PROVIDER,
+                plugin_id=_NOTION_PLUGIN,
+            )
+            integration_secret = credentials.get("integration_secret")
+            if not isinstance(integration_secret, str) or not integration_secret:
+                raise EstimateSourceNotFoundError(credential_id)
+            tokens[credential_id] = integration_secret
+        return tokens
+
+    @staticmethod
+    def _load_process_rule(
+        session: Session,
+        command: ExistingDocumentsEstimateCommand,
+    ) -> Mapping[str, object]:
+        process_rule_id = next(
+            (document.dataset_process_rule_id for document in command.documents if document.dataset_process_rule_id),
+            None,
+        )
+        if process_rule_id is None:
+            return {"mode": "automatic", "rules": {}}
+        process_rule = session.scalar(
+            select(DatasetProcessRule)
+            .where(
+                DatasetProcessRule.id == process_rule_id,
+                DatasetProcessRule.dataset_id == command.dataset.id,
+            )
+            .limit(1)
+        )
+        if process_rule is None:
+            raise EstimateSourceNotFoundError(process_rule_id)
+        rules = json.loads(process_rule.rules) if process_rule.rules else {}
+        return {"mode": str(process_rule.mode), "rules": rules}
+
+    @staticmethod
+    def _string_value(info: Mapping[str, object] | None, key: str) -> str | None:
+        if info is None:
+            return None
+        value = info.get(key)
+        return value if isinstance(value, str) and value else None
+
+    @classmethod
+    def _required_string(cls, info: Mapping[str, object], key: str) -> str:
+        value = cls._string_value(info, key)
+        if value is None:
+            raise EstimateSourceNotFoundError(key)
+        return value

api/services/knowledge/application.py

@@ -0,0 +1,268 @@
+"""Framework-neutral application services for knowledge use cases."""
+
+from typing import Protocol
+
+from core.entities.knowledge_entities import IndexingEstimate
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef
+from machinery.context import RequestContext
+from services.entities.knowledge_entities.indexing_estimate import (
+    EstimateCommand,
+    ExistingDocumentsEstimateCommand,
+    NewSourcesEstimateCommand,
+    normalize_process_rule,
+)
+from services.entities.knowledge_entities.records import DatasetAccessSnapshot, DatasetRecord, DocumentRecord
+
+_OWNER_ROLE = "owner"
+_ONLY_ME = "only_me"
+_ALL_TEAM_MEMBERS = "all_team_members"
+_PARTIAL_MEMBERS = "partial_members"
+
+
+class DatasetAccessReader(Protocol):
+    """Load tenant-scoped dataset state needed by the access policy."""
+
+    def get_access_snapshot(
+        self,
+        *,
+        workspace_id: str,
+        dataset_id: str,
+        actor_id: str,
+    ) -> DatasetAccessSnapshot | None: ...
+
+
+class WorkspaceRoleReader(Protocol):
+    """Read the account's legacy role in a workspace."""
+
+    def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None: ...
+
+
+class DatasetAccessError(Exception):
+    """Base class for framework-neutral dataset access failures."""
+
+
+class DatasetNotFoundError(DatasetAccessError):
+    """Raised when a dataset does not belong to the active workspace."""
+
+    def __init__(self) -> None:
+        super().__init__("Dataset not found")
+
+
+class DatasetAccessDeniedError(DatasetAccessError):
+    """Raised when legacy dataset permission rules deny the actor."""
+
+    def __init__(self) -> None:
+        super().__init__("You do not have permission to access this dataset")
+
+
+def can_access_dataset(snapshot: DatasetAccessSnapshot, *, actor_id: str, workspace_role: str | None) -> bool:
+    """Evaluate legacy dataset visibility without database or framework access."""
+
+    dataset = snapshot.dataset
+    if workspace_role == _OWNER_ROLE or dataset.maintainer_id == actor_id:
+        return True
+    if dataset.permission == _ALL_TEAM_MEMBERS:
+        return True
+    if dataset.permission == _PARTIAL_MEMBERS:
+        return snapshot.actor_has_partial_access
+    if dataset.permission == _ONLY_ME:
+        return False
+    return False
+
+
+class DatasetAccessService:
+    """Resolve an owned dataset and enforce legacy permissions when configured."""
+
+    def __init__(
+        self,
+        *,
+        datasets: DatasetAccessReader,
+        workspace_roles: WorkspaceRoleReader,
+        legacy_permissions_enabled: bool,
+    ) -> None:
+        self._datasets = datasets
+        self._workspace_roles = workspace_roles
+        self._legacy_permissions_enabled = legacy_permissions_enabled
+
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord:
+        snapshot = self._datasets.get_access_snapshot(
+            workspace_id=context.active_workspace_id,
+            dataset_id=dataset_id,
+            actor_id=context.account_id,
+        )
+        if snapshot is None:
+            raise DatasetNotFoundError()
+        if not self._legacy_permissions_enabled:
+            return snapshot.dataset
+
+        workspace_role = self._workspace_roles.get_legacy_role(
+            workspace_id=context.active_workspace_id,
+            account_id=context.account_id,
+        )
+        if can_access_dataset(snapshot, actor_id=context.account_id, workspace_role=workspace_role):
+            return snapshot.dataset
+        raise DatasetAccessDeniedError()
+
+
+class DatasetAccess(Protocol):
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord: ...
+
+
+class SyncDocumentReader(Protocol):
+    def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]: ...
+
+    def get_by_ref(self, document_ref: DocumentRef) -> DocumentRecord | None: ...
+
+
+class DocumentSyncDispatcher(Protocol):
+    def dispatch(self, *, dataset_id: str, document_id: str) -> None: ...
+
+
+class DocumentSyncError(Exception):
+    """Base class for document synchronization failures."""
+
+
+class SyncDocumentNotFoundError(DocumentSyncError):
+    def __init__(self) -> None:
+        super().__init__("Document not found")
+
+
+class SyncDocumentSourceError(DocumentSyncError):
+    def __init__(self) -> None:
+        super().__init__("Document is not notion type")
+
+
+class DocumentSyncApplicationService:
+    def __init__(
+        self,
+        *,
+        dataset_access: DatasetAccess,
+        documents: SyncDocumentReader,
+        dispatcher: DocumentSyncDispatcher,
+    ) -> None:
+        self._dataset_access = dataset_access
+        self._documents = documents
+        self._dispatcher = dispatcher
+
+    def sync_dataset(self, context: RequestContext, dataset_id: str) -> int:
+        dataset = self._dataset_access.require_accessible(context, dataset_id)
+        dataset_ref = DatasetRef(context.active_workspace_id, dataset.id)
+        document_refs = self._documents.list_active_notion_refs(dataset_ref)
+        for document_ref in document_refs:
+            self._dispatcher.dispatch(dataset_id=dataset_ref.dataset_id, document_id=document_ref.document_id)
+        return len(document_refs)
+
+    def sync_document(self, context: RequestContext, dataset_id: str, document_id: str) -> None:
+        dataset = self._dataset_access.require_accessible(context, dataset_id)
+        document_ref = DatasetRef(context.active_workspace_id, dataset.id).document(document_id)
+        document = self._documents.get_by_ref(document_ref)
+        if document is None:
+            raise SyncDocumentNotFoundError()
+        if document.data_source_type != "notion_import":
+            raise SyncDocumentSourceError()
+        self._dispatcher.dispatch(dataset_id=dataset.id, document_id=document.id)
+
+
+class EstimateDocumentReader(Protocol):
+    def get_by_ref(self, document_ref: DocumentRef) -> DocumentRecord | None: ...
+
+    def list_by_batch(self, dataset_ref: DatasetRef, batch: str) -> tuple[DocumentRecord, ...]: ...
+
+
+class IndexingEstimateGateway(Protocol):
+    def estimate(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        command: EstimateCommand,
+    ) -> IndexingEstimate: ...
+
+
+class IndexingEstimateApplicationError(Exception):
+    """Base class for framework-neutral estimate failures."""
+
+
+class EstimateDocumentNotFoundError(IndexingEstimateApplicationError):
+    def __init__(self) -> None:
+        super().__init__("Document not found")
+
+
+class EstimateDocumentAlreadyFinishedError(IndexingEstimateApplicationError):
+    def __init__(self) -> None:
+        super().__init__("Document already finished")
+
+
+class IndexingEstimateApplicationService:
+    def __init__(
+        self,
+        *,
+        dataset_access: DatasetAccess,
+        documents: EstimateDocumentReader,
+        gateway: IndexingEstimateGateway,
+    ) -> None:
+        self._dataset_access = dataset_access
+        self._documents = documents
+        self._gateway = gateway
+
+    def estimate_new_sources(
+        self,
+        context: RequestContext,
+        command: NewSourcesEstimateCommand,
+    ) -> IndexingEstimate:
+        dataset_id = command.dataset_id
+        if dataset_id is not None:
+            dataset = self._dataset_access.require_accessible(context, dataset_id)
+            dataset_id = dataset.id
+        normalized = NewSourcesEstimateCommand(
+            sources=command.sources,
+            process_rule=normalize_process_rule(command.process_rule),
+            doc_form=command.doc_form,
+            doc_language=command.doc_language,
+            dataset_id=dataset_id,
+            indexing_technique=command.indexing_technique,
+        )
+        return self._gateway.estimate(
+            workspace_id=context.active_workspace_id,
+            actor_id=context.account_id,
+            command=normalized,
+        )
+
+    def estimate_document(
+        self,
+        context: RequestContext,
+        *,
+        dataset_id: str,
+        document_id: str,
+    ) -> IndexingEstimate:
+        dataset = self._dataset_access.require_accessible(context, dataset_id)
+        document = self._documents.get_by_ref(DatasetRef(context.active_workspace_id, dataset.id).document(document_id))
+        if document is None:
+            raise EstimateDocumentNotFoundError()
+        self._require_unfinished((document,))
+        return self._estimate_documents(context, dataset, (document,))
+
+    def estimate_batch(self, context: RequestContext, *, dataset_id: str, batch: str) -> IndexingEstimate:
+        dataset = self._dataset_access.require_accessible(context, dataset_id)
+        documents = self._documents.list_by_batch(DatasetRef(context.active_workspace_id, dataset.id), batch)
+        self._require_unfinished(documents)
+        return self._estimate_documents(context, dataset, documents)
+
+    def _estimate_documents(
+        self,
+        context: RequestContext,
+        dataset: DatasetRecord,
+        documents: tuple[DocumentRecord, ...],
+    ) -> IndexingEstimate:
+        if not documents:
+            return IndexingEstimate(total_segments=0, preview=[])
+        return self._gateway.estimate(
+            workspace_id=context.active_workspace_id,
+            actor_id=context.account_id,
+            command=ExistingDocumentsEstimateCommand(dataset=dataset, documents=documents),
+        )
+
+    @staticmethod
+    def _require_unfinished(documents: tuple[DocumentRecord, ...]) -> None:
+        if any(document.indexing_status in {"completed", "error"} for document in documents):
+            raise EstimateDocumentAlreadyFinishedError()

api/services/rag_pipeline/pipeline_generate_service.py

@@ -6,11 +6,12 @@ from sqlalchemy.orm import Session
 from configs import dify_config
 from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
 from core.app.entities.app_invoke_entities import InvokeFrom
+from core.rag.entities.dataset_reference import DocumentRef
 from models.dataset import Pipeline
 from models.enums import IndexingStatus
 from models.model import Account, App, EndUser
 from models.workflow import Workflow
-from services.dataset_ref_service import DatasetRefService, DocumentRef
+from services.dataset_ref_service import DatasetRefService
 from services.rag_pipeline.rag_pipeline import RagPipelineService
 
 

api/services/tools/builtin_tools_manage_service.py

@@ -377,6 +377,7 @@ class BuiltinToolManageService:
             if user is not None:
                 visible_query = CredentialPermissionService.apply_visibility_filter(
                     visible_query,
+                    tenant_id=tenant_id,
                     model_id_column=BuiltinToolProvider.id,
                     model_user_id_column=BuiltinToolProvider.user_id,
                     model_visibility_column=BuiltinToolProvider.visibility,

api/services/trigger/trigger_provider_service.py

@@ -93,6 +93,7 @@ class TriggerProviderService:
             if user is not None:
                 query = CredentialPermissionService.apply_visibility_filter(
                     query,
+                    tenant_id=tenant_id,
                     model_id_column=TriggerSubscription.id,
                     model_user_id_column=TriggerSubscription.user_id,
                     model_visibility_column=TriggerSubscription.visibility,

api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_bearer_auth.py

@@ -9,7 +9,7 @@ from sqlalchemy.orm import Session
 
 from models.source import DataSourceApiKeyAuthBinding
 from services.auth.errors import DataSourceApiKeyAuthCredentialValidationError
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthCredentials,
 )

api/tests/test_containers_integration_tests/controllers/console/auth/test_data_source_oauth.py

@@ -7,8 +7,8 @@ from sqlalchemy.orm import Session
 
 from extensions.ext_application_services import ApplicationServices
 from models.source import DataSourceOauthBinding
-from services.data_source_oauth_service import InvalidDataSourceOAuthProviderError
-from services.entities.data_source_oauth_entities import DataSourceOAuthCallback
+from services.data_source.oauth_service import InvalidDataSourceOAuthProviderError
+from services.entities.data_source.oauth import DataSourceOAuthCallback
 from tests.test_containers_integration_tests.controllers.console.helpers import (
     authenticate_console_client,
     create_console_account_and_tenant,
@@ -17,7 +17,7 @@ from tests.test_containers_integration_tests.controllers.console.helpers import
 
 def _services(service: MagicMock) -> MagicMock:
     services = create_autospec(ApplicationServices, instance=True, spec_set=True)
-    services.resolve_data_source_oauth.return_value = service
+    services.data_sources.resolve_oauth.return_value = service
     return services
 
 
@@ -54,7 +54,7 @@ def test_get_oauth_url_invalid_provider(
 
     service = MagicMock()
     services = _services(service)
-    services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError
+    services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError
     with patch(
         "controllers.console.auth.data_source_oauth.application_services",
         return_value=services,
@@ -99,7 +99,7 @@ def test_oauth_callback_missing_code(test_client_with_containers: FlaskClient) -
 def test_oauth_callback_invalid_provider(test_client_with_containers: FlaskClient) -> None:
     service = MagicMock()
     services = _services(service)
-    services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError
+    services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError
     with patch(
         "controllers.console.auth.data_source_oauth.application_services",
         return_value=services,

api/tests/test_containers_integration_tests/services/auth/test_auth_integration.py

@@ -12,14 +12,14 @@ from sqlalchemy.orm import Session, sessionmaker
 
 from machinery.context import RequestContext
 from models.source import DataSourceApiKeyAuthBinding
-from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
+from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
 from services.auth.data_source_api_key_auth_gateways import (
     ProviderApiKeyAuthCredentialValidator,
     TenantApiKeyAuthCredentialEncryptor,
 )
 from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService
 from services.auth.errors import DataSourceApiKeyAuthProviderUnavailableError
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthCredentials,
 )

api/tests/test_containers_integration_tests/services/test_dataset_service_permissions.py

@@ -9,6 +9,7 @@ import pytest
 from flask import Flask
 from sqlalchemy.orm import Session
 
+from core.rag.entities.dataset_reference import DatasetRef
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
 from models.dataset import (
@@ -20,7 +21,7 @@ from models.dataset import (
     DatasetPermissionEnum,
 )
 from models.enums import DataSourceType
-from services.dataset_ref_service import DatasetRef, DatasetRefService
+from services.dataset_ref_service import DatasetRefService
 from services.dataset_service import DatasetCollectionBindingService, DatasetPermissionService, DatasetService
 from services.errors.account import NoPermissionError
 

api/tests/unit_tests/controllers/console/auth/test_data_source_bearer_auth.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from inspect import unwrap
-from types import SimpleNamespace
 from unittest.mock import create_autospec, patch
 from uuid import UUID
 
@@ -30,7 +30,7 @@ from services.auth.errors import (
     InvalidDataSourceApiKeyAuthCredentialsError,
     UnsupportedDataSourceApiKeyAuthProviderError,
 )
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthBindingRecord,
     DataSourceApiKeyAuthCredentials,
@@ -46,8 +46,18 @@ def _request_context() -> RequestContext:
     )
 
 
-def _application_services(service: DataSourceApiKeyAuthService) -> SimpleNamespace:
-    return SimpleNamespace(data_source_api_key_auth=service)
+@dataclass(frozen=True)
+class _DataSources:
+    api_key_auth: DataSourceApiKeyAuthService
+
+
+@dataclass(frozen=True)
+class _ApplicationServices:
+    data_sources: _DataSources
+
+
+def _application_services(service: DataSourceApiKeyAuthService) -> _ApplicationServices:
+    return _ApplicationServices(data_sources=_DataSources(api_key_auth=service))
 
 
 def test_list_data_source_auth_passes_context_and_serializes_application_result() -> None:

api/tests/unit_tests/controllers/console/auth/test_data_source_oauth.py

@@ -13,12 +13,12 @@ from controllers.console.auth.data_source_oauth import (
 )
 from extensions.ext_application_services import ApplicationServices
 from machinery.context import RequestContext
-from services.data_source_oauth_service import (
+from services.data_source.oauth_service import (
     DataSourceOAuthError,
     InvalidDataSourceOAuthCodeError,
     InvalidDataSourceOAuthProviderError,
 )
-from services.entities.data_source_oauth_entities import DataSourceOAuthCallback
+from services.entities.data_source.oauth import DataSourceOAuthCallback
 from tests.unit_tests.config_override import config_overrides_context
 
 
@@ -33,7 +33,7 @@ def _request_context() -> RequestContext:
 
 def _services(service: MagicMock) -> MagicMock:
     services = create_autospec(ApplicationServices, instance=True, spec_set=True)
-    services.resolve_data_source_oauth.return_value = service
+    services.data_sources.resolve_oauth.return_value = service
     return services
 
 
@@ -55,7 +55,7 @@ def test_get_delegates_to_application_service_and_serializes_response() -> None:
 def test_get_maps_unknown_provider_to_bad_request() -> None:
     service = MagicMock()
     services = _services(service)
-    services.resolve_data_source_oauth.side_effect = InvalidDataSourceOAuthProviderError
+    services.data_sources.resolve_oauth.side_effect = InvalidDataSourceOAuthProviderError
 
     with patch(
         "controllers.console.auth.data_source_oauth.application_services",

api/tests/unit_tests/controllers/console/datasets/test_datasets.py

@@ -3,7 +3,6 @@ import json
 from collections.abc import Callable
 from contextlib import ExitStack
 from inspect import unwrap
-from types import SimpleNamespace
 from unittest.mock import ANY, MagicMock, PropertyMock, call, patch
 
 import pytest
@@ -41,6 +40,7 @@ from core.entities.knowledge_entities import IndexingEstimate
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.provider_manager import ProviderManager
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.entities.dataset_reference import DatasetRef
 from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from extensions.storage.storage_type import StorageType
@@ -48,7 +48,6 @@ from models.account import Account, TenantAccountRole
 from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment
 from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus
 from models.model import ApiToken, App, AppMode, IconType, UploadFile
-from services.dataset_ref_service import DatasetRef
 from services.dataset_service import DatasetPermissionService, DatasetService
 from services.enterprise import rbac_service as enterprise_rbac_service
 
@@ -300,7 +299,7 @@ class TestDatasetList(_UsesSQLiteSession):
                 ),
                 patch(
                     "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources",
-                    return_value=SimpleNamespace(resource_ids=[]),
+                    return_value=enterprise_rbac_service.ResourceWhitelistResources(resource_ids=[]),
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
@@ -327,7 +326,7 @@ class TestDatasetList(_UsesSQLiteSession):
                 ),
                 patch(
                     "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources",
-                    return_value=SimpleNamespace(unrestricted=True, resource_ids=[]),
+                    return_value=enterprise_rbac_service.ResourceWhitelistResources(unrestricted=True, resource_ids=[]),
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
@@ -353,7 +352,9 @@ class TestDatasetList(_UsesSQLiteSession):
                 ),
                 patch(
                     "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources",
-                    return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]),
+                    return_value=enterprise_rbac_service.ResourceWhitelistResources(
+                        resource_ids=["dataset-whitelist-only"]
+                    ),
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
@@ -393,7 +394,9 @@ class TestDatasetList(_UsesSQLiteSession):
                 ),
                 patch(
                     "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources",
-                    return_value=SimpleNamespace(resource_ids=["dataset-whitelist-only"]),
+                    return_value=enterprise_rbac_service.ResourceWhitelistResources(
+                        resource_ids=["dataset-whitelist-only"]
+                    ),
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
@@ -418,7 +421,7 @@ class TestDatasetList(_UsesSQLiteSession):
                 ),
                 patch(
                     "controllers.console.datasets.datasets.enterprise_rbac_service.RBACService.DatasetAccess.whitelist_resources",
-                    return_value=SimpleNamespace(resource_ids=[]),
+                    return_value=enterprise_rbac_service.ResourceWhitelistResources(resource_ids=[]),
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):

api/tests/unit_tests/controllers/console/datasets/test_datasets_document.py

@@ -1,6 +1,6 @@
 import datetime
+import json
 from inspect import unwrap
-from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -46,11 +46,11 @@ from controllers.console.datasets.error import (
     InvalidMetadataError,
 )
 from core.entities.knowledge_entities import IndexingEstimate
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef
 from core.rag.index_processor.constant.index_type import IndexStructureType
-from models.dataset import Dataset, DatasetPermissionEnum
+from models.dataset import Dataset, DatasetPermissionEnum, DatasetProcessRule
 from models.dataset import Document as DatasetDocument
-from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
-from services.dataset_ref_service import DatasetRef, DocumentRef
+from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, ProcessRuleMode
 from services.vector_space_admission_service import (
     VECTOR_SPACE_ADMISSION_ERROR_CODE,
     format_vector_space_admission_error,
@@ -196,6 +196,15 @@ def make_document(**overrides):
     return DatasetDocument(**attrs)
 
 
+def make_process_rule(rules: dict[str, object] | None) -> DatasetProcessRule:
+    return DatasetProcessRule(
+        dataset_id="ds-1",
+        mode=ProcessRuleMode.CUSTOM,
+        rules=json.dumps(rules) if rules is not None else None,
+        created_by="u1",
+    )
+
+
 @pytest.fixture
 def tenant_ctx():
     return (MagicMock(is_dataset_editor=True, id="u1"), "tenant-1")
@@ -271,9 +280,8 @@ class TestGetProcessRuleApi:
         document = MagicMock(dataset_id="ds-1")
         session = MagicMock()
         dataset = MagicMock()
-        dataset.get_latest_process_rule.return_value = SimpleNamespace(
-            mode="custom",
-            rules_dict={"segmentation": {"delimiter": "---", "max_tokens": 123}},
+        dataset.get_latest_process_rule.return_value = make_process_rule(
+            {"segmentation": {"delimiter": "---", "max_tokens": 123}}
         )
 
         with (
@@ -305,7 +313,7 @@ class TestGetProcessRuleApi:
         user, _ = patch_tenant
         session = MagicMock()
         dataset = MagicMock()
-        dataset.get_latest_process_rule.return_value = SimpleNamespace(mode="custom", rules_dict=None)
+        dataset.get_latest_process_rule.return_value = make_process_rule(None)
         with (
             app.test_request_context("/?document_id=doc-1"),
             patch(
@@ -1543,7 +1551,7 @@ class TestDocumentPermissionCases:
         document = MagicMock(dataset_id="ds-1")
         session = MagicMock()
         dataset = MagicMock()
-        dataset.get_latest_process_rule.return_value = SimpleNamespace(mode="custom", rules_dict={"a": 1})
+        dataset.get_latest_process_rule.return_value = make_process_rule({"a": 1})
         with (
             app.test_request_context("/?document_id=doc-1"),
             patch(

api/tests/unit_tests/controllers/service_api/dataset/test_document.py

@@ -47,6 +47,7 @@ from controllers.service_api.dataset.document import (
     InvalidMetadataError,
 )
 from controllers.service_api.dataset.error import ArchivedDocumentImmutableError
+from core.rag.entities.dataset_reference import DatasetRef
 from core.rag.index_processor.constant.index_type import IndexStructureType
 from extensions.storage.storage_type import StorageType
 from models.account import Account
@@ -61,7 +62,6 @@ from models.enums import (
     SegmentStatus,
 )
 from models.model import ApiToken, UploadFile
-from services.dataset_ref_service import DatasetRef
 from services.dataset_service import DocumentService
 from services.entities.knowledge_entities.knowledge_entities import ProcessRule, RetrievalModel
 from services.errors.file import FileTooLargeError as FileTooLargeServiceError

api/tests/unit_tests/core/rag/extractor/test_extract_processor.py

@@ -1,12 +1,13 @@
 from datetime import datetime
 from pathlib import Path
-from types import SimpleNamespace
 
+import httpx
 import pytest
 from sqlalchemy.orm import Session
 
 import core.rag.extractor.extract_processor as processor_module
 from core.rag.extractor.entity.datasource_type import DatasourceType
+from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
 from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.models.document import Document
 from extensions.storage.storage_type import StorageType
@@ -81,8 +82,6 @@ def _patch_all_extractors(monkeypatch) -> _ExtractorFactory:
 
 class TestExtractProcessorLoaders:
     def test_load_from_upload_file_return_docs_and_text(self, monkeypatch: pytest.MonkeyPatch):
-        monkeypatch.setattr(processor_module, "ExtractSetting", lambda **kwargs: SimpleNamespace(**kwargs))
-
         monkeypatch.setattr(
             ExtractProcessor,
             "extract",
@@ -120,9 +119,8 @@ class TestExtractProcessorLoaders:
     def test_load_from_url_builds_temp_file_with_correct_suffix(
         self, monkeypatch: pytest.MonkeyPatch, url, headers, expected_suffix
     ):
-        response = SimpleNamespace(headers=headers, content=b"body")
+        response = httpx.Response(200, headers=headers, content=b"body")
         monkeypatch.setattr(processor_module.remote_fetcher, "make_request", lambda *args, **kwargs: response)
-        monkeypatch.setattr(processor_module, "ExtractSetting", lambda **kwargs: SimpleNamespace(**kwargs))
 
         captured = {}
 
@@ -144,7 +142,7 @@ class TestExtractProcessorLoaders:
 
     def test_load_from_url_extracts_long_text_without_upload_file(self, monkeypatch: pytest.MonkeyPatch):
         content = "a" * 100_000
-        response = SimpleNamespace(headers={"Content-Type": "text/plain"}, content=content.encode())
+        response = httpx.Response(200, headers={"Content-Type": "text/plain"}, content=content.encode())
         monkeypatch.setattr(processor_module.remote_fetcher, "make_request", lambda *args, **kwargs: response)
         apply_config_overrides(monkeypatch, ETL_TYPE="SelfHosted")
 
@@ -179,7 +177,7 @@ class TestExtractProcessorFileRouting:
         monkeypatch.setattr(processor_module.storage, "download", fake_download)
         monkeypatch.setattr(processor_module.tempfile, "_get_candidate_names", lambda: iter(["candidate-name"]))
 
-        setting = SimpleNamespace(
+        setting = ExtractSetting(
             datasource_type=DatasourceType.FILE,
             upload_file=_upload_file(key=f"uploaded{extension}"),
         )
@@ -260,7 +258,7 @@ class TestExtractProcessorFileRouting:
         assert kwargs["session"] is unbound_session
 
     def test_extract_requires_upload_file_when_file_path_not_provided(self):
-        setting = SimpleNamespace(datasource_type=DatasourceType.FILE, upload_file=None)
+        setting = ExtractSetting(datasource_type=DatasourceType.FILE, upload_file=None)
 
         with pytest.raises(AssertionError, match="upload_file is required"):
             ExtractProcessor.extract(setting)
@@ -270,20 +268,22 @@ class TestExtractProcessorDatasourceRouting:
     def test_extract_routes_notion_datasource(self, monkeypatch: pytest.MonkeyPatch):
         factory = _patch_all_extractors(monkeypatch)
 
-        notion_info = SimpleNamespace(
+        notion_info = NotionInfo(
             notion_workspace_id="ws",
             notion_obj_id="obj",
             notion_page_type="page",
-            document="doc",
+            document=None,
             tenant_id="tenant",
             credential_id="cred",
+            notion_access_token="secret",
         )
-        setting = SimpleNamespace(datasource_type=DatasourceType.NOTION, notion_info=notion_info)
+        setting = ExtractSetting(datasource_type=DatasourceType.NOTION, notion_info=notion_info)
 
         docs = ExtractProcessor.extract(setting)
 
         assert docs[0].page_content == "extracted-by-NotionExtractor"
         assert factory.calls[-1][0] == "NotionExtractor"
+        assert factory.calls[-1][2]["notion_access_token"] == "secret"
 
     @pytest.mark.parametrize(
         ("provider", "expected"),
@@ -298,7 +298,7 @@ class TestExtractProcessorDatasourceRouting:
     ):
         factory = _patch_all_extractors(monkeypatch)
 
-        website_info = SimpleNamespace(
+        website_info = WebsiteInfo(
             provider=provider,
             url="https://example.com",
             job_id="job",
@@ -306,14 +306,14 @@ class TestExtractProcessorDatasourceRouting:
             mode="crawl",
             only_main_content=True,
         )
-        setting = SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=website_info)
+        setting = ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=website_info)
 
         docs = ExtractProcessor.extract(setting)
         assert docs[0].page_content == f"extracted-by-{expected}"
         assert factory.calls[-1][0] == expected
 
     def test_extract_unsupported_website_provider(self):
-        bad_provider = SimpleNamespace(
+        bad_provider = WebsiteInfo(
             provider="unknown",
             url="https://example.com",
             job_id="job",
@@ -321,19 +321,19 @@ class TestExtractProcessorDatasourceRouting:
             mode="crawl",
             only_main_content=True,
         )
-        setting = SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=bad_provider)
+        setting = ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=bad_provider)
 
         with pytest.raises(ValueError, match="Unsupported website provider"):
             ExtractProcessor.extract(setting)
 
     def test_extract_unsupported_datasource_type(self):
         with pytest.raises(ValueError, match="Unsupported datasource type"):
-            ExtractProcessor.extract(SimpleNamespace(datasource_type="unknown"))
+            ExtractProcessor.extract(ExtractSetting(datasource_type="unknown"))
 
     def test_extract_requires_notion_info(self):
         with pytest.raises(AssertionError, match="notion_info is required"):
-            ExtractProcessor.extract(SimpleNamespace(datasource_type=DatasourceType.NOTION, notion_info=None))
+            ExtractProcessor.extract(ExtractSetting(datasource_type=DatasourceType.NOTION, notion_info=None))
 
     def test_extract_requires_website_info(self):
         with pytest.raises(AssertionError, match="website_info is required"):
-            ExtractProcessor.extract(SimpleNamespace(datasource_type=DatasourceType.WEBSITE, website_info=None))
+            ExtractProcessor.extract(ExtractSetting(datasource_type=DatasourceType.WEBSITE, website_info=None))

api/tests/unit_tests/extensions/test_ext_application_services.py

@@ -1,7 +1,6 @@
 """Tests for application-service dependency wiring."""
 
 import json
-from types import SimpleNamespace
 from typing import cast
 from unittest.mock import MagicMock, call, patch
 from uuid import uuid4
@@ -49,11 +48,14 @@ from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthS
 from services.billing_portal_service import BillingPortalService
 from services.billing_service import BillingService
 from services.compliance_download_service import ComplianceDownloadService
+from services.data_source.binding_application_service import DataSourceBindingApplicationService
+from services.data_source.notion_import_application_service import NotionImportApplicationService
 from services.enterprise.enterprise_service import WebAppSettings
 from services.entities.mail_entities import InnerMailMessage
 from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError
 from services.file_service import FileService
 from services.init_validation_service import InvalidInitializationPasswordError
+from services.knowledge.application import DocumentSyncApplicationService, IndexingEstimateApplicationService
 from services.partner_tenant_binding_service import PartnerTenantBindingService
 from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache
 from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveService
@@ -91,7 +93,7 @@ def test_build_application_services_configures_init_validation(
         database_client=sqlite_session_factory,
         deployment_edition=deployment_edition,
         initialization_password=initialization_password,
-        redis=MagicMock(spec=RedisClientWrapper),
+        redis=cast(RedisClientWrapper, object()),
     )
 
     assert services.init_validation.is_validated(session_validated=session_validated) is expected
@@ -488,7 +490,36 @@ def test_build_application_services_wires_data_source_api_key_auth(
         redis=MagicMock(spec=RedisClientWrapper),
     )
 
-    assert isinstance(services.data_source_api_key_auth, DataSourceApiKeyAuthService)
+    assert isinstance(services.data_sources.api_key_auth, DataSourceApiKeyAuthService)
+
+
+def test_build_application_services_groups_dataset_services_and_reuses_repositories(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    services = ext_application_services.build_application_services(
+        database_client=sqlite_session_factory,
+        deployment_edition=DeploymentEdition.COMMUNITY,
+        initialization_password="",
+        redis=cast(RedisClientWrapper, object()),
+    )
+
+    assert isinstance(services.data_sources.bindings, DataSourceBindingApplicationService)
+    assert isinstance(services.data_sources.notion_imports, NotionImportApplicationService)
+    assert isinstance(services.knowledge.document_sync, DocumentSyncApplicationService)
+    assert isinstance(services.knowledge.indexing_estimates, IndexingEstimateApplicationService)
+    assert services.data_sources.bindings._bindings is services.data_sources.oauth["notion"]._bindings
+    assert services.data_sources.notion_imports._dataset_access is services.knowledge.document_sync._dataset_access
+    assert services.data_sources.notion_imports._dataset_access is services.knowledge.indexing_estimates._dataset_access
+    assert services.data_sources.notion_imports._documents is services.knowledge.document_sync._documents
+    assert services.data_sources.notion_imports._documents is services.knowledge.indexing_estimates._documents
+    assert (
+        services.data_sources.notion_imports._source._credentials
+        is services.knowledge.indexing_estimates._gateway._credentials
+    )
+    assert (
+        services.workspace_member_queries._members
+        is services.data_sources.notion_imports._dataset_access._workspace_roles
+    )
 
 
 @pytest.mark.parametrize(
@@ -576,7 +607,7 @@ def test_build_application_services_adapts_enterprise_webapp_access_mode(
         patch("extensions.ext_application_services.SystemFeatureService.is_webapp_auth_enabled", return_value=True),
         patch(
             "extensions.ext_application_services.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
-            return_value=SimpleNamespace(access_mode="private_all"),
+            return_value=WebAppSettings(accessMode="private_all"),
         ) as get_access_mode,
     ):
         services = ext_application_services.build_application_services(
@@ -636,7 +667,7 @@ def test_build_application_services_maps_invalid_access_mode_to_unavailable(
         patch("extensions.ext_application_services.SystemFeatureService.is_webapp_auth_enabled", return_value=True),
         patch(
             "extensions.ext_application_services.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
-            return_value=SimpleNamespace(access_mode="invalid"),
+            return_value=WebAppSettings(accessMode="invalid"),
         ),
     ):
         services = ext_application_services.build_application_services(
@@ -685,7 +716,7 @@ def test_build_application_services_wires_webapp_permission(
         ) as enabled,
         patch(
             "extensions.ext_application_services.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
-            return_value=SimpleNamespace(access_mode="private"),
+            return_value=WebAppSettings(accessMode="private"),
         ) as get_access_mode,
         patch(
             "extensions.ext_application_services.EnterpriseService.WebAppAuth.is_user_allowed_to_access_webapp",

api/tests/unit_tests/repositories/data_source/__init__.py

@@ -0,0 +1 @@
+"""Tests for data-source persistence adapters."""

api/tests/unit_tests/repositories/data_source/test_api_key_auth_repository.py

@@ -5,8 +5,8 @@ from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
 from models.source import DataSourceApiKeyAuthBinding
-from repositories.data_source_api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from repositories.data_source.api_key_auth_repository import SQLAlchemyDataSourceApiKeyAuthBindingRepository
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 def _binding(

api/tests/unit_tests/repositories/data_source/test_credential_repository.py

@@ -0,0 +1,223 @@
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from models.credential_permission import CredentialPermission, CredentialType
+from models.enums import PermissionEnum
+from models.oauth import DatasourceProvider
+from repositories.data_source.credential_repository import SQLAlchemyDatasourceCredentialRepository
+
+
+def _provider(
+    *,
+    workspace_id: str = "workspace-1",
+    owner_id: str = "owner-1",
+    provider: str = "notion_datasource",
+    plugin_id: str = "langgenius/notion_datasource",
+    visibility: PermissionEnum = PermissionEnum.ONLY_ME,
+) -> DatasourceProvider:
+    return DatasourceProvider(
+        tenant_id=workspace_id,
+        name="Notion",
+        provider=provider,
+        plugin_id=plugin_id,
+        auth_type="oauth2",
+        encrypted_credentials={"integration_secret": "encrypted"},
+        user_id=owner_id,
+        visibility=visibility,
+        expires_at=-1,
+    )
+
+
+def test_get_visible_scopes_workspace_provider_plugin_and_actor(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        credential = _provider(visibility=PermissionEnum.PARTIAL_TEAM)
+        session.add(credential)
+        session.flush()
+        session.add(
+            CredentialPermission(
+                credential_id=credential.id,
+                credential_type=CredentialType.DATASOURCE_PROVIDER,
+                account_id="member-1",
+                tenant_id="workspace-1",
+            )
+        )
+        credential_id = credential.id
+
+    record = repository.get_visible(
+        workspace_id="workspace-1",
+        actor_id="member-1",
+        credential_id=credential_id,
+        provider="notion_datasource",
+        plugin_id="langgenius/notion_datasource",
+    )
+
+    assert record is not None
+    assert record.encrypted_credentials == {"integration_secret": "encrypted"}
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-2",
+            actor_id="member-1",
+            credential_id=credential_id,
+            provider="notion_datasource",
+            plugin_id="langgenius/notion_datasource",
+        )
+        is None
+    )
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-1",
+            actor_id="member-1",
+            credential_id=credential_id,
+            provider="other",
+            plugin_id="langgenius/notion_datasource",
+        )
+        is None
+    )
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-1",
+            actor_id="member-1",
+            credential_id=credential_id,
+            provider="notion_datasource",
+            plugin_id="other/plugin",
+        )
+        is None
+    )
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-1",
+            actor_id="outsider-1",
+            credential_id=credential_id,
+            provider="notion_datasource",
+            plugin_id="langgenius/notion_datasource",
+        )
+        is None
+    )
+
+
+def test_get_visible_ignores_revoked_partial_permission(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        credential = _provider(visibility=PermissionEnum.PARTIAL_TEAM)
+        session.add(credential)
+        session.flush()
+        session.add(
+            CredentialPermission(
+                credential_id=credential.id,
+                credential_type=CredentialType.DATASOURCE_PROVIDER,
+                account_id="member-1",
+                tenant_id="workspace-1",
+                has_permission=False,
+            )
+        )
+        credential_id = credential.id
+
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-1",
+            actor_id="member-1",
+            credential_id=credential_id,
+            provider="notion_datasource",
+            plugin_id="langgenius/notion_datasource",
+        )
+        is None
+    )
+
+
+def test_get_visible_does_not_apply_stale_share_to_private_credential(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        credential = _provider(visibility=PermissionEnum.ONLY_ME)
+        session.add(credential)
+        session.flush()
+        session.add(
+            CredentialPermission(
+                credential_id=credential.id,
+                credential_type=CredentialType.DATASOURCE_PROVIDER,
+                account_id="member-1",
+                tenant_id="workspace-1",
+            )
+        )
+        credential_id = credential.id
+
+    assert (
+        repository.get_visible(
+            workspace_id="workspace-1",
+            actor_id="member-1",
+            credential_id=credential_id,
+            provider="notion_datasource",
+            plugin_id="langgenius/notion_datasource",
+        )
+        is None
+    )
+
+
+def test_update_if_unchanged_detects_concurrent_change(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        credential = _provider()
+        session.add(credential)
+        session.flush()
+        credential_id = credential.id
+
+    record = repository.get_visible(
+        workspace_id="workspace-1",
+        actor_id="owner-1",
+        credential_id=credential_id,
+        provider="notion_datasource",
+        plugin_id="langgenius/notion_datasource",
+    )
+    assert record is not None
+
+    with sqlite_session_factory.begin() as session:
+        stored = session.scalar(select(DatasourceProvider).where(DatasourceProvider.id == credential_id))
+        assert stored is not None
+        stored.encrypted_credentials = {"integration_secret": "concurrent"}
+        stored.expires_at = 200
+
+    assert not repository.update_if_unchanged(
+        record=record,
+        encrypted_credentials={"integration_secret": "new"},
+        expires_at=100,
+    )
+
+
+def test_update_if_unchanged_persists_refreshed_snapshot(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDatasourceCredentialRepository(session_factory=sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        credential = _provider()
+        session.add(credential)
+        session.flush()
+        credential_id = credential.id
+
+    record = repository.get_visible(
+        workspace_id="workspace-1",
+        actor_id="owner-1",
+        credential_id=credential_id,
+        provider="notion_datasource",
+        plugin_id="langgenius/notion_datasource",
+    )
+    assert record is not None
+    assert repository.update_if_unchanged(
+        record=record,
+        encrypted_credentials={"integration_secret": "new"},
+        expires_at=100,
+    )
+
+    with sqlite_session_factory() as session:
+        stored = session.get(DatasourceProvider, credential_id)
+        assert stored is not None
+        assert stored.encrypted_credentials == {"integration_secret": "new"}
+        assert stored.expires_at == 100
+        assert stored.updated_at >= record.updated_at

api/tests/unit_tests/repositories/data_source/test_oauth_binding_repository.py

@@ -0,0 +1,247 @@
+from sqlalchemy import select, update
+from sqlalchemy.orm import Session, sessionmaker
+
+from models.source import DataSourceOauthBinding
+from repositories.data_source.oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
+from services.entities.data_source.oauth import DataSourceOAuthAuthorization
+
+
+def test_upsert_authorization_inserts_and_reenables_existing_binding(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
+    authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []})
+
+    repository.upsert_authorization(
+        workspace_id="workspace-1",
+        provider="notion",
+        authorization=authorization,
+    )
+    with sqlite_session_factory.begin() as session:
+        binding = session.scalar(select(DataSourceOauthBinding))
+        assert binding is not None
+        binding.disabled = True
+
+    updated_authorization = DataSourceOAuthAuthorization(
+        access_token="token",
+        source_info={"pages": [{"page_id": "page-1"}]},
+    )
+    repository.upsert_authorization(
+        workspace_id="workspace-1",
+        provider="notion",
+        authorization=updated_authorization,
+    )
+
+    with sqlite_session_factory() as session:
+        bindings = session.scalars(select(DataSourceOauthBinding)).all()
+        assert len(bindings) == 1
+        assert bindings[0].disabled is False
+        assert bindings[0].source_info == updated_authorization.source_info
+
+
+def test_get_enabled_scopes_by_workspace_provider_and_state(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        enabled = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="enabled-token",
+            source_info={"pages": []},
+            disabled=False,
+        )
+        disabled = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="disabled-token",
+            source_info={"pages": []},
+            disabled=True,
+        )
+        session.add_all([enabled, disabled])
+
+    record = repository.get_enabled(
+        workspace_id="workspace-1",
+        provider="notion",
+        binding_id=enabled.id,
+    )
+
+    assert record is not None
+    assert record.id == enabled.id
+    assert record.access_token == "enabled-token"
+    assert (
+        repository.get_enabled(
+            workspace_id="workspace-2",
+            provider="notion",
+            binding_id=enabled.id,
+        )
+        is None
+    )
+    assert (
+        repository.get_enabled(
+            workspace_id="workspace-1",
+            provider="other",
+            binding_id=enabled.id,
+        )
+        is None
+    )
+    assert (
+        repository.get_enabled(
+            workspace_id="workspace-1",
+            provider="notion",
+            binding_id=disabled.id,
+        )
+        is None
+    )
+
+
+def test_update_source_info_only_updates_matching_enabled_binding(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        binding = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="token",
+            source_info={"pages": []},
+            disabled=False,
+        )
+        session.add(binding)
+
+    assert repository.update_source_info(
+        workspace_id="workspace-1",
+        provider="notion",
+        binding_id=binding.id,
+        source_info={"pages": [{"page_id": "page-1"}]},
+    )
+    assert not repository.update_source_info(
+        workspace_id="workspace-2",
+        provider="notion",
+        binding_id=binding.id,
+        source_info={"pages": []},
+    )
+    assert not repository.update_source_info(
+        workspace_id="workspace-1",
+        provider="other",
+        binding_id=binding.id,
+        source_info={"pages": []},
+    )
+
+    with sqlite_session_factory() as session:
+        updated = session.get(DataSourceOauthBinding, binding.id)
+        assert updated is not None
+        assert updated.source_info == {"pages": [{"page_id": "page-1"}]}
+
+    with sqlite_session_factory.begin() as session:
+        stored = session.get(DataSourceOauthBinding, binding.id)
+        assert stored is not None
+        stored.disabled = True
+
+    assert not repository.update_source_info(
+        workspace_id="workspace-1",
+        provider="notion",
+        binding_id=binding.id,
+        source_info={"pages": []},
+    )
+
+
+def test_binding_management_lists_enabled_and_changes_state_atomically(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        enabled = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="enabled-token",
+            source_info={"workspace_id": "notion-workspace"},
+            disabled=False,
+        )
+        disabled = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="disabled-token",
+            source_info={},
+            disabled=True,
+        )
+        session.add_all([enabled, disabled])
+
+    summaries = repository.list_enabled_bindings(workspace_id="workspace-1")
+
+    assert [summary.id for summary in summaries] == [enabled.id]
+    assert (
+        repository.change_disabled_state(workspace_id="workspace-1", binding_id=enabled.id, disabled=True) == "updated"
+    )
+    assert (
+        repository.change_disabled_state(workspace_id="workspace-1", binding_id=enabled.id, disabled=True)
+        == "already_disabled"
+    )
+    assert (
+        repository.change_disabled_state(workspace_id="workspace-1", binding_id=disabled.id, disabled=False)
+        == "updated"
+    )
+    assert (
+        repository.change_disabled_state(workspace_id="workspace-2", binding_id=disabled.id, disabled=False)
+        == "not_found"
+    )
+
+
+def test_null_disabled_state_is_treated_as_enabled(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
+    with sqlite_session_factory.begin() as session:
+        binding = DataSourceOauthBinding(
+            tenant_id="workspace-1",
+            provider="notion",
+            access_token="token",
+            source_info={"workspace_id": "notion-workspace"},
+            disabled=False,
+        )
+        session.add(binding)
+        session.flush()
+        binding_id = binding.id
+        session.execute(
+            update(DataSourceOauthBinding).where(DataSourceOauthBinding.id == binding_id).values(disabled=None)
+        )
+
+    assert (
+        repository.get_enabled(
+            workspace_id="workspace-1",
+            provider="notion",
+            binding_id=binding_id,
+        )
+        is not None
+    )
+    assert [summary.id for summary in repository.list_enabled_bindings(workspace_id="workspace-1")] == [binding_id]
+    assert (
+        repository.change_disabled_state(
+            workspace_id="workspace-1",
+            binding_id=binding_id,
+            disabled=False,
+        )
+        == "already_enabled"
+    )
+    assert repository.update_source_info(
+        workspace_id="workspace-1",
+        provider="notion",
+        binding_id=binding_id,
+        source_info={"workspace_id": "updated-workspace"},
+    )
+    assert (
+        repository.change_disabled_state(
+            workspace_id="workspace-1",
+            binding_id=binding_id,
+            disabled=True,
+        )
+        == "updated"
+    )
+    assert (
+        repository.get_enabled(
+            workspace_id="workspace-1",
+            provider="notion",
+            binding_id=binding_id,
+        )
+        is None
+    )

api/tests/unit_tests/repositories/knowledge/__init__.py

@@ -0,0 +1 @@
+"""Tests for knowledge persistence adapters."""

api/tests/unit_tests/repositories/knowledge/test_dataset_repository.py

@@ -0,0 +1,155 @@
+import pytest
+from sqlalchemy import select
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.rag.entities.dataset_reference import DatasetRef
+from models.dataset import Dataset, DatasetPermission
+from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
+from services.knowledge.application import DatasetNotFoundError
+
+
+def _dataset(dataset_id: str, workspace_id: str, *, permission: str = "only_me") -> Dataset:
+    return Dataset(
+        id=dataset_id,
+        tenant_id=workspace_id,
+        name=f"Dataset {dataset_id}",
+        description="",
+        provider="vendor",
+        permission=permission,
+        created_by="account-1",
+        maintainer="account-1",
+        data_source_type="notion_import",
+        indexing_technique="economy",
+    )
+
+
+def test_get_by_ref_and_access_snapshot_are_tenant_scoped(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1", permission="partial_members"))
+        session.add(
+            DatasetPermission(
+                tenant_id="workspace-1",
+                dataset_id="dataset-1",
+                account_id="account-2",
+            )
+        )
+
+    repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory)
+
+    record = repository.get_by_ref(DatasetRef("workspace-1", "dataset-1"))
+    snapshot = repository.get_access_snapshot(
+        workspace_id="workspace-1",
+        dataset_id="dataset-1",
+        actor_id="account-2",
+    )
+
+    assert record is not None
+    assert record.workspace_id == "workspace-1"
+    assert snapshot is not None
+    assert snapshot.dataset == record
+    assert snapshot.actor_has_partial_access is True
+    assert repository.get_by_ref(DatasetRef("workspace-2", "dataset-1")) is None
+    assert (
+        repository.get_access_snapshot(
+            workspace_id="workspace-2",
+            dataset_id="dataset-1",
+            actor_id="account-2",
+        )
+        is None
+    )
+
+
+def test_partial_member_queries_ignore_revoked_permissions(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1", permission="partial_members"))
+        session.add_all(
+            [
+                DatasetPermission(
+                    tenant_id="workspace-1",
+                    dataset_id="dataset-1",
+                    account_id="granted-account",
+                    has_permission=True,
+                ),
+                DatasetPermission(
+                    tenant_id="workspace-1",
+                    dataset_id="dataset-1",
+                    account_id="revoked-account",
+                    has_permission=False,
+                ),
+            ]
+        )
+
+    repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory)
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+
+    assert set(repository.list_partial_member_ids(dataset_ref)) == {"granted-account"}
+    revoked_snapshot = repository.get_access_snapshot(
+        workspace_id="workspace-1",
+        dataset_id="dataset-1",
+        actor_id="revoked-account",
+    )
+    assert revoked_snapshot is not None
+    assert revoked_snapshot.actor_has_partial_access is False
+
+
+def test_replace_partial_member_ids_is_atomic_deduplicated_and_tenant_scoped(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1", permission="partial_members"))
+        session.add_all(
+            [
+                DatasetPermission(
+                    tenant_id="workspace-1",
+                    dataset_id="dataset-1",
+                    account_id="old-account",
+                ),
+                DatasetPermission(
+                    tenant_id="workspace-2",
+                    dataset_id="dataset-1",
+                    account_id="other-workspace-account",
+                ),
+            ]
+        )
+
+    repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory)
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+    repository.replace_partial_member_ids(dataset_ref, ["account-2", "account-2", "account-3"])
+
+    assert set(repository.list_partial_member_ids(dataset_ref)) == {"account-2", "account-3"}
+    with sqlite_session_factory() as session:
+        other_workspace_ids = session.scalars(
+            select(DatasetPermission.account_id).where(DatasetPermission.tenant_id == "workspace-2")
+        ).all()
+        assert other_workspace_ids == ["other-workspace-account"]
+
+    repository.replace_partial_member_ids(dataset_ref, [])
+
+    assert repository.list_partial_member_ids(dataset_ref) == ()
+    with sqlite_session_factory() as session:
+        assert session.scalar(select(DatasetPermission).where(DatasetPermission.tenant_id == "workspace-2")) is not None
+
+
+def test_replace_partial_member_ids_rejects_missing_owner_without_changing_permissions(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1", permission="partial_members"))
+        session.add(
+            DatasetPermission(
+                tenant_id="workspace-1",
+                dataset_id="dataset-1",
+                account_id="account-1",
+            )
+        )
+
+    repository = SQLAlchemyDatasetRepository(session_factory=sqlite_session_factory)
+
+    with pytest.raises(DatasetNotFoundError):
+        repository.replace_partial_member_ids(DatasetRef("workspace-2", "dataset-1"), ["account-2"])
+
+    assert repository.list_partial_member_ids(DatasetRef("workspace-1", "dataset-1")) == ("account-1",)

api/tests/unit_tests/repositories/knowledge/test_document_repository.py

@@ -0,0 +1,185 @@
+import json
+from datetime import datetime
+
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.rag.entities.dataset_reference import DatasetRef
+from core.rag.index_processor.constant.index_type import IndexStructureType
+from models.dataset import Dataset, Document
+from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
+from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository
+
+
+def _dataset(dataset_id: str, workspace_id: str) -> Dataset:
+    return Dataset(
+        id=dataset_id,
+        tenant_id=workspace_id,
+        name=f"Dataset {dataset_id}",
+        description="",
+        provider="vendor",
+        created_by="account-1",
+        maintainer="account-1",
+        data_source_type="notion_import",
+        indexing_technique="economy",
+    )
+
+
+def _document(
+    document_id: str,
+    *,
+    workspace_id: str = "workspace-1",
+    dataset_id: str = "dataset-1",
+    source_type: str = DataSourceType.NOTION_IMPORT,
+    enabled: bool = True,
+    archived: bool = False,
+    status: str = IndexingStatus.COMPLETED,
+    batch: str = "batch-1",
+) -> Document:
+    source_info = (
+        {"notion_page_id": f"page-{document_id}"}
+        if source_type == DataSourceType.NOTION_IMPORT
+        else {"upload_file_id": f"file-{document_id}"}
+    )
+    return Document(
+        id=document_id,
+        tenant_id=workspace_id,
+        dataset_id=dataset_id,
+        position=1,
+        data_source_type=source_type,
+        data_source_info=json.dumps(source_info),
+        batch=batch,
+        name=f"Document {document_id}",
+        created_from=DocumentCreatedFrom.API,
+        created_by="account-1",
+        created_at=datetime(2026, 1, 1),
+        updated_at=datetime(2026, 1, 1),
+        indexing_status=status,
+        doc_form=IndexStructureType.PARAGRAPH_INDEX,
+        word_count=1,
+        enabled=enabled,
+        archived=archived,
+        is_paused=False,
+    )
+
+
+def test_document_queries_enforce_the_complete_owner_chain(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add_all(
+            [
+                _dataset("dataset-1", "workspace-1"),
+                _dataset("dataset-2", "workspace-2"),
+                _document("document-1"),
+                _document("document-2", workspace_id="workspace-2", dataset_id="dataset-2"),
+            ]
+        )
+
+    repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory)
+    document_ref = DatasetRef("workspace-1", "dataset-1").document("document-1")
+
+    record = repository.get_by_ref(document_ref)
+
+    assert record is not None
+    assert record.workspace_id == "workspace-1"
+    assert record.data_source_info == {"notion_page_id": "page-document-1"}
+    assert repository.get_by_ref(DatasetRef("workspace-2", "dataset-1").document("document-1")) is None
+    assert repository.get_by_id_for_workspace(workspace_id="workspace-2", document_id="document-1") is None
+
+
+def test_active_notion_queries_exclude_other_sources_and_inactive_documents(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1"))
+        session.add_all(
+            [
+                _document("notion"),
+                _document("upload", source_type=DataSourceType.UPLOAD_FILE),
+                _document("disabled", enabled=False),
+                _document("archived", archived=True),
+            ]
+        )
+
+    repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory)
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+
+    assert repository.list_active_notion_refs(dataset_ref) == (dataset_ref.document("notion"),)
+    assert repository.list_bound_notion_page_ids(dataset_ref) == frozenset({"page-notion"})
+
+
+def test_document_records_normalize_invalid_mapping_data_and_ignore_invalid_notion_page_ids(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    invalid_json = _document("invalid-json")
+    invalid_json.data_source_info = "{invalid"
+    missing_info = _document("missing-info")
+    missing_info.data_source_info = None
+    invalid_page_id = _document("invalid-page-id")
+    invalid_page_id.data_source_info = json.dumps({"notion_page_id": 42})
+    invalid_page_id.doc_metadata = {"rank": 1}
+    with sqlite_session_factory.begin() as session:
+        session.add(_dataset("dataset-1", "workspace-1"))
+        session.add_all([invalid_json, missing_info, invalid_page_id])
+
+    repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory)
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+
+    invalid_record = repository.get_by_ref(dataset_ref.document("invalid-json"))
+    missing_record = repository.get_by_ref(dataset_ref.document("missing-info"))
+    assert invalid_record is not None
+    assert invalid_record.data_source_info == {}
+    assert missing_record is not None
+    assert missing_record.data_source_info is None
+    metadata_record = repository.get_by_ref(dataset_ref.document("invalid-page-id"))
+    assert metadata_record is not None
+    assert metadata_record.doc_metadata == {"rank": 1}
+    assert repository.list_bound_notion_page_ids(dataset_ref) == frozenset()
+
+
+def test_batch_and_available_queries_keep_explicit_state_semantics(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add_all(
+            [
+                _dataset("dataset-1", "workspace-1"),
+                _dataset("dataset-2", "workspace-1"),
+                _dataset("dataset-3", "workspace-2"),
+            ]
+        )
+        session.add_all(
+            [
+                _document("available"),
+                _document("waiting", status=IndexingStatus.WAITING),
+                _document("disabled", enabled=False),
+                _document("archived", archived=True),
+                _document("other-batch", batch="batch-2"),
+                _document("other-dataset", dataset_id="dataset-2"),
+                _document("other-workspace", workspace_id="workspace-2", dataset_id="dataset-3"),
+            ]
+        )
+
+    repository = SQLAlchemyDocumentRepository(session_factory=sqlite_session_factory)
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+
+    assert repository.list_by_refs(dataset_ref, []) == ()
+    assert repository.list_available_by_refs(dataset_ref, []) == ()
+    assert {record.id for record in repository.list_by_batch(dataset_ref, "batch-1")} == {
+        "available",
+        "waiting",
+        "disabled",
+        "archived",
+    }
+    assert [record.id for record in repository.list_by_refs(dataset_ref, ["waiting", "available", "missing"])] == [
+        "waiting",
+        "available",
+    ]
+    assert [
+        record.id
+        for record in repository.list_available_by_refs(
+            dataset_ref,
+            ["waiting", "disabled", "archived", "available"],
+        )
+    ] == ["available"]
+    assert {record.id for record in repository.list_working_by_dataset(dataset_ref)} == {"available", "other-batch"}

api/tests/unit_tests/repositories/test_data_source_oauth_binding_repository.py

@@ -1,120 +0,0 @@
-from sqlalchemy import select
-from sqlalchemy.orm import Session, sessionmaker
-
-from models.source import DataSourceOauthBinding
-from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSourceOAuthBindingRepository
-from services.entities.data_source_oauth_entities import DataSourceOAuthAuthorization
-
-
-def test_upsert_authorization_inserts_and_reenables_existing_binding(
-    sqlite_session_factory: sessionmaker[Session],
-) -> None:
-    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
-    authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []})
-
-    repository.upsert_authorization(
-        workspace_id="workspace-1",
-        provider="notion",
-        authorization=authorization,
-    )
-    with sqlite_session_factory.begin() as session:
-        binding = session.scalar(select(DataSourceOauthBinding))
-        assert binding is not None
-        binding.disabled = True
-
-    updated_authorization = DataSourceOAuthAuthorization(
-        access_token="token",
-        source_info={"pages": [{"page_id": "page-1"}]},
-    )
-    repository.upsert_authorization(
-        workspace_id="workspace-1",
-        provider="notion",
-        authorization=updated_authorization,
-    )
-
-    with sqlite_session_factory() as session:
-        bindings = session.scalars(select(DataSourceOauthBinding)).all()
-        assert len(bindings) == 1
-        assert bindings[0].disabled is False
-        assert bindings[0].source_info == updated_authorization.source_info
-
-
-def test_get_enabled_scopes_by_workspace_provider_and_state(
-    sqlite_session_factory: sessionmaker[Session],
-) -> None:
-    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
-    with sqlite_session_factory.begin() as session:
-        enabled = DataSourceOauthBinding(
-            tenant_id="workspace-1",
-            provider="notion",
-            access_token="enabled-token",
-            source_info={"pages": []},
-            disabled=False,
-        )
-        disabled = DataSourceOauthBinding(
-            tenant_id="workspace-1",
-            provider="notion",
-            access_token="disabled-token",
-            source_info={"pages": []},
-            disabled=True,
-        )
-        session.add_all([enabled, disabled])
-
-    record = repository.get_enabled(
-        workspace_id="workspace-1",
-        provider="notion",
-        binding_id=enabled.id,
-    )
-
-    assert record is not None
-    assert record.id == enabled.id
-    assert record.access_token == "enabled-token"
-    assert (
-        repository.get_enabled(
-            workspace_id="workspace-2",
-            provider="notion",
-            binding_id=enabled.id,
-        )
-        is None
-    )
-    assert (
-        repository.get_enabled(
-            workspace_id="workspace-1",
-            provider="notion",
-            binding_id=disabled.id,
-        )
-        is None
-    )
-
-
-def test_update_source_info_only_updates_matching_enabled_binding(
-    sqlite_session_factory: sessionmaker[Session],
-) -> None:
-    repository = SQLAlchemyDataSourceOAuthBindingRepository(sqlite_session_factory)
-    with sqlite_session_factory.begin() as session:
-        binding = DataSourceOauthBinding(
-            tenant_id="workspace-1",
-            provider="notion",
-            access_token="token",
-            source_info={"pages": []},
-            disabled=False,
-        )
-        session.add(binding)
-
-    assert repository.update_source_info(
-        workspace_id="workspace-1",
-        provider="notion",
-        binding_id=binding.id,
-        source_info={"pages": [{"page_id": "page-1"}]},
-    )
-    assert not repository.update_source_info(
-        workspace_id="workspace-2",
-        provider="notion",
-        binding_id=binding.id,
-        source_info={"pages": []},
-    )
-
-    with sqlite_session_factory() as session:
-        updated = session.get(DataSourceOauthBinding, binding.id)
-        assert updated is not None
-        assert updated.source_info == {"pages": [{"page_id": "page-1"}]}

api/tests/unit_tests/repositories/test_workspace_member_query_repository.py

@@ -124,3 +124,34 @@ def test_list_for_workspace_returns_empty_tuple_without_membership(
     result = WorkspaceMemberQueryRepository(sqlite_session_factory).list_for_workspace("workspace-1")
 
     assert result == ()
+
+
+def test_get_legacy_role_is_scoped_to_workspace_and_account(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    created_at = datetime(2026, 1, 1)
+    account = make_account("member", status=AccountStatus.ACTIVE, created_at=created_at)
+    workspace = make_tenant("workspace-1")
+    other_workspace = make_tenant("workspace-2")
+    with sqlite_session_factory.begin() as session:
+        session.add_all(
+            [
+                account,
+                workspace,
+                other_workspace,
+                TenantAccountJoin(
+                    tenant_id=workspace.id,
+                    account_id=account.id,
+                    role=TenantAccountRole.DATASET_OPERATOR,
+                ),
+            ]
+        )
+
+    repository = WorkspaceMemberQueryRepository(sqlite_session_factory)
+
+    assert (
+        repository.get_legacy_role(workspace_id=workspace.id, account_id=account.id)
+        == TenantAccountRole.DATASET_OPERATOR.value
+    )
+    assert repository.get_legacy_role(workspace_id=other_workspace.id, account_id=account.id) is None
+    assert repository.get_legacy_role(workspace_id=workspace.id, account_id="other-account") is None

api/tests/unit_tests/services/auth/test_data_source_api_key_auth_gateways.py

@@ -9,7 +9,7 @@ from services.auth.errors import (
     DataSourceApiKeyAuthProviderUnavailableError,
     UnsupportedDataSourceApiKeyAuthProviderError,
 )
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 class TestProviderApiKeyAuthCredentialValidator:

api/tests/unit_tests/services/auth/test_data_source_api_key_auth_service.py

@@ -5,7 +5,7 @@ from unittest.mock import MagicMock
 
 from machinery.context import RequestContext
 from services.auth.data_source_api_key_auth_service import DataSourceApiKeyAuthService
-from services.entities.data_source_api_key_auth_entities import (
+from services.entities.data_source.api_key_auth import (
     DataSourceApiKeyAuthBindingCreate,
     DataSourceApiKeyAuthBindingRecord,
     DataSourceApiKeyAuthCredentials,

api/tests/unit_tests/services/auth/test_firecrawl_auth.py

@@ -10,7 +10,7 @@ from services.auth.errors import (
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
 from services.auth.firecrawl.firecrawl import FirecrawlAuth
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 def _credentials(

api/tests/unit_tests/services/auth/test_jina_auth.py

@@ -10,7 +10,7 @@ from services.auth.errors import (
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
 from services.auth.jina.jina import JinaAuth
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 def _credentials(

api/tests/unit_tests/services/auth/test_watercrawl_auth.py

@@ -10,7 +10,7 @@ from services.auth.errors import (
     InvalidDataSourceApiKeyAuthCredentialsError,
 )
 from services.auth.watercrawl.watercrawl import WatercrawlAuth
-from services.entities.data_source_api_key_auth_entities import DataSourceApiKeyAuthCredentials
+from services.entities.data_source.api_key_auth import DataSourceApiKeyAuthCredentials
 
 
 def _credentials(

api/tests/unit_tests/services/data_source/__init__.py

@@ -0,0 +1 @@
+"""Tests for data-source application services and adapters."""

api/tests/unit_tests/services/data_source/test_binding_application_service.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass, field
+from datetime import datetime
+
+import pytest
+
+from machinery.context import RequestContext
+from services.data_source.binding_application_service import (
+    BindingMutationResult,
+    DataSourceBindingApplicationService,
+    DataSourceBindingNotFoundError,
+    DataSourceBindingStateError,
+)
+from services.entities.data_source.oauth import DataSourceBindingSummary
+
+
+def _context() -> RequestContext:
+    return RequestContext("request-1", None, "account-1", "workspace-1")
+
+
+@dataclass
+class InMemoryBindingStore:
+    bindings: tuple[DataSourceBindingSummary, ...] = ()
+    mutation_result: BindingMutationResult = "updated"
+    changes: list[tuple[str, str, bool]] = field(default_factory=list)
+
+    def list_enabled_bindings(self, *, workspace_id: str) -> tuple[DataSourceBindingSummary, ...]:
+        assert workspace_id == "workspace-1"
+        return self.bindings
+
+    def change_disabled_state(
+        self,
+        *,
+        workspace_id: str,
+        binding_id: str,
+        disabled: bool,
+    ) -> BindingMutationResult:
+        self.changes.append((workspace_id, binding_id, disabled))
+        return self.mutation_result
+
+
+def test_list_integrations_returns_secret_free_records() -> None:
+    binding = DataSourceBindingSummary(
+        id="binding-1",
+        provider="notion",
+        created_at=datetime(2026, 1, 1),
+        disabled=False,
+        source_info={"workspace_id": "notion-workspace"},
+    )
+    unsupported = DataSourceBindingSummary(
+        id="binding-2",
+        provider="unsupported",
+        created_at=datetime(2026, 1, 1),
+        disabled=False,
+        source_info={},
+    )
+    service = DataSourceBindingApplicationService(bindings=InMemoryBindingStore(bindings=(binding, unsupported)))
+
+    assert service.list_integrations(_context()) == (binding,)
+
+
+@pytest.mark.parametrize(("operation", "disabled"), [("enable", False), ("disable", True)])
+def test_change_state_passes_stable_context(operation: str, disabled: bool) -> None:
+    store = InMemoryBindingStore()
+    service = DataSourceBindingApplicationService(bindings=store)
+
+    getattr(service, operation)(_context(), "binding-1")
+
+    assert store.changes == [("workspace-1", "binding-1", disabled)]
+
+
+@pytest.mark.parametrize(
+    ("operation", "result", "error", "message"),
+    [
+        ("enable", "not_found", DataSourceBindingNotFoundError, "Data source binding not found"),
+        ("enable", "already_enabled", DataSourceBindingStateError, "Data source is already enabled"),
+        ("disable", "already_disabled", DataSourceBindingStateError, "Data source is already disabled"),
+    ],
+)
+def test_change_state_translates_repository_outcomes(
+    operation: str,
+    result: BindingMutationResult,
+    error: type[Exception],
+    message: str,
+) -> None:
+    service = DataSourceBindingApplicationService(bindings=InMemoryBindingStore(mutation_result=result))
+
+    with pytest.raises(error, match=message):
+        getattr(service, operation)(_context(), "binding-1")

api/tests/unit_tests/services/data_source/test_credential_gateway.py

@@ -0,0 +1,499 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, field, replace
+from datetime import datetime, timedelta
+from typing import cast
+
+import pytest
+from sqlalchemy.orm import Session, sessionmaker
+
+import services.data_source.credential_gateway as credential_gateway_module
+from core.datasource.entities.datasource_entities import (
+    DatasourceProviderEntityWithPlugin,
+    DatasourceProviderIdentity,
+    DatasourceProviderType,
+)
+from core.entities.provider_entities import ProviderConfig, ProviderConfigType
+from core.plugin.entities.oauth import OAuthSchema
+from core.plugin.entities.plugin_daemon import PluginDatasourceProviderEntity, PluginOAuthCredentialsResponse
+from core.plugin.impl.datasource import PluginDatasourceManager
+from core.plugin.impl.oauth import OAuthHandler
+from models.oauth import DatasourceOauthParamConfig, DatasourceOauthTenantParamConfig
+from services.data_source.credential_gateway import (
+    ActorAwareDatasourceCredentialGateway,
+    DatasourceCredentialConcurrentUpdateError,
+    DatasourceCredentialError,
+    DatasourceCredentialNotFoundError,
+    DatasourceCredentialRefreshError,
+    OAuthDatasourceCredentialRefresher,
+    PluginDatasourceCredentialCodec,
+    RefreshedDatasourceCredential,
+)
+from services.entities.data_source.credential import DatasourceCredentialRecord
+
+
+def _record(
+    *,
+    auth_type: str = "oauth2",
+    expires_at: int = -1,
+    updated_at: datetime | None = None,
+) -> DatasourceCredentialRecord:
+    return DatasourceCredentialRecord(
+        id="credential-1",
+        workspace_id="workspace-1",
+        name="Notion",
+        provider="notion_datasource",
+        plugin_id="langgenius/notion_datasource",
+        auth_type=auth_type,
+        encrypted_credentials={"integration_secret": "encrypted"},
+        expires_at=expires_at,
+        updated_at=updated_at or datetime(2026, 1, 1),
+    )
+
+
+@dataclass
+class ScriptedCredentialStore:
+    records: list[DatasourceCredentialRecord | None]
+    update_result: bool = True
+    get_calls: list[tuple[str, str, str, str, str]] = field(default_factory=list)
+    updates: list[tuple[DatasourceCredentialRecord, Mapping[str, object], int]] = field(default_factory=list)
+    events: list[str] = field(default_factory=list)
+
+    def get_visible(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> DatasourceCredentialRecord | None:
+        self.events.append("read")
+        self.get_calls.append((workspace_id, actor_id, credential_id, provider, plugin_id))
+        return self.records.pop(0)
+
+    def update_if_unchanged(
+        self,
+        *,
+        record: DatasourceCredentialRecord,
+        encrypted_credentials: Mapping[str, object],
+        expires_at: int,
+    ) -> bool:
+        self.events.append("update")
+        self.updates.append((record, encrypted_credentials, expires_at))
+        return self.update_result
+
+
+@dataclass
+class ReversibleCredentialCodec:
+    def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]:
+        return {"integration_secret": f"plain:{record.encrypted_credentials['integration_secret']}"}
+
+    def encrypt(self, _record: DatasourceCredentialRecord, credentials: Mapping[str, object]) -> dict[str, object]:
+        return {"integration_secret": f"cipher:{credentials['integration_secret']}"}
+
+
+@dataclass
+class RecordingCredentialRefresher:
+    calls: list[tuple[str, str, DatasourceCredentialRecord, Mapping[str, object]]] = field(default_factory=list)
+    events: list[str] = field(default_factory=list)
+
+    def refresh(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        record: DatasourceCredentialRecord,
+        credentials: Mapping[str, object],
+    ) -> RefreshedDatasourceCredential:
+        self.events.append("refresh")
+        self.calls.append((workspace_id, actor_id, record, credentials))
+        return RefreshedDatasourceCredential({"integration_secret": "refreshed"}, 500)
+
+
+@dataclass
+class StaticProviderManager:
+    provider: PluginDatasourceProviderEntity
+    calls: list[tuple[str, str]] = field(default_factory=list)
+
+    def fetch_datasource_provider(self, *, tenant_id: str, provider_id: str) -> PluginDatasourceProviderEntity:
+        self.calls.append((tenant_id, provider_id))
+        return self.provider
+
+
+@dataclass
+class RecordingOAuthHandler:
+    response: PluginOAuthCredentialsResponse | None = None
+    error: Exception | None = None
+    calls: list[dict[str, object]] = field(default_factory=list)
+
+    def refresh_credentials(
+        self,
+        tenant_id: str,
+        user_id: str,
+        plugin_id: str,
+        provider: str,
+        redirect_uri: str,
+        system_credentials: Mapping[str, object],
+        credentials: Mapping[str, object],
+    ) -> PluginOAuthCredentialsResponse:
+        self.calls.append(
+            {
+                "tenant_id": tenant_id,
+                "user_id": user_id,
+                "plugin_id": plugin_id,
+                "provider": provider,
+                "redirect_uri": redirect_uri,
+                "system_credentials": system_credentials,
+                "credentials": credentials,
+            }
+        )
+        if self.error is not None:
+            raise self.error
+        assert self.response is not None
+        return self.response
+
+
+@dataclass
+class RecordingCredentialEncrypter:
+    decrypted: Mapping[str, object]
+    calls: list[Mapping[str, object]] = field(default_factory=list)
+
+    def decrypt(self, credentials: Mapping[str, object]) -> Mapping[str, object]:
+        self.calls.append(credentials)
+        return self.decrypted
+
+
+def _provider(*, with_oauth_schema: bool = True) -> PluginDatasourceProviderEntity:
+    oauth_schema = (
+        OAuthSchema(
+            client_schema=[ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="client_secret")],
+            credentials_schema=[
+                ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="integration_secret"),
+                ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="workspace_name"),
+            ],
+        )
+        if with_oauth_schema
+        else None
+    )
+    return PluginDatasourceProviderEntity(
+        provider="notion_datasource",
+        plugin_unique_identifier="langgenius/notion_datasource:1.0.0",
+        plugin_id="langgenius/notion_datasource",
+        declaration=DatasourceProviderEntityWithPlugin(
+            identity=DatasourceProviderIdentity(
+                author="langgenius",
+                name="notion_datasource",
+                description={"en_US": "Notion", "zh_Hans": "Notion"},
+                icon="notion.svg",
+                label={"en_US": "Notion", "zh_Hans": "Notion"},
+            ),
+            credentials_schema=[
+                ProviderConfig(type=ProviderConfigType.SECRET_INPUT, name="api_secret"),
+                ProviderConfig(type=ProviderConfigType.TEXT_INPUT, name="label"),
+            ],
+            oauth_schema=oauth_schema,
+            provider_type=DatasourceProviderType.ONLINE_DOCUMENT,
+        ),
+    )
+
+
+def _gateway(
+    store: ScriptedCredentialStore, refresher: RecordingCredentialRefresher | None = None
+) -> tuple[ActorAwareDatasourceCredentialGateway, RecordingCredentialRefresher]:
+    refresh = refresher or RecordingCredentialRefresher()
+    return (
+        ActorAwareDatasourceCredentialGateway(
+            credentials=store,
+            codec=ReversibleCredentialCodec(),
+            refresher=refresh,
+            now=lambda: 100,
+        ),
+        refresh,
+    )
+
+
+def _resolve(gateway: ActorAwareDatasourceCredentialGateway) -> dict[str, object]:
+    return gateway.resolve(
+        workspace_id="workspace-1",
+        actor_id="actor-1",
+        credential_id="credential-1",
+        provider="notion_datasource",
+        plugin_id="langgenius/notion_datasource",
+    )
+
+
+def test_resolve_fails_closed_when_credential_is_not_visible() -> None:
+    gateway, _ = _gateway(ScriptedCredentialStore(records=[None]))
+
+    with pytest.raises(DatasourceCredentialNotFoundError):
+        _resolve(gateway)
+
+
+def test_resolve_decrypts_without_refreshing_non_expiring_credential() -> None:
+    store = ScriptedCredentialStore(records=[_record()])
+    gateway, refresher = _gateway(store)
+
+    assert _resolve(gateway) == {"integration_secret": "plain:encrypted"}
+    assert refresher.calls == []
+    assert store.updates == []
+
+
+def test_refresh_runs_between_read_and_conditional_update() -> None:
+    record = _record(expires_at=120)
+    events: list[str] = []
+    store = ScriptedCredentialStore(records=[record], events=events)
+    gateway, refresher = _gateway(store, RecordingCredentialRefresher(events=events))
+
+    assert _resolve(gateway) == {"integration_secret": "refreshed"}
+    assert refresher.calls == [("workspace-1", "actor-1", record, {"integration_secret": "plain:encrypted"})]
+    assert store.updates == [(record, {"integration_secret": "cipher:refreshed"}, 500)]
+    assert events == ["read", "refresh", "update"]
+
+
+def test_concurrent_refresh_uses_newer_non_expiring_snapshot() -> None:
+    stale = _record(expires_at=120)
+    latest = replace(
+        stale,
+        encrypted_credentials={"integration_secret": "latest"},
+        expires_at=500,
+        updated_at=stale.updated_at + timedelta(seconds=1),
+    )
+    store = ScriptedCredentialStore(records=[stale, latest], update_result=False)
+    gateway, _ = _gateway(store)
+
+    assert _resolve(gateway) == {"integration_secret": "plain:latest"}
+    assert store.get_calls == [
+        ("workspace-1", "actor-1", "credential-1", "notion_datasource", "langgenius/notion_datasource"),
+        ("workspace-1", "actor-1", "credential-1", "notion_datasource", "langgenius/notion_datasource"),
+    ]
+
+
+def test_concurrent_refresh_fails_closed_if_latest_snapshot_is_still_expired() -> None:
+    stale = _record(expires_at=120)
+    latest = replace(stale, updated_at=stale.updated_at + timedelta(seconds=1))
+    gateway, _ = _gateway(ScriptedCredentialStore(records=[stale, latest], update_result=False))
+
+    with pytest.raises(DatasourceCredentialConcurrentUpdateError):
+        _resolve(gateway)
+
+
+@pytest.mark.parametrize(
+    ("auth_type", "secret_name"),
+    [("api-key", "api_secret"), ("oauth2", "integration_secret")],
+)
+def test_plugin_codec_encrypts_and_decrypts_only_declared_secret_fields(
+    monkeypatch: pytest.MonkeyPatch,
+    auth_type: str,
+    secret_name: str,
+) -> None:
+    manager = StaticProviderManager(_provider())
+    codec = PluginDatasourceCredentialCodec(provider_manager=cast(PluginDatasourceManager, manager))
+    record = replace(
+        _record(auth_type=auth_type),
+        encrypted_credentials={secret_name: "cipher", "label": "plain"},
+    )
+    decrypt_calls: list[tuple[str, str]] = []
+    encrypt_calls: list[tuple[str, str]] = []
+
+    def decrypt_token(workspace_id: str, value: str) -> str:
+        decrypt_calls.append((workspace_id, value))
+        return f"decrypted:{value}"
+
+    def encrypt_token(workspace_id: str, value: str) -> str:
+        encrypt_calls.append((workspace_id, value))
+        return f"encrypted:{value}"
+
+    monkeypatch.setattr(credential_gateway_module.encrypter, "decrypt_token", decrypt_token)
+    monkeypatch.setattr(credential_gateway_module.encrypter, "encrypt_token", encrypt_token)
+
+    assert codec.decrypt(record) == {secret_name: "decrypted:cipher", "label": "plain"}
+    assert codec.encrypt(record, {secret_name: "secret", "label": "plain"}) == {
+        secret_name: "encrypted:secret",
+        "label": "plain",
+    }
+    assert decrypt_calls == [("workspace-1", "cipher")]
+    assert encrypt_calls == [("workspace-1", "secret")]
+
+    non_string_record = replace(record, encrypted_credentials={secret_name: 42})
+    assert codec.decrypt(non_string_record) == {secret_name: 42}
+    assert codec.encrypt(record, {secret_name: None}) == {secret_name: None}
+    assert decrypt_calls == [("workspace-1", "cipher")]
+    assert encrypt_calls == [("workspace-1", "secret")]
+
+
+def test_plugin_codec_rejects_unsupported_credential_type() -> None:
+    codec = PluginDatasourceCredentialCodec(
+        provider_manager=cast(PluginDatasourceManager, StaticProviderManager(_provider()))
+    )
+
+    with pytest.raises(DatasourceCredentialError, match="Unsupported credential type"):
+        codec.decrypt(_record(auth_type="unauthorized"))
+
+
+def test_oauth_refresher_prefers_enabled_tenant_client_and_returns_refreshed_snapshot(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add_all(
+            [
+                DatasourceOauthTenantParamConfig(
+                    tenant_id="workspace-1",
+                    provider="notion_datasource",
+                    plugin_id="langgenius/notion_datasource",
+                    client_params={"client_secret": "tenant-cipher"},
+                    enabled=True,
+                ),
+                DatasourceOauthParamConfig(
+                    provider="notion_datasource",
+                    plugin_id="langgenius/notion_datasource",
+                    system_credentials={"client_secret": "system-secret"},
+                ),
+            ]
+        )
+    manager = StaticProviderManager(_provider())
+    handler = RecordingOAuthHandler(
+        response=PluginOAuthCredentialsResponse(
+            credentials={"integration_secret": "refreshed-secret"},
+            expires_at=500,
+        )
+    )
+    encrypter = RecordingCredentialEncrypter({"client_secret": "tenant-secret"})
+
+    def create_encrypter(**kwargs: object) -> tuple[RecordingCredentialEncrypter, None]:
+        assert kwargs["tenant_id"] == "workspace-1"
+        return encrypter, None
+
+    monkeypatch.setattr(credential_gateway_module, "create_provider_encrypter", create_encrypter)
+    refresher = OAuthDatasourceCredentialRefresher(
+        session_factory=sqlite_session_factory,
+        provider_manager=cast(PluginDatasourceManager, manager),
+        oauth_handler=cast(OAuthHandler, handler),
+    )
+
+    result = refresher.refresh(
+        workspace_id="workspace-1",
+        actor_id="actor-1",
+        record=_record(),
+        credentials={"integration_secret": "old-secret"},
+    )
+
+    assert result == RefreshedDatasourceCredential({"integration_secret": "refreshed-secret"}, 500)
+    assert encrypter.calls == [{"client_secret": "tenant-cipher"}]
+    assert handler.calls[0]["system_credentials"] == {"client_secret": "tenant-secret"}
+    assert handler.calls[0]["credentials"] == {"integration_secret": "old-secret"}
+    assert handler.calls[0]["user_id"] == "actor-1"
+
+
+def test_oauth_refresher_rejects_tenant_client_when_provider_has_no_oauth_schema(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(
+            DatasourceOauthTenantParamConfig(
+                tenant_id="workspace-1",
+                provider="notion_datasource",
+                plugin_id="langgenius/notion_datasource",
+                client_params={"client_secret": "tenant-cipher"},
+                enabled=True,
+            )
+        )
+    provider = _provider(with_oauth_schema=False)
+    handler = RecordingOAuthHandler()
+    refresher = OAuthDatasourceCredentialRefresher(
+        session_factory=sqlite_session_factory,
+        provider_manager=cast(PluginDatasourceManager, StaticProviderManager(provider)),
+        oauth_handler=cast(OAuthHandler, handler),
+    )
+
+    with pytest.raises(DatasourceCredentialError, match="oauth schema not found"):
+        refresher.refresh(
+            workspace_id="workspace-1",
+            actor_id="actor-1",
+            record=_record(),
+            credentials={"integration_secret": "old"},
+        )
+
+    assert handler.calls == []
+
+
+@pytest.mark.parametrize("verified", [True, False])
+def test_oauth_refresher_uses_system_client_only_for_verified_plugin(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_session_factory: sessionmaker[Session],
+    verified: bool,
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add_all(
+            [
+                DatasourceOauthTenantParamConfig(
+                    tenant_id="workspace-1",
+                    provider="notion_datasource",
+                    plugin_id="langgenius/notion_datasource",
+                    client_params={"client_secret": "disabled-tenant-secret"},
+                    enabled=False,
+                ),
+                DatasourceOauthParamConfig(
+                    provider="notion_datasource",
+                    plugin_id="langgenius/notion_datasource",
+                    system_credentials={"client_secret": "system-secret"},
+                ),
+            ]
+        )
+    manager = StaticProviderManager(_provider())
+    handler = RecordingOAuthHandler(
+        response=PluginOAuthCredentialsResponse(credentials={"integration_secret": "refreshed"}, expires_at=500)
+    )
+    verification_calls: list[tuple[str, str]] = []
+
+    def is_plugin_verified(workspace_id: str, plugin_unique_identifier: str) -> bool:
+        verification_calls.append((workspace_id, plugin_unique_identifier))
+        return verified
+
+    monkeypatch.setattr(credential_gateway_module.PluginService, "is_plugin_verified", is_plugin_verified)
+    refresher = OAuthDatasourceCredentialRefresher(
+        session_factory=sqlite_session_factory,
+        provider_manager=cast(PluginDatasourceManager, manager),
+        oauth_handler=cast(OAuthHandler, handler),
+    )
+
+    if not verified:
+        with pytest.raises(DatasourceCredentialError, match="OAuth client is not configured"):
+            refresher.refresh(
+                workspace_id="workspace-1",
+                actor_id="actor-1",
+                record=_record(),
+                credentials={"integration_secret": "old"},
+            )
+        assert handler.calls == []
+    else:
+        refresher.refresh(
+            workspace_id="workspace-1",
+            actor_id="actor-1",
+            record=_record(),
+            credentials={"integration_secret": "old"},
+        )
+        assert handler.calls[0]["system_credentials"] == {"client_secret": "system-secret"}
+
+    assert verification_calls == [("workspace-1", "langgenius/notion_datasource:1.0.0")]
+
+
+def test_oauth_refresher_translates_plugin_refresh_failure(
+    monkeypatch: pytest.MonkeyPatch,
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    handler = RecordingOAuthHandler(error=RuntimeError("plugin unavailable"))
+    refresher = OAuthDatasourceCredentialRefresher(
+        session_factory=sqlite_session_factory,
+        provider_manager=cast(PluginDatasourceManager, StaticProviderManager(_provider())),
+        oauth_handler=cast(OAuthHandler, handler),
+    )
+    monkeypatch.setattr(refresher, "_oauth_client", lambda _workspace_id, _provider_id: {})
+
+    with pytest.raises(DatasourceCredentialRefreshError, match="credential-1"):
+        refresher.refresh(
+            workspace_id="workspace-1",
+            actor_id="actor-1",
+            record=_record(),
+            credentials={"integration_secret": "old"},
+        )

api/tests/unit_tests/services/data_source/test_notion_import_adapters.py

@@ -0,0 +1,210 @@
+from collections.abc import Generator
+from dataclasses import dataclass, field
+
+import pytest
+
+from core.datasource.__base.datasource_runtime import DatasourceRuntime
+from core.datasource.entities.datasource_entities import (
+    DatasourceProviderType,
+    OnlineDocumentInfo,
+    OnlineDocumentPage,
+    OnlineDocumentPagesMessage,
+)
+from core.rag.models.document import Document
+from services.data_source.notion_import_adapters import NotionCredentialDataError, PluginNotionSourceGateway
+
+
+@dataclass
+class RecordingCredentialResolver:
+    credentials: dict[str, object]
+    calls: list[tuple[str, str, str, str, str]] = field(default_factory=list)
+
+    def resolve(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> dict[str, object]:
+        self.calls.append((workspace_id, actor_id, credential_id, provider, plugin_id))
+        return self.credentials
+
+
+@dataclass
+class RecordingDatasourceRuntime:
+    messages: tuple[OnlineDocumentPagesMessage, ...]
+    runtime: DatasourceRuntime = field(default_factory=lambda: DatasourceRuntime(tenant_id="workspace-1"))
+    users: list[str] = field(default_factory=list)
+
+    def get_online_document_pages(
+        self, user_id: str, datasource_parameters: dict[str, object], provider_type: str
+    ) -> Generator[OnlineDocumentPagesMessage, None, None]:
+        assert datasource_parameters == {}
+        assert provider_type == DatasourceProviderType.ONLINE_DOCUMENT
+        self.users.append(user_id)
+        yield from self.messages
+
+    def datasource_provider_type(self) -> str:
+        return DatasourceProviderType.ONLINE_DOCUMENT
+
+
+@dataclass
+class StaticRuntimeLoader:
+    runtime: RecordingDatasourceRuntime
+    calls: list[tuple[str, str, str, DatasourceProviderType]] = field(default_factory=list)
+
+    def __call__(
+        self,
+        *,
+        provider_id: str,
+        datasource_name: str,
+        tenant_id: str,
+        datasource_type: DatasourceProviderType,
+    ) -> object:
+        self.calls.append((provider_id, datasource_name, tenant_id, datasource_type))
+        return self.runtime
+
+
+@dataclass
+class StaticExtractor:
+    documents: list[Document]
+
+    def extract(self) -> list[Document]:
+        return self.documents
+
+
+@dataclass
+class RecordingExtractorFactory:
+    calls: list[dict[str, object]] = field(default_factory=list)
+
+    def __call__(
+        self,
+        notion_workspace_id: str,
+        notion_obj_id: str,
+        notion_page_type: str,
+        tenant_id: str,
+        document_model: object | None = None,
+        notion_access_token: str | None = None,
+        credential_id: str | None = None,
+    ) -> StaticExtractor:
+        self.calls.append(
+            {
+                "notion_workspace_id": notion_workspace_id,
+                "notion_obj_id": notion_obj_id,
+                "notion_page_type": notion_page_type,
+                "tenant_id": tenant_id,
+                "document_model": document_model,
+                "notion_access_token": notion_access_token,
+                "credential_id": credential_id,
+            }
+        )
+        return StaticExtractor([Document(page_content="one"), Document(page_content="two")])
+
+
+def _page(page_id: str, *, with_icon: bool = True) -> OnlineDocumentPage:
+    return OnlineDocumentPage(
+        page_id=page_id,
+        page_name=f"Page {page_id}",
+        page_icon={"type": "emoji", "emoji": "📄"} if with_icon else None,
+        type="page",
+        last_edited_time="2026-01-01T00:00:00Z",
+        parent_id=None,
+    )
+
+
+def test_list_authorized_pages_preserves_and_merges_workspace_groups() -> None:
+    runtime = RecordingDatasourceRuntime(
+        messages=(
+            OnlineDocumentPagesMessage(
+                result=[
+                    OnlineDocumentInfo(
+                        workspace_id="w1",
+                        workspace_name="One",
+                        workspace_icon=None,
+                        total=1,
+                        pages=[_page("p1")],
+                    )
+                ]
+            ),
+            OnlineDocumentPagesMessage(
+                result=[
+                    OnlineDocumentInfo(
+                        workspace_id="w2",
+                        workspace_name="Two",
+                        workspace_icon=None,
+                        total=1,
+                        pages=[_page("p2", with_icon=False)],
+                    ),
+                    OnlineDocumentInfo(
+                        workspace_id="w1",
+                        workspace_name="One",
+                        workspace_icon=None,
+                        total=1,
+                        pages=[_page("p3")],
+                    ),
+                ]
+            ),
+        )
+    )
+    resolver = RecordingCredentialResolver({"integration_secret": "secret"})
+    loader = StaticRuntimeLoader(runtime)
+    gateway = PluginNotionSourceGateway(credentials=resolver, runtime_loader=loader)
+
+    workspaces = gateway.list_authorized_pages(
+        workspace_id="workspace-1", actor_id="actor-1", credential_id="credential-1"
+    )
+
+    assert [workspace.workspace_id for workspace in workspaces] == ["w1", "w2"]
+    assert [page.page_id for page in workspaces[0].pages] == ["p1", "p3"]
+    assert workspaces[0].pages[0].page_icon is not None
+    assert workspaces[0].pages[0].page_icon.emoji == "📄"
+    assert workspaces[1].pages[0].page_icon is None
+    assert runtime.runtime.credentials == {"integration_secret": "secret"}
+    assert runtime.users == ["actor-1"]
+    assert resolver.calls == [
+        (
+            "workspace-1",
+            "actor-1",
+            "credential-1",
+            "notion_datasource",
+            "langgenius/notion_datasource",
+        )
+    ]
+
+
+def test_preview_requires_secret_and_passes_it_only_to_extractor() -> None:
+    factory = RecordingExtractorFactory()
+    gateway = PluginNotionSourceGateway(
+        credentials=RecordingCredentialResolver({"integration_secret": "secret"}),
+        runtime_loader=StaticRuntimeLoader(RecordingDatasourceRuntime(())),
+        extractor_factory=factory,
+    )
+
+    content = gateway.preview_page(
+        workspace_id="workspace-1",
+        actor_id="actor-1",
+        credential_id="credential-1",
+        page_id="page-1",
+        page_type="page",
+    )
+
+    assert content == "one\ntwo"
+    assert factory.calls[0]["notion_access_token"] == "secret"
+
+
+def test_preview_fails_closed_when_secret_is_missing() -> None:
+    gateway = PluginNotionSourceGateway(
+        credentials=RecordingCredentialResolver({}),
+        runtime_loader=StaticRuntimeLoader(RecordingDatasourceRuntime(())),
+    )
+
+    with pytest.raises(NotionCredentialDataError):
+        gateway.preview_page(
+            workspace_id="workspace-1",
+            actor_id="actor-1",
+            credential_id="credential-1",
+            page_id="page-1",
+            page_type="page",
+        )

api/tests/unit_tests/services/data_source/test_notion_import_application_service.py

@@ -0,0 +1,146 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from core.rag.entities.dataset_reference import DatasetRef
+from machinery.context import RequestContext
+from services.data_source.notion_import_application_service import (
+    DatasetIsNotNotionSourceError,
+    NotionImportApplicationService,
+)
+from services.entities.data_source.notion_import import AuthorizedNotionPage, NotionWorkspace
+from services.entities.knowledge_entities.records import DatasetRecord
+
+
+def _context() -> RequestContext:
+    return RequestContext("request-1", None, "account-1", "workspace-1")
+
+
+def _dataset(*, data_source_type: str = "notion_import") -> DatasetRecord:
+    return DatasetRecord(
+        id="dataset-1",
+        workspace_id="workspace-1",
+        maintainer_id="account-1",
+        permission="only_me",
+        data_source_type=data_source_type,
+        indexing_technique="high_quality",
+        embedding_model=None,
+        embedding_model_provider=None,
+    )
+
+
+@dataclass
+class DatasetAccessStub:
+    dataset: DatasetRecord = field(default_factory=_dataset)
+    calls: list[str] = field(default_factory=list)
+
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord:
+        assert context == _context()
+        self.calls.append(dataset_id)
+        return self.dataset
+
+
+@dataclass
+class DocumentBindingReaderStub:
+    page_ids: frozenset[str] = frozenset()
+    refs: list[DatasetRef] = field(default_factory=list)
+
+    def list_bound_notion_page_ids(self, dataset_ref: DatasetRef) -> frozenset[str]:
+        self.refs.append(dataset_ref)
+        return self.page_ids
+
+
+@dataclass
+class NotionSourceStub:
+    workspaces: tuple[NotionWorkspace, ...] = ()
+    list_calls: list[tuple[str, str, str]] = field(default_factory=list)
+    preview_calls: list[tuple[str, str, str, str, str]] = field(default_factory=list)
+
+    def list_authorized_pages(
+        self, *, workspace_id: str, actor_id: str, credential_id: str
+    ) -> tuple[NotionWorkspace, ...]:
+        self.list_calls.append((workspace_id, actor_id, credential_id))
+        return self.workspaces
+
+    def preview_page(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        page_id: str,
+        page_type: str,
+    ) -> str:
+        self.preview_calls.append((workspace_id, actor_id, credential_id, page_id, page_type))
+        return "preview"
+
+
+def _service(
+    *,
+    access: DatasetAccessStub | None = None,
+    documents: DocumentBindingReaderStub | None = None,
+    source: NotionSourceStub | None = None,
+) -> NotionImportApplicationService:
+    return NotionImportApplicationService(
+        dataset_access=access or DatasetAccessStub(),
+        documents=documents or DocumentBindingReaderStub(),
+        source=source or NotionSourceStub(),
+    )
+
+
+def test_list_pages_preserves_workspace_groups_and_marks_bound_pages() -> None:
+    page_1 = AuthorizedNotionPage("page-1", "One", None, None, "page")
+    page_2 = AuthorizedNotionPage("page-2", "Two", None, "page-1", "page")
+    source = NotionSourceStub(
+        workspaces=(
+            NotionWorkspace("notion-workspace-1", "First", None, (page_1,)),
+            NotionWorkspace("notion-workspace-2", "Second", None, (page_2,)),
+        )
+    )
+    documents = DocumentBindingReaderStub(page_ids=frozenset({"page-2"}))
+
+    result = _service(documents=documents, source=source).list_pages(
+        _context(), credential_id="credential-1", dataset_id="dataset-1"
+    )
+
+    assert [workspace.workspace_id for workspace in result.workspaces] == [
+        "notion-workspace-1",
+        "notion-workspace-2",
+    ]
+    assert result.workspaces[0].pages[0].is_bound is False
+    assert result.workspaces[1].pages[0].is_bound is True
+    assert documents.refs == [DatasetRef("workspace-1", "dataset-1")]
+    assert source.list_calls == [("workspace-1", "account-1", "credential-1")]
+
+
+def test_list_pages_without_dataset_does_not_read_document_bindings() -> None:
+    access = DatasetAccessStub()
+    documents = DocumentBindingReaderStub()
+
+    _service(access=access, documents=documents).list_pages(_context(), credential_id="credential-1")
+
+    assert access.calls == []
+    assert documents.refs == []
+
+
+def test_list_pages_rejects_non_notion_dataset_before_external_call() -> None:
+    access = DatasetAccessStub(dataset=_dataset(data_source_type="upload_file"))
+    source = NotionSourceStub()
+
+    with pytest.raises(DatasetIsNotNotionSourceError):
+        _service(access=access, source=source).list_pages(
+            _context(), credential_id="credential-1", dataset_id="dataset-1"
+        )
+
+    assert source.list_calls == []
+
+
+def test_preview_passes_explicit_actor_and_workspace() -> None:
+    source = NotionSourceStub()
+
+    content = _service(source=source).preview_page(
+        _context(), credential_id="credential-1", page_id="page-1", page_type="page"
+    )
+
+    assert content == "preview"
+    assert source.preview_calls == [("workspace-1", "account-1", "credential-1", "page-1", "page")]

api/tests/unit_tests/services/data_source/test_notion_oauth_gateway.py

@@ -0,0 +1,401 @@
+import json
+from dataclasses import dataclass
+from typing import Never, cast
+
+import httpx
+import pytest
+
+from core.helper.ssrf_proxy import SSRFProxy
+from core.tools.errors import ToolSSRFError
+from services.data_source.notion_oauth_gateway import NotionDataSourceGateway
+from services.data_source.oauth_service import DataSourceOAuthError
+
+
+@dataclass(frozen=True)
+class RejectingSSRFClient:
+    error: Exception
+
+    def post(self, *_args: object, **_kwargs: object) -> Never:
+        raise self.error
+
+
+def _gateway(handler: httpx.MockTransport) -> tuple[NotionDataSourceGateway, httpx.Client]:
+    client = httpx.Client(transport=handler)
+    return (
+        NotionDataSourceGateway(
+            client_id="client-id",
+            client_secret="client-secret",
+            redirect_uri="https://api.example/oauth/data-source/callback/notion",
+            http_client=client,
+        ),
+        client,
+    )
+
+
+def test_get_authorization_url_encodes_configured_values() -> None:
+    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(500)))
+    try:
+        result = gateway.get_authorization_url()
+    finally:
+        client.close()
+
+    assert result == (
+        "https://api.notion.com/v1/oauth/authorize?client_id=client-id&response_type=code&"
+        "redirect_uri=https%3A%2F%2Fapi.example%2Foauth%2Fdata-source%2Fcallback%2Fnotion&owner=user"
+    )
+
+
+def test_authorize_exchanges_code_and_builds_source_info_without_persistence() -> None:
+    search_types: list[str] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/v1/oauth/token":
+            assert request.headers["authorization"].startswith("Basic ")
+            return httpx.Response(
+                200,
+                json={
+                    "access_token": "notion-token",
+                    "workspace_name": "Notion Workspace",
+                    "workspace_icon": "https://example/icon.png",
+                    "workspace_id": "notion-workspace",
+                },
+            )
+        if request.url.path == "/v1/search":
+            payload = json.loads(request.content)
+            search_types.append(payload["filter"]["value"])
+            return httpx.Response(200, json={"results": [], "has_more": False})
+        raise AssertionError(f"Unexpected Notion request: {request.url}")
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+    try:
+        authorization = gateway.authorize("authorization-code")
+    finally:
+        client.close()
+
+    assert authorization.access_token == "notion-token"
+    assert authorization.source_info == {
+        "workspace_name": "Notion Workspace",
+        "workspace_icon": "https://example/icon.png",
+        "workspace_id": "notion-workspace",
+        "pages": [],
+        "total": 0,
+    }
+    assert search_types == ["page", "database"]
+
+
+@pytest.mark.parametrize("token_payload", [{}, {"access_token": ""}, {"access_token": 42}])
+def test_authorize_rejects_missing_or_invalid_access_token(token_payload: dict[str, object]) -> None:
+    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=token_payload)))
+    try:
+        with pytest.raises(DataSourceOAuthError, match="did not include an access token"):
+            gateway.authorize("authorization-code")
+    finally:
+        client.close()
+
+
+def test_authorize_internal_reads_workspace_name_and_builds_source_info() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/v1/users/me":
+            assert request.headers["authorization"] == "Bearer notion-token"
+            return httpx.Response(
+                200,
+                json={"object": "user", "type": "bot", "bot": {"workspace_name": "Internal Workspace"}},
+            )
+        if request.url.path == "/v1/search":
+            return httpx.Response(200, json={"results": [], "has_more": False})
+        raise AssertionError(f"Unexpected Notion request: {request.url}")
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+    try:
+        authorization = gateway.authorize_internal("notion-token", "notion-workspace")
+    finally:
+        client.close()
+
+    assert authorization.access_token == "notion-token"
+    assert authorization.source_info == {
+        "workspace_name": "Internal Workspace",
+        "workspace_icon": None,
+        "workspace_id": "notion-workspace",
+        "pages": [],
+        "total": 0,
+    }
+
+
+@pytest.mark.parametrize(
+    ("operation", "expected_message"),
+    [
+        ("authorize", "Notion OAuth response was invalid"),
+        ("authorize_internal", "Notion integration response was invalid"),
+        ("refresh", "Notion integration response was invalid"),
+    ],
+)
+def test_public_operations_translate_malformed_page_payload(
+    operation: str,
+    expected_message: str,
+) -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/v1/oauth/token":
+            return httpx.Response(200, json={"access_token": "notion-token"})
+        if request.url.path == "/v1/users/me":
+            return httpx.Response(200, json={"object": "user", "type": "bot", "bot": {}})
+        if request.url.path == "/v1/search":
+            payload = json.loads(request.content)
+            results = [{}] if payload["filter"]["value"] == "page" else []
+            return httpx.Response(200, json={"results": results, "has_more": False})
+        raise AssertionError(f"Unexpected Notion request: {request.url}")
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+
+    def invoke() -> None:
+        if operation == "authorize":
+            gateway.authorize("authorization-code")
+        elif operation == "authorize_internal":
+            gateway.authorize_internal("notion-token", "notion-workspace")
+        else:
+            gateway.refresh("notion-token", {})
+
+    try:
+        with pytest.raises(DataSourceOAuthError, match=expected_message):
+            invoke()
+    finally:
+        client.close()
+
+
+def test_authorized_pages_map_pagination_icons_and_parent_chains() -> None:
+    page_requests: list[dict[str, object]] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        if request.url.path == "/v1/search":
+            payload = json.loads(request.content)
+            object_type = payload["filter"]["value"]
+            if object_type == "page":
+                page_requests.append(payload)
+                if "start_cursor" not in payload:
+                    return httpx.Response(
+                        200,
+                        json={
+                            "results": [
+                                {
+                                    "id": "page-1",
+                                    "properties": {"Name": {"title": [{"plain_text": "First page"}]}},
+                                    "icon": {"type": "external", "external": {"url": "/page-icon"}},
+                                    "parent": {"type": "block_id", "block_id": "child-block"},
+                                }
+                            ],
+                            "has_more": True,
+                            "next_cursor": "next-page",
+                        },
+                    )
+                return httpx.Response(
+                    200,
+                    json={
+                        "results": [
+                            {
+                                "id": "page-2",
+                                "properties": {},
+                                "icon": {"type": "emoji", "emoji": "📄"},
+                                "parent": {"type": "workspace", "workspace": True},
+                            }
+                        ],
+                        "has_more": False,
+                    },
+                )
+            return httpx.Response(
+                200,
+                json={
+                    "results": [
+                        {
+                            "id": "database-1",
+                            "title": [{"plain_text": "Database"}],
+                            "icon": {"type": "emoji", "emoji": "🗃️"},
+                            "parent": {"type": "page_id", "page_id": "page-1"},
+                        },
+                        {
+                            "id": "database-2",
+                            "title": [],
+                            "icon": None,
+                            "parent": {"type": "workspace", "workspace": True},
+                        },
+                    ],
+                    "has_more": False,
+                },
+            )
+        if request.url.path == "/v1/blocks/child-block":
+            return httpx.Response(200, json={"parent": {"type": "block_id", "block_id": "root-block"}})
+        if request.url.path == "/v1/blocks/root-block":
+            return httpx.Response(200, json={"parent": {"type": "page_id", "page_id": "parent-page"}})
+        raise AssertionError(f"Unexpected Notion request: {request.url}")
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+    try:
+        pages = gateway.get_authorized_pages("notion-token")
+    finally:
+        client.close()
+
+    assert pages == [
+        {
+            "page_id": "page-1",
+            "page_name": "First page",
+            "page_icon": {"type": "url", "url": "https://www.notion.so/page-icon"},
+            "parent_id": "parent-page",
+            "type": "page",
+        },
+        {
+            "page_id": "page-2",
+            "page_name": "Untitled",
+            "page_icon": {"type": "emoji", "emoji": "📄"},
+            "parent_id": "root",
+            "type": "page",
+        },
+        {
+            "page_id": "database-1",
+            "page_name": "Database",
+            "page_icon": {"type": "emoji", "emoji": "🗃️"},
+            "parent_id": "page-1",
+            "type": "database",
+        },
+        {
+            "page_id": "database-2",
+            "page_name": "Untitled",
+            "page_icon": None,
+            "parent_id": "root",
+            "type": "database",
+        },
+    ]
+    assert page_requests == [
+        {"filter": {"value": "page", "property": "object"}},
+        {"filter": {"value": "page", "property": "object"}, "start_cursor": "next-page"},
+    ]
+
+
+def test_refresh_preserves_workspace_metadata() -> None:
+    gateway, client = _gateway(
+        httpx.MockTransport(
+            lambda request: (
+                httpx.Response(200, json={"results": [], "has_more": False})
+                if request.url.path == "/v1/search"
+                else httpx.Response(404)
+            )
+        )
+    )
+    try:
+        source_info = gateway.refresh(
+            "notion-token",
+            {
+                "workspace_name": "Workspace",
+                "workspace_icon": None,
+                "workspace_id": "notion-workspace",
+                "pages": [
+                    {
+                        "page_id": "old",
+                        "page_name": "Old page",
+                        "page_icon": None,
+                        "parent_id": "root",
+                        "type": "page",
+                    }
+                ],
+                "total": 1,
+            },
+        )
+    finally:
+        client.close()
+
+    assert source_info == {
+        "workspace_name": "Workspace",
+        "workspace_icon": None,
+        "workspace_id": "notion-workspace",
+        "pages": [],
+        "total": 0,
+    }
+
+
+def test_search_rejects_missing_cursor_when_more_results_are_declared() -> None:
+    requests = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal requests
+        requests += 1
+        return httpx.Response(200, json={"results": [], "has_more": True, "next_cursor": None})
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+    try:
+        with pytest.raises(DataSourceOAuthError, match="valid next_cursor"):
+            gateway.notion_page_search("notion-token")
+    finally:
+        client.close()
+
+    assert requests == 1
+
+
+def test_search_rejects_non_list_results() -> None:
+    gateway, client = _gateway(
+        httpx.MockTransport(lambda _request: httpx.Response(200, json={"results": {}, "has_more": False}))
+    )
+    try:
+        with pytest.raises(DataSourceOAuthError, match="search response was invalid"):
+            gateway.notion_page_search("notion-token")
+    finally:
+        client.close()
+
+
+@pytest.mark.parametrize(
+    ("payload", "expected_message"),
+    [
+        ({"parent": []}, "block response was invalid"),
+        ({"parent": {"type": "page_id"}}, "block parent was invalid"),
+        ({"parent": {"type": "page_id", "page_id": 42}}, "parent identifier was invalid"),
+    ],
+)
+def test_block_parent_rejects_invalid_provider_payload(
+    payload: dict[str, object],
+    expected_message: str,
+) -> None:
+    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=payload)))
+    try:
+        with pytest.raises(DataSourceOAuthError, match=expected_message):
+            gateway.notion_block_parent_page_id("notion-token", "block-1")
+    finally:
+        client.close()
+
+
+def test_provider_http_failure_is_translated_to_application_boundary_error() -> None:
+    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(503, text="unavailable")))
+    try:
+        with pytest.raises(DataSourceOAuthError, match="unsuccessful response"):
+            gateway.authorize("authorization-code")
+    finally:
+        client.close()
+
+
+def test_provider_non_object_json_is_rejected() -> None:
+    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(200, json=[])))
+    try:
+        with pytest.raises(DataSourceOAuthError, match="response was invalid"):
+            gateway.authorize("authorization-code")
+    finally:
+        client.close()
+
+
+def test_provider_get_failure_is_translated_to_application_boundary_error() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        raise httpx.ConnectError("offline", request=request)
+
+    gateway, client = _gateway(httpx.MockTransport(handler))
+    try:
+        with pytest.raises(DataSourceOAuthError, match="Notion request failed"):
+            gateway.authorize_internal("notion-token", "notion-workspace")
+    finally:
+        client.close()
+
+
+def test_ssrf_rejection_is_translated_to_application_boundary_error() -> None:
+    client = RejectingSSRFClient(ToolSSRFError("blocked"))
+    gateway = NotionDataSourceGateway(
+        client_id="client-id",
+        client_secret="client-secret",
+        redirect_uri="https://api.example/oauth/data-source/callback/notion",
+        http_client=cast(SSRFProxy, client),
+    )
+
+    with pytest.raises(DataSourceOAuthError, match="Notion request failed"):
+        gateway.authorize("authorization-code")

api/tests/unit_tests/services/data_source/test_oauth_service.py

@@ -0,0 +1,236 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+
+import pytest
+
+from machinery.context import RequestContext
+from services.data_source.oauth_service import (
+    DataSourceOAuthConfigurationError,
+    DataSourceOAuthError,
+    DataSourceOAuthService,
+    InvalidDataSourceOAuthCodeError,
+)
+from services.entities.data_source.oauth import (
+    DataSourceOAuthAuthorization,
+    DataSourceOAuthBindingRecord,
+    DataSourceOAuthCallback,
+)
+
+
+def _authorization() -> DataSourceOAuthAuthorization:
+    return DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []})
+
+
+@dataclass
+class RecordingProviderGateway:
+    authorization_url: str = "https://notion.example/authorize"
+    authorization: DataSourceOAuthAuthorization = field(default_factory=_authorization)
+    refreshed_source_info: Mapping[str, object] = field(default_factory=lambda: {"pages": [{"page_id": "page-1"}]})
+    events: list[tuple[str, object]] = field(default_factory=list)
+
+    def get_authorization_url(self) -> str:
+        self.events.append(("get_authorization_url", ()))
+        return self.authorization_url
+
+    def authorize(self, code: str) -> DataSourceOAuthAuthorization:
+        self.events.append(("authorize", code))
+        return self.authorization
+
+    def authorize_internal(self, access_token: str, workspace_id: str) -> DataSourceOAuthAuthorization:
+        self.events.append(("authorize_internal", (access_token, workspace_id)))
+        return self.authorization
+
+    def refresh(self, access_token: str, source_info: Mapping[str, object]) -> Mapping[str, object]:
+        self.events.append(("refresh", (access_token, source_info)))
+        return self.refreshed_source_info
+
+
+@dataclass
+class InMemoryBindingRepository:
+    binding: DataSourceOAuthBindingRecord | None = None
+    update_result: bool = True
+    events: list[tuple[str, object]] = field(default_factory=list)
+
+    def upsert_authorization(
+        self,
+        *,
+        workspace_id: str,
+        provider: str,
+        authorization: DataSourceOAuthAuthorization,
+    ) -> None:
+        self.events.append(("persist", (workspace_id, provider, authorization)))
+
+    def get_enabled(
+        self,
+        *,
+        workspace_id: str,
+        provider: str,
+        binding_id: str,
+    ) -> DataSourceOAuthBindingRecord | None:
+        self.events.append(("read", (workspace_id, provider, binding_id)))
+        return self.binding
+
+    def update_source_info(
+        self,
+        *,
+        workspace_id: str,
+        provider: str,
+        binding_id: str,
+        source_info: Mapping[str, object],
+    ) -> bool:
+        self.events.append(("update", (workspace_id, provider, binding_id, source_info)))
+        return self.update_result
+
+
+@pytest.fixture
+def request_context() -> RequestContext:
+    return RequestContext(
+        request_id="request-1",
+        trace_id="trace-1",
+        account_id="account-1",
+        active_workspace_id="workspace-1",
+    )
+
+
+def test_start_authorization_delegates_to_external_provider(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    provider = RecordingProviderGateway(events=events)
+    bindings = InMemoryBindingRepository(events=events)
+    service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings)
+
+    result = service.start_authorization(request_context)
+
+    assert result == "https://notion.example/authorize"
+    assert events == [("get_authorization_url", ())]
+
+
+def test_start_authorization_persists_internal_authorization(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    authorization = DataSourceOAuthAuthorization(access_token="secret", source_info={"pages": []})
+    provider = RecordingProviderGateway(authorization=authorization, events=events)
+    bindings = InMemoryBindingRepository(events=events)
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=provider,
+        bindings=bindings,
+        is_internal_provider=True,
+        internal_access_token="secret",
+    )
+
+    result = service.start_authorization(request_context)
+
+    assert result == "internal"
+    assert events == [
+        ("authorize_internal", ("secret", "workspace-1")),
+        ("persist", ("workspace-1", "notion", authorization)),
+    ]
+
+
+def test_start_authorization_rejects_missing_internal_secret(request_context: RequestContext) -> None:
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=RecordingProviderGateway(),
+        bindings=InMemoryBindingRepository(),
+        is_internal_provider=True,
+    )
+
+    with pytest.raises(DataSourceOAuthConfigurationError, match="Internal secret is not set"):
+        service.start_authorization(request_context)
+
+
+@pytest.mark.parametrize(
+    ("code", "error", "expected_error"),
+    [
+        ("auth/code", None, None),
+        (None, "access denied", "access denied"),
+        (None, None, "Access denied"),
+    ],
+)
+def test_complete_callback_returns_framework_neutral_result(
+    code: str | None,
+    error: str | None,
+    expected_error: str | None,
+) -> None:
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=RecordingProviderGateway(),
+        bindings=InMemoryBindingRepository(),
+    )
+
+    assert service.complete_callback(code=code, error=error) == DataSourceOAuthCallback(
+        provider="notion",
+        code=code,
+        error=expected_error,
+    )
+
+
+def test_bind_authorizes_before_persisting(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    authorization = _authorization()
+    provider = RecordingProviderGateway(authorization=authorization, events=events)
+    bindings = InMemoryBindingRepository(events=events)
+    service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings)
+
+    service.bind(request_context, code="code-1")
+
+    assert events == [
+        ("authorize", "code-1"),
+        ("persist", ("workspace-1", "notion", authorization)),
+    ]
+
+
+def test_bind_rejects_empty_code(request_context: RequestContext) -> None:
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=RecordingProviderGateway(),
+        bindings=InMemoryBindingRepository(),
+    )
+
+    with pytest.raises(InvalidDataSourceOAuthCodeError):
+        service.bind(request_context, code="")
+
+
+def test_sync_reads_before_provider_io_and_persists_afterward(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []})
+    refreshed_source_info = {"pages": [{"page_id": "page-1"}]}
+    bindings = InMemoryBindingRepository(binding=binding, events=events)
+    provider = RecordingProviderGateway(refreshed_source_info=refreshed_source_info, events=events)
+    service = DataSourceOAuthService(provider_name="notion", provider_gateway=provider, bindings=bindings)
+
+    service.sync(request_context, binding_id="binding-1")
+
+    assert events == [
+        ("read", ("workspace-1", "notion", "binding-1")),
+        ("refresh", ("token", {"pages": []})),
+        ("update", ("workspace-1", "notion", "binding-1", refreshed_source_info)),
+    ]
+
+
+def test_sync_rejects_missing_binding_without_provider_io(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=RecordingProviderGateway(events=events),
+        bindings=InMemoryBindingRepository(events=events),
+    )
+
+    with pytest.raises(DataSourceOAuthError, match="Data source binding not found"):
+        service.sync(request_context, binding_id="missing")
+
+    assert events == [("read", ("workspace-1", "notion", "missing"))]
+
+
+def test_sync_rejects_binding_removed_before_update(request_context: RequestContext) -> None:
+    events: list[tuple[str, object]] = []
+    binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []})
+    service = DataSourceOAuthService(
+        provider_name="notion",
+        provider_gateway=RecordingProviderGateway(events=events),
+        bindings=InMemoryBindingRepository(binding=binding, update_result=False, events=events),
+    )
+
+    with pytest.raises(DataSourceOAuthError, match="Data source binding not found"):
+        service.sync(request_context, binding_id="binding-1")
+
+    assert [event for event, _ in events] == ["read", "refresh", "update"]

api/tests/unit_tests/services/knowledge/__init__.py

@@ -0,0 +1 @@
+"""Tests for knowledge application services and adapters."""

api/tests/unit_tests/services/knowledge/test_dataset_access_service.py

@@ -0,0 +1,149 @@
+from dataclasses import dataclass
+
+import pytest
+
+from machinery.context import RequestContext
+from services.entities.knowledge_entities.records import DatasetAccessSnapshot, DatasetRecord
+from services.knowledge.application import (
+    DatasetAccessDeniedError,
+    DatasetAccessService,
+    DatasetNotFoundError,
+)
+
+
+def _context(*, actor_id: str = "actor-1", workspace_id: str = "workspace-1") -> RequestContext:
+    return RequestContext(
+        request_id="request-1",
+        trace_id=None,
+        account_id=actor_id,
+        active_workspace_id=workspace_id,
+    )
+
+
+def _snapshot(
+    *,
+    permission: str,
+    maintainer_id: str = "maintainer-1",
+    partial_access: bool = False,
+) -> DatasetAccessSnapshot:
+    return DatasetAccessSnapshot(
+        dataset=DatasetRecord(
+            id="dataset-1",
+            workspace_id="workspace-1",
+            maintainer_id=maintainer_id,
+            permission=permission,
+            data_source_type="notion_import",
+            indexing_technique="economy",
+            embedding_model=None,
+            embedding_model_provider=None,
+        ),
+        actor_has_partial_access=partial_access,
+    )
+
+
+@dataclass
+class DatasetReaderStub:
+    snapshot: DatasetAccessSnapshot | None
+
+    def get_access_snapshot(
+        self,
+        *,
+        workspace_id: str,
+        dataset_id: str,
+        actor_id: str,
+    ) -> DatasetAccessSnapshot | None:
+        assert workspace_id == "workspace-1"
+        assert dataset_id == "dataset-1"
+        assert actor_id == "actor-1"
+        return self.snapshot
+
+
+@dataclass
+class WorkspaceRoleReaderStub:
+    role: str | None
+    calls: int = 0
+
+    def get_legacy_role(self, *, workspace_id: str, account_id: str) -> str | None:
+        assert workspace_id == "workspace-1"
+        assert account_id == "actor-1"
+        self.calls += 1
+        return self.role
+
+
+@pytest.mark.parametrize(
+    ("permission", "role", "maintainer_id", "partial_access"),
+    [
+        ("only_me", "owner", "maintainer-1", False),
+        ("only_me", "normal", "actor-1", False),
+        ("all_team_members", "normal", "maintainer-1", False),
+        ("partial_members", "normal", "maintainer-1", True),
+    ],
+)
+def test_require_accessible_allows_each_legacy_access_path(
+    permission: str,
+    role: str,
+    maintainer_id: str,
+    partial_access: bool,
+) -> None:
+    roles = WorkspaceRoleReaderStub(role)
+    service = DatasetAccessService(
+        datasets=DatasetReaderStub(
+            _snapshot(
+                permission=permission,
+                maintainer_id=maintainer_id,
+                partial_access=partial_access,
+            )
+        ),
+        workspace_roles=roles,
+        legacy_permissions_enabled=True,
+    )
+
+    result = service.require_accessible(_context(), "dataset-1")
+
+    assert result.id == "dataset-1"
+    assert roles.calls == 1
+
+
+@pytest.mark.parametrize(
+    ("permission", "partial_access"),
+    [
+        ("only_me", False),
+        ("partial_members", False),
+        ("unknown", True),
+    ],
+)
+def test_require_accessible_denies_unrecognized_or_ungranted_legacy_access(
+    permission: str,
+    partial_access: bool,
+) -> None:
+    service = DatasetAccessService(
+        datasets=DatasetReaderStub(_snapshot(permission=permission, partial_access=partial_access)),
+        workspace_roles=WorkspaceRoleReaderStub("normal"),
+        legacy_permissions_enabled=True,
+    )
+
+    with pytest.raises(DatasetAccessDeniedError):
+        service.require_accessible(_context(), "dataset-1")
+
+
+def test_require_accessible_only_enforces_owner_chain_when_legacy_permissions_are_disabled() -> None:
+    roles = WorkspaceRoleReaderStub("normal")
+    service = DatasetAccessService(
+        datasets=DatasetReaderStub(_snapshot(permission="only_me")),
+        workspace_roles=roles,
+        legacy_permissions_enabled=False,
+    )
+
+    assert service.require_accessible(_context(), "dataset-1").id == "dataset-1"
+    assert roles.calls == 0
+
+
+def test_require_accessible_hides_missing_or_cross_workspace_dataset() -> None:
+    service = DatasetAccessService(
+        datasets=DatasetReaderStub(None),
+        workspace_roles=WorkspaceRoleReaderStub("owner"),
+        legacy_permissions_enabled=True,
+    )
+
+    with pytest.raises(DatasetNotFoundError):
+        service.require_accessible(_context(), "dataset-1")

api/tests/unit_tests/services/knowledge/test_document_sync_application_service.py

@@ -0,0 +1,166 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef
+from machinery.context import RequestContext
+from services.entities.knowledge_entities.records import DatasetRecord, DocumentRecord
+from services.knowledge.adapters import CeleryDocumentSyncDispatcher
+from services.knowledge.application import (
+    DatasetAccessDeniedError,
+    DocumentSyncApplicationService,
+    SyncDocumentNotFoundError,
+    SyncDocumentSourceError,
+)
+
+
+def _context() -> RequestContext:
+    return RequestContext("request-1", None, "account-1", "workspace-1")
+
+
+def _dataset() -> DatasetRecord:
+    return DatasetRecord(
+        id="dataset-1",
+        workspace_id="workspace-1",
+        maintainer_id="account-1",
+        permission="only_me",
+        data_source_type="notion_import",
+        indexing_technique="high_quality",
+        embedding_model=None,
+        embedding_model_provider=None,
+    )
+
+
+def _document(*, data_source_type: str = "notion_import") -> DocumentRecord:
+    return DocumentRecord(
+        id="document-1",
+        workspace_id="workspace-1",
+        dataset_id="dataset-1",
+        name="Document",
+        data_source_type=data_source_type,
+        data_source_info={"notion_page_id": "page-1"},
+        enabled=True,
+        archived=False,
+        indexing_status="completed",
+        batch="batch-1",
+        doc_form="text_model",
+        doc_language="English",
+        dataset_process_rule_id=None,
+        need_summary=False,
+        doc_metadata=None,
+    )
+
+
+@dataclass
+class DatasetAccessStub:
+    error: Exception | None = None
+    calls: list[str] = field(default_factory=list)
+
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord:
+        assert context == _context()
+        self.calls.append(dataset_id)
+        if self.error is not None:
+            raise self.error
+        return _dataset()
+
+
+@dataclass
+class DocumentReaderStub:
+    document: DocumentRecord | None = field(default_factory=_document)
+    active_refs: tuple[DocumentRef, ...] = ()
+    list_calls: list[DatasetRef] = field(default_factory=list)
+    get_calls: list[DocumentRef] = field(default_factory=list)
+
+    def list_active_notion_refs(self, dataset_ref: DatasetRef) -> tuple[DocumentRef, ...]:
+        self.list_calls.append(dataset_ref)
+        return self.active_refs
+
+    def get_by_ref(self, document_ref: DocumentRef) -> DocumentRecord | None:
+        self.get_calls.append(document_ref)
+        return self.document
+
+
+@dataclass
+class DispatcherRecorder:
+    calls: list[tuple[str, str]] = field(default_factory=list)
+
+    def dispatch(self, *, dataset_id: str, document_id: str) -> None:
+        self.calls.append((dataset_id, document_id))
+
+
+def test_celery_dispatcher_forwards_owned_identifiers_to_injected_delay() -> None:
+    calls: list[tuple[str, str]] = []
+    dispatcher = CeleryDocumentSyncDispatcher(
+        delay=lambda dataset_id, document_id: calls.append((dataset_id, document_id))
+    )
+
+    dispatcher.dispatch(dataset_id="dataset-1", document_id="document-1")
+
+    assert calls == [("dataset-1", "document-1")]
+
+
+def test_sync_dataset_dispatches_snapshot_after_tenant_scoped_read() -> None:
+    dataset_ref = DatasetRef("workspace-1", "dataset-1")
+    documents = DocumentReaderStub(active_refs=(dataset_ref.document("document-1"), dataset_ref.document("document-2")))
+    dispatcher = DispatcherRecorder()
+    service = DocumentSyncApplicationService(
+        dataset_access=DatasetAccessStub(), documents=documents, dispatcher=dispatcher
+    )
+
+    count = service.sync_dataset(_context(), "dataset-1")
+
+    assert count == 2
+    assert documents.list_calls == [dataset_ref]
+    assert dispatcher.calls == [("dataset-1", "document-1"), ("dataset-1", "document-2")]
+
+
+def test_sync_document_validates_owner_chain_and_source_before_dispatch() -> None:
+    documents = DocumentReaderStub()
+    dispatcher = DispatcherRecorder()
+    service = DocumentSyncApplicationService(
+        dataset_access=DatasetAccessStub(), documents=documents, dispatcher=dispatcher
+    )
+
+    service.sync_document(_context(), "dataset-1", "document-1")
+
+    assert documents.get_calls == [DatasetRef("workspace-1", "dataset-1").document("document-1")]
+    assert dispatcher.calls == [("dataset-1", "document-1")]
+
+
+@pytest.mark.parametrize(
+    ("document", "error"),
+    [
+        (None, SyncDocumentNotFoundError),
+        (_document(data_source_type="upload_file"), SyncDocumentSourceError),
+    ],
+)
+def test_sync_document_rejects_invalid_document_before_dispatch(
+    document: DocumentRecord | None, error: type[Exception]
+) -> None:
+    dispatcher = DispatcherRecorder()
+    service = DocumentSyncApplicationService(
+        dataset_access=DatasetAccessStub(),
+        documents=DocumentReaderStub(document=document),
+        dispatcher=dispatcher,
+    )
+
+    with pytest.raises(error):
+        service.sync_document(_context(), "dataset-1", "document-1")
+
+    assert dispatcher.calls == []
+
+
+def test_sync_stops_before_document_read_and_dispatch_when_dataset_access_is_denied() -> None:
+    documents = DocumentReaderStub()
+    dispatcher = DispatcherRecorder()
+    service = DocumentSyncApplicationService(
+        dataset_access=DatasetAccessStub(error=DatasetAccessDeniedError()),
+        documents=documents,
+        dispatcher=dispatcher,
+    )
+
+    with pytest.raises(DatasetAccessDeniedError):
+        service.sync_document(_context(), "dataset-1", "document-1")
+
+    assert documents.get_calls == []
+    assert dispatcher.calls == []

api/tests/unit_tests/services/knowledge/test_indexing_estimate_adapters.py

@@ -0,0 +1,568 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
+
+import pytest
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.entities.knowledge_entities import IndexingEstimate
+from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
+from core.plugin.impl.exc import PluginDaemonClientSideError
+from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo
+from extensions.storage.storage_type import StorageType
+from models.dataset import DatasetProcessRule
+from models.enums import CreatorUserRole, ProcessRuleMode
+from models.model import UploadFile
+from services.entities.knowledge_entities.indexing_estimate import (
+    ExistingDocumentsEstimateCommand,
+    NewSourcesEstimateCommand,
+    NotionEstimateSource,
+    UploadFileEstimateSource,
+    WebsiteEstimateSource,
+)
+from services.entities.knowledge_entities.records import DatasetRecord, DocumentRecord
+from services.knowledge.adapters import (
+    EstimateSourceNotFoundError,
+    IndexingEstimateExecutionError,
+    IndexingEstimateProviderUnavailableError,
+    IndexingRunnerEstimateGateway,
+    UnsupportedEstimateSourceError,
+)
+
+
+@dataclass(frozen=True)
+class RunnerCall:
+    tenant_id: str
+    source_types: tuple[str, ...]
+    file_ids: tuple[str | None, ...]
+    notion_page_ids: tuple[str | None, ...]
+    notion_tokens: tuple[str | None, ...]
+    website_urls: tuple[str | None, ...]
+    process_rule: Mapping[str, object]
+    dataset_id: str | None
+
+
+@dataclass
+class RecordingIndexingRunner:
+    calls: list[RunnerCall] = field(default_factory=list)
+    error: Exception | None = None
+
+    def indexing_estimate(
+        self,
+        tenant_id: str,
+        extract_settings: list[ExtractSetting],
+        tmp_processing_rule: Mapping[str, object],
+        doc_form: str | None = None,
+        doc_language: str = "English",
+        dataset_id: str | None = None,
+        indexing_technique: str = "economy",
+        *,
+        session: Session,
+    ) -> IndexingEstimate:
+        del doc_form, doc_language, indexing_technique, session
+        if self.error is not None:
+            raise self.error
+        self.calls.append(
+            RunnerCall(
+                tenant_id=tenant_id,
+                source_types=tuple(str(setting.datasource_type) for setting in extract_settings),
+                file_ids=tuple(setting.upload_file.id if setting.upload_file else None for setting in extract_settings),
+                notion_page_ids=tuple(
+                    setting.notion_info.notion_obj_id if setting.notion_info else None for setting in extract_settings
+                ),
+                notion_tokens=tuple(
+                    setting.notion_info.notion_access_token if setting.notion_info else None
+                    for setting in extract_settings
+                ),
+                website_urls=tuple(
+                    setting.website_info.url if setting.website_info else None for setting in extract_settings
+                ),
+                process_rule=tmp_processing_rule,
+                dataset_id=dataset_id,
+            )
+        )
+        return IndexingEstimate(total_segments=len(extract_settings), preview=[])
+
+
+@dataclass
+class StaticCredentialResolver:
+    credentials: Mapping[str, object] = field(default_factory=lambda: {"integration_secret": "secret"})
+    calls: list[tuple[str, str, str, str, str]] = field(default_factory=list)
+
+    def resolve(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        credential_id: str,
+        provider: str,
+        plugin_id: str,
+    ) -> dict[str, object]:
+        self.calls.append((workspace_id, actor_id, credential_id, provider, plugin_id))
+        return dict(self.credentials)
+
+
+def test_notion_access_token_is_excluded_from_model_serialization_and_repr() -> None:
+    info = NotionInfo(
+        credential_id="credential-1",
+        notion_workspace_id="workspace",
+        notion_obj_id="page-1",
+        notion_page_type="page",
+        tenant_id="workspace-1",
+        notion_access_token="secret",
+    )
+
+    assert "notion_access_token" not in info.model_dump()
+    assert "secret" not in repr(info)
+
+
+def _upload(file_id: str) -> UploadFile:
+    upload = UploadFile(
+        tenant_id="workspace-1",
+        storage_type=StorageType.LOCAL,
+        key=f"{file_id}.txt",
+        name=f"{file_id}.txt",
+        size=10,
+        extension="txt",
+        mime_type="text/plain",
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by="account-1",
+        created_at=datetime(2026, 1, 1),
+        used=False,
+    )
+    upload.id = file_id
+    return upload
+
+
+def _dataset() -> DatasetRecord:
+    return DatasetRecord(
+        id="dataset-1",
+        workspace_id="workspace-1",
+        maintainer_id="account-1",
+        permission="only_me",
+        data_source_type="upload_file",
+        indexing_technique="high_quality",
+        embedding_model=None,
+        embedding_model_provider=None,
+    )
+
+
+def _document(
+    document_id: str,
+    source_type: str,
+    source_info: Mapping[str, object] | None,
+    *,
+    process_rule_id: str | None = "rule-1",
+) -> DocumentRecord:
+    return DocumentRecord(
+        id=document_id,
+        workspace_id="workspace-1",
+        dataset_id="dataset-1",
+        name=document_id,
+        data_source_type=source_type,
+        data_source_info=source_info,
+        enabled=True,
+        archived=False,
+        indexing_status="waiting",
+        batch="batch-1",
+        doc_form="text_model",
+        doc_language="English",
+        dataset_process_rule_id=process_rule_id,
+        need_summary=False,
+        doc_metadata=None,
+    )
+
+
+def test_gateway_maps_all_new_source_variants(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_upload("file-1"))
+    runner = RecordingIndexingRunner()
+    credentials = StaticCredentialResolver()
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=credentials,
+        runner_factory=lambda: runner,
+    )
+
+    result = gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=NewSourcesEstimateCommand(
+            sources=(
+                UploadFileEstimateSource("file-1"),
+                NotionEstimateSource("notion-workspace", "page-1", "page", "credential-1"),
+                WebsiteEstimateSource("firecrawl", "job-1", "https://example.com"),
+            ),
+            process_rule={"mode": "automatic", "rules": {}},
+            dataset_id="dataset-1",
+        ),
+    )
+
+    assert result.total_segments == 3
+    assert runner.calls == [
+        RunnerCall(
+            tenant_id="workspace-1",
+            source_types=("upload_file", "notion_import", "website_crawl"),
+            file_ids=("file-1", None, None),
+            notion_page_ids=(None, "page-1", None),
+            notion_tokens=(None, "secret", None),
+            website_urls=(None, None, "https://example.com"),
+            process_rule={"mode": "automatic", "rules": {}},
+            dataset_id="dataset-1",
+        )
+    ]
+    assert credentials.calls == [
+        (
+            "workspace-1",
+            "account-1",
+            "credential-1",
+            "notion_datasource",
+            "langgenius/notion_datasource",
+        )
+    ]
+
+
+def test_gateway_maps_entire_existing_document_batch_in_one_runner_call(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_upload("file-1"))
+        rule = DatasetProcessRule(
+            dataset_id="dataset-1",
+            mode=ProcessRuleMode.AUTOMATIC,
+            rules=None,
+            created_by="account-1",
+        )
+        rule.id = "rule-1"
+        session.add(rule)
+    runner = RecordingIndexingRunner()
+    credentials = StaticCredentialResolver()
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=credentials,
+        runner_factory=lambda: runner,
+    )
+
+    result = gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=ExistingDocumentsEstimateCommand(
+            dataset=_dataset(),
+            documents=(
+                _document("document-1", "upload_file", {"upload_file_id": "file-1"}),
+                _document(
+                    "document-2",
+                    "notion_import",
+                    {
+                        "credential_id": "credential-1",
+                        "notion_workspace_id": "notion-workspace",
+                        "notion_page_id": "page-1",
+                        "type": "page",
+                    },
+                ),
+                _document(
+                    "document-3",
+                    "website_crawl",
+                    {
+                        "provider": "firecrawl",
+                        "job_id": "job-1",
+                        "url": "https://example.com",
+                        "mode": "crawl",
+                        "only_main_content": True,
+                    },
+                ),
+            ),
+        ),
+    )
+
+    assert result.total_segments == 3
+    assert len(runner.calls) == 1
+    assert runner.calls[0].source_types == ("upload_file", "notion_import", "website_crawl")
+
+
+def test_gateway_resolves_duplicate_notion_credential_once(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    runner = RecordingIndexingRunner()
+    credentials = StaticCredentialResolver()
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=credentials,
+        runner_factory=lambda: runner,
+    )
+
+    gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=NewSourcesEstimateCommand(
+            sources=(
+                NotionEstimateSource("notion-workspace", "page-1", "page", "credential-1"),
+                NotionEstimateSource("notion-workspace", "page-2", "page", "credential-1"),
+            ),
+            process_rule={"mode": "automatic", "rules": {}},
+        ),
+    )
+
+    assert len(credentials.calls) == 1
+    assert runner.calls[0].notion_page_ids == ("page-1", "page-2")
+
+
+def test_gateway_returns_zero_for_empty_sources_without_creating_runner(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    runner_created = False
+
+    def create_runner() -> RecordingIndexingRunner:
+        nonlocal runner_created
+        runner_created = True
+        return RecordingIndexingRunner()
+
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=create_runner,
+    )
+
+    result = gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=NewSourcesEstimateCommand(sources=(), process_rule={"mode": "automatic", "rules": {}}),
+    )
+
+    assert result == IndexingEstimate(total_segments=0, preview=[])
+    assert runner_created is False
+
+
+def test_gateway_rejects_upload_file_outside_workspace(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    with sqlite_session_factory.begin() as session:
+        session.add(_upload("file-1"))
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=RecordingIndexingRunner,
+    )
+
+    with pytest.raises(EstimateSourceNotFoundError, match="file-1"):
+        gateway.estimate(
+            workspace_id="workspace-2",
+            actor_id="account-1",
+            command=NewSourcesEstimateCommand(
+                sources=(UploadFileEstimateSource("file-1"),),
+                process_rule={"mode": "automatic", "rules": {}},
+            ),
+        )
+
+
+def test_gateway_rejects_existing_upload_file_outside_workspace(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    foreign_upload = _upload("file-1")
+    foreign_upload.tenant_id = "workspace-2"
+    with sqlite_session_factory.begin() as session:
+        session.add(foreign_upload)
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=RecordingIndexingRunner,
+    )
+
+    with pytest.raises(EstimateSourceNotFoundError, match="file-1"):
+        gateway.estimate(
+            workspace_id="workspace-1",
+            actor_id="account-1",
+            command=ExistingDocumentsEstimateCommand(
+                dataset=_dataset(),
+                documents=(_document("document-1", "upload_file", {"upload_file_id": "file-1"}),),
+            ),
+        )
+
+
+def test_gateway_uses_automatic_rule_when_existing_documents_have_no_saved_rule(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    runner = RecordingIndexingRunner()
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=lambda: runner,
+    )
+
+    result = gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=ExistingDocumentsEstimateCommand(
+            dataset=_dataset(),
+            documents=(
+                _document(
+                    "document-1",
+                    "website_crawl",
+                    {
+                        "provider": "firecrawl",
+                        "job_id": "job-1",
+                        "url": "https://example.com",
+                        "mode": "crawl",
+                    },
+                    process_rule_id=None,
+                ),
+            ),
+        ),
+    )
+
+    assert result.total_segments == 1
+    assert runner.calls[0].process_rule == {"mode": "automatic", "rules": {}}
+
+
+def test_gateway_skips_existing_document_without_source_info(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    runner_created = False
+
+    def create_runner() -> RecordingIndexingRunner:
+        nonlocal runner_created
+        runner_created = True
+        return RecordingIndexingRunner()
+
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=create_runner,
+    )
+
+    result = gateway.estimate(
+        workspace_id="workspace-1",
+        actor_id="account-1",
+        command=ExistingDocumentsEstimateCommand(
+            dataset=_dataset(),
+            documents=(_document("document-1", "upload_file", None, process_rule_id=None),),
+        ),
+    )
+
+    assert result == IndexingEstimate(total_segments=0, preview=[])
+    assert runner_created is False
+
+
+def test_gateway_rejects_notion_credential_without_secret(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(credentials={}),
+        runner_factory=RecordingIndexingRunner,
+    )
+
+    with pytest.raises(EstimateSourceNotFoundError, match="credential-1"):
+        gateway.estimate(
+            workspace_id="workspace-1",
+            actor_id="account-1",
+            command=NewSourcesEstimateCommand(
+                sources=(NotionEstimateSource("notion-workspace", "page-1", "page", "credential-1"),),
+                process_rule={"mode": "automatic", "rules": {}},
+            ),
+        )
+
+
+@pytest.mark.parametrize(
+    ("document", "expected_message"),
+    [
+        (
+            _document(
+                "document-1",
+                "notion_import",
+                {
+                    "credential_id": "credential-1",
+                    "notion_page_id": "page-1",
+                    "type": "page",
+                },
+            ),
+            "notion_workspace_id",
+        ),
+        (_document("document-1", "unsupported", {"source": "value"}), "unsupported"),
+    ],
+)
+def test_gateway_rejects_invalid_existing_document_source(
+    sqlite_session_factory: sessionmaker[Session],
+    document: DocumentRecord,
+    expected_message: str,
+) -> None:
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=RecordingIndexingRunner,
+    )
+    expected_error = (
+        EstimateSourceNotFoundError if document.data_source_type == "notion_import" else UnsupportedEstimateSourceError
+    )
+
+    with pytest.raises(expected_error, match=expected_message):
+        gateway.estimate(
+            workspace_id="workspace-1",
+            actor_id="account-1",
+            command=ExistingDocumentsEstimateCommand(dataset=_dataset(), documents=(document,)),
+        )
+
+
+def test_gateway_rejects_missing_existing_process_rule(
+    sqlite_session_factory: sessionmaker[Session],
+) -> None:
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=RecordingIndexingRunner,
+    )
+
+    with pytest.raises(EstimateSourceNotFoundError, match="rule-1"):
+        gateway.estimate(
+            workspace_id="workspace-1",
+            actor_id="account-1",
+            command=ExistingDocumentsEstimateCommand(
+                dataset=_dataset(),
+                documents=(
+                    _document(
+                        "document-1",
+                        "website_crawl",
+                        {
+                            "provider": "firecrawl",
+                            "job_id": "job-1",
+                            "url": "https://example.com",
+                            "mode": "crawl",
+                        },
+                    ),
+                ),
+            ),
+        )
+
+
+@pytest.mark.parametrize(
+    ("runner_error", "expected_error"),
+    [
+        (LLMBadRequestError("provider unavailable"), IndexingEstimateProviderUnavailableError),
+        (ProviderTokenNotInitError("provider unavailable"), IndexingEstimateProviderUnavailableError),
+        (PluginDaemonClientSideError("provider unavailable"), IndexingEstimateProviderUnavailableError),
+        (RuntimeError("runner failed"), IndexingEstimateExecutionError),
+    ],
+)
+def test_gateway_translates_runner_errors(
+    sqlite_session_factory: sessionmaker[Session],
+    runner_error: Exception,
+    expected_error: type[Exception],
+) -> None:
+    runner = RecordingIndexingRunner(error=runner_error)
+    gateway = IndexingRunnerEstimateGateway(
+        session_factory=sqlite_session_factory,
+        credentials=StaticCredentialResolver(),
+        runner_factory=lambda: runner,
+    )
+
+    with pytest.raises(expected_error, match="provider unavailable|runner failed"):
+        gateway.estimate(
+            workspace_id="workspace-1",
+            actor_id="account-1",
+            command=NewSourcesEstimateCommand(
+                sources=(WebsiteEstimateSource("firecrawl", "job-1", "https://example.com"),),
+                process_rule={"mode": "automatic", "rules": {}},
+            ),
+        )

api/tests/unit_tests/services/knowledge/test_indexing_estimate_application_service.py

@@ -0,0 +1,230 @@
+from dataclasses import dataclass, field
+
+import pytest
+
+from core.entities.knowledge_entities import IndexingEstimate
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef
+from machinery.context import RequestContext
+from services.entities.knowledge_entities.indexing_estimate import (
+    EstimateCommand,
+    ExistingDocumentsEstimateCommand,
+    NewEstimateSource,
+    NewSourcesEstimateCommand,
+    NotionEstimateSource,
+    UploadFileEstimateSource,
+    WebsiteEstimateSource,
+)
+from services.entities.knowledge_entities.records import DatasetRecord, DocumentRecord
+from services.knowledge.application import (
+    DatasetAccessDeniedError,
+    EstimateDocumentAlreadyFinishedError,
+    EstimateDocumentNotFoundError,
+    IndexingEstimateApplicationService,
+)
+
+
+def _context() -> RequestContext:
+    return RequestContext("request-1", None, "account-1", "workspace-1")
+
+
+def _dataset() -> DatasetRecord:
+    return DatasetRecord(
+        id="dataset-1",
+        workspace_id="workspace-1",
+        maintainer_id="account-1",
+        permission="only_me",
+        data_source_type="upload_file",
+        indexing_technique="high_quality",
+        embedding_model=None,
+        embedding_model_provider=None,
+    )
+
+
+def _document(document_id: str = "document-1", *, status: str = "waiting") -> DocumentRecord:
+    return DocumentRecord(
+        id=document_id,
+        workspace_id="workspace-1",
+        dataset_id="dataset-1",
+        name=document_id,
+        data_source_type="upload_file",
+        data_source_info={"upload_file_id": f"file-{document_id}"},
+        enabled=True,
+        archived=False,
+        indexing_status=status,
+        batch="batch-1",
+        doc_form="text_model",
+        doc_language="English",
+        dataset_process_rule_id="rule-1",
+        need_summary=False,
+        doc_metadata=None,
+    )
+
+
+@dataclass
+class DatasetAccessStub:
+    dataset: DatasetRecord = field(default_factory=_dataset)
+    error: Exception | None = None
+    calls: list[str] = field(default_factory=list)
+
+    def require_accessible(self, context: RequestContext, dataset_id: str) -> DatasetRecord:
+        assert context == _context()
+        self.calls.append(dataset_id)
+        if self.error is not None:
+            raise self.error
+        return self.dataset
+
+
+@dataclass
+class DocumentReaderStub:
+    document: DocumentRecord | None = field(default_factory=_document)
+    batch: tuple[DocumentRecord, ...] = ()
+
+    def get_by_ref(self, document_ref: DocumentRef) -> DocumentRecord | None:
+        assert document_ref.dataset == DatasetRef("workspace-1", "dataset-1")
+        return self.document
+
+    def list_by_batch(self, dataset_ref: DatasetRef, batch: str) -> tuple[DocumentRecord, ...]:
+        assert dataset_ref == DatasetRef("workspace-1", "dataset-1")
+        assert batch == "batch-1"
+        return self.batch
+
+
+@dataclass
+class GatewayRecorder:
+    calls: list[tuple[str, str, EstimateCommand]] = field(default_factory=list)
+
+    def estimate(
+        self,
+        *,
+        workspace_id: str,
+        actor_id: str,
+        command: EstimateCommand,
+    ) -> IndexingEstimate:
+        self.calls.append((workspace_id, actor_id, command))
+        return IndexingEstimate(total_segments=len(getattr(command, "sources", ())), preview=[])
+
+
+def _service(
+    *,
+    access: DatasetAccessStub | None = None,
+    documents: DocumentReaderStub | None = None,
+    gateway: GatewayRecorder | None = None,
+) -> tuple[IndexingEstimateApplicationService, GatewayRecorder]:
+    recorder = gateway or GatewayRecorder()
+    return (
+        IndexingEstimateApplicationService(
+            dataset_access=access or DatasetAccessStub(),
+            documents=documents or DocumentReaderStub(),
+            gateway=recorder,
+        ),
+        recorder,
+    )
+
+
+@pytest.mark.parametrize(
+    "source",
+    [
+        UploadFileEstimateSource("file-1"),
+        NotionEstimateSource("notion-workspace", "page-1", "page", "credential-1"),
+        WebsiteEstimateSource("firecrawl", "job-1", "https://example.com"),
+    ],
+)
+def test_estimate_new_sources_supports_each_explicit_variant(source: NewEstimateSource) -> None:
+    service, gateway = _service()
+    command = NewSourcesEstimateCommand(
+        sources=(source,),
+        process_rule={"mode": "automatic", "summary_index_setting": {"enable": None}},
+    )
+
+    service.estimate_new_sources(_context(), command)
+
+    normalized = gateway.calls[0][2]
+    assert isinstance(normalized, NewSourcesEstimateCommand)
+    assert normalized.sources == (source,)
+    assert normalized.process_rule == {"mode": "automatic", "rules": {}}
+
+
+def test_estimate_new_sources_validates_optional_dataset_before_gateway_call() -> None:
+    access = DatasetAccessStub()
+    service, gateway = _service(access=access)
+
+    service.estimate_new_sources(
+        _context(),
+        NewSourcesEstimateCommand(
+            sources=(UploadFileEstimateSource("file-1"),),
+            process_rule={"mode": "automatic"},
+            dataset_id="dataset-1",
+        ),
+    )
+
+    assert access.calls == ["dataset-1"]
+    command = gateway.calls[0][2]
+    assert isinstance(command, NewSourcesEstimateCommand)
+    assert command.dataset_id == "dataset-1"
+
+
+def test_estimate_new_sources_stops_when_dataset_access_is_denied() -> None:
+    access = DatasetAccessStub(error=DatasetAccessDeniedError())
+    service, gateway = _service(access=access)
+
+    with pytest.raises(DatasetAccessDeniedError):
+        service.estimate_new_sources(
+            _context(),
+            NewSourcesEstimateCommand(
+                sources=(UploadFileEstimateSource("file-1"),),
+                process_rule={"mode": "automatic"},
+                dataset_id="dataset-1",
+            ),
+        )
+
+    assert gateway.calls == []
+
+
+def test_estimate_batch_sends_all_documents_in_one_gateway_call() -> None:
+    documents = (_document("document-1"), _document("document-2"))
+    service, gateway = _service(documents=DocumentReaderStub(batch=documents))
+
+    service.estimate_batch(_context(), dataset_id="dataset-1", batch="batch-1")
+
+    assert len(gateway.calls) == 1
+    command = gateway.calls[0][2]
+    assert isinstance(command, ExistingDocumentsEstimateCommand)
+    assert command.documents == documents
+
+
+def test_estimate_empty_batch_returns_zero_without_gateway_call() -> None:
+    service, gateway = _service(documents=DocumentReaderStub(batch=()))
+
+    result = service.estimate_batch(_context(), dataset_id="dataset-1", batch="batch-1")
+
+    assert result == IndexingEstimate(total_segments=0, preview=[])
+    assert gateway.calls == []
+
+
+@pytest.mark.parametrize("status", ["completed", "error"])
+def test_estimate_rejects_finished_document(status: str) -> None:
+    service, gateway = _service(documents=DocumentReaderStub(document=_document(status=status)))
+
+    with pytest.raises(EstimateDocumentAlreadyFinishedError):
+        service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1")
+
+    assert gateway.calls == []
+
+
+def test_estimate_rejects_missing_document() -> None:
+    service, _ = _service(documents=DocumentReaderStub(document=None))
+
+    with pytest.raises(EstimateDocumentNotFoundError):
+        service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1")
+
+
+def test_estimate_single_document_calls_gateway_with_owned_document() -> None:
+    document = _document()
+    service, gateway = _service(documents=DocumentReaderStub(document=document))
+
+    service.estimate_document(_context(), dataset_id="dataset-1", document_id="document-1")
+
+    assert len(gateway.calls) == 1
+    command = gateway.calls[0][2]
+    assert isinstance(command, ExistingDocumentsEstimateCommand)
+    assert command.documents == (document,)

api/tests/unit_tests/services/knowledge/test_indexing_estimate_entities.py

@@ -0,0 +1,126 @@
+import pytest
+
+from services.entities.knowledge_entities.indexing_estimate import (
+    EstimateValidationError,
+    normalize_indexing_estimate_args,
+    normalize_process_rule,
+)
+
+
+def test_normalize_indexing_estimate_args_requires_info_list() -> None:
+    with pytest.raises(EstimateValidationError, match="Field required"):
+        normalize_indexing_estimate_args({})
+
+
+def test_normalize_process_rule_sets_empty_rules_for_automatic_mode() -> None:
+    normalized = normalize_process_rule({"mode": "automatic", "rules": {"ignored": True}})
+
+    assert normalized["rules"] == {}
+
+
+@pytest.mark.parametrize("summary_setting", [None, {"enable": None}])
+def test_normalize_process_rule_treats_absent_summary_enable_as_no_summary(summary_setting: object) -> None:
+    normalized = normalize_process_rule(
+        {
+            "mode": "automatic",
+            "summary_index_setting": summary_setting,
+        }
+    )
+
+    assert "summary_index_setting" not in normalized
+
+
+def test_normalize_process_rule_rejects_unknown_pre_processing_rule_id() -> None:
+    with pytest.raises(EstimateValidationError):
+        normalize_process_rule(
+            {
+                "mode": "custom",
+                "rules": {
+                    "pre_processing_rules": [{"id": "unknown", "enabled": True}],
+                    "segmentation": {"separator": "\n", "max_tokens": 128},
+                },
+            }
+        )
+
+
+def test_normalize_process_rule_deduplicates_custom_rules() -> None:
+    normalized = normalize_process_rule(
+        {
+            "mode": "custom",
+            "rules": {
+                "pre_processing_rules": [
+                    {"id": "remove_stopwords", "enabled": True},
+                    {"id": "remove_stopwords", "enabled": False},
+                ],
+                "segmentation": {"separator": "\n", "max_tokens": 128},
+            },
+        }
+    )
+
+    assert normalized["rules"]["pre_processing_rules"] == [{"id": "remove_stopwords", "enabled": False}]
+
+
+def test_normalize_process_rule_drops_hierarchical_fields_from_custom_mode() -> None:
+    normalized = normalize_process_rule(
+        {
+            "mode": "custom",
+            "rules": {
+                "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
+                "segmentation": {"separator": "\n", "max_tokens": 128},
+                "parent_mode": "full-doc",
+                "subchunk_segmentation": {"separator": "###", "max_tokens": 64},
+            },
+        }
+    )
+
+    assert normalized["rules"] == {
+        "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
+        "segmentation": {"separator": "\n", "max_tokens": 128},
+    }
+
+
+def test_normalize_process_rule_requires_summary_index_provider_name() -> None:
+    with pytest.raises(EstimateValidationError, match="Field required"):
+        normalize_process_rule(
+            {
+                "mode": "custom",
+                "rules": {
+                    "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
+                    "segmentation": {"separator": "\n", "max_tokens": 128},
+                },
+                "summary_index_setting": {"enable": True, "model_name": "summary-model"},
+            }
+        )
+
+
+def test_normalize_process_rule_preserves_hierarchical_fields() -> None:
+    normalized = normalize_process_rule(
+        {
+            "mode": "hierarchical",
+            "rules": {
+                "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
+                "segmentation": {"separator": "\n", "max_tokens": 512},
+                "parent_mode": "full-doc",
+                "subchunk_segmentation": {"separator": "###", "max_tokens": 128},
+            },
+        }
+    )
+
+    assert normalized["rules"]["parent_mode"] == "full-doc"
+    assert normalized["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128}
+
+
+def test_normalize_process_rule_defaults_hierarchical_parent_mode() -> None:
+    normalized = normalize_process_rule(
+        {
+            "mode": "hierarchical",
+            "rules": {
+                "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
+                "segmentation": {"separator": "\n", "max_tokens": 512},
+                "subchunk_segmentation": {"separator": "###", "max_tokens": 128},
+            },
+        }
+    )
+
+    assert normalized["rules"]["parent_mode"] == "paragraph"
+    assert normalized["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128}

api/tests/unit_tests/services/test_credential_permission_service.py

@@ -132,6 +132,7 @@ class TestApplyVisibilityFilter:
 
         query = CredentialPermissionService.apply_visibility_filter(
             select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id),
+            tenant_id=tenant_id,
             model_id_column=TriggerSubscription.id,
             model_user_id_column=TriggerSubscription.user_id,
             model_visibility_column=TriggerSubscription.visibility,
@@ -190,6 +191,7 @@ class TestApplyVisibilityFilter:
 
         query = CredentialPermissionService.apply_visibility_filter(
             select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id),
+            tenant_id=tenant_id,
             model_id_column=TriggerSubscription.id,
             model_user_id_column=TriggerSubscription.user_id,
             model_visibility_column=TriggerSubscription.visibility,
@@ -199,3 +201,41 @@ class TestApplyVisibilityFilter:
 
         visible_ids = {subscription.id for subscription in sqlite_session.scalars(query)}
         assert visible_ids == {team_subscription.id, owned_subscription.id, shared_subscription.id}
+
+    def test_actor_entry_point_scopes_partial_permissions_by_tenant(
+        self,
+        sqlite_session: Session,
+        tenant_id: str,
+        user_id: str,
+        other_user_id: str,
+    ) -> None:
+        other_tenant_id = str(uuid4())
+        shared_subscription = _subscription(
+            tenant_id=tenant_id,
+            owner_id=other_user_id,
+            name="shared",
+            visibility=PermissionEnum.PARTIAL_TEAM,
+        )
+        sqlite_session.add(shared_subscription)
+        sqlite_session.flush()
+        sqlite_session.add(
+            CredentialPermission(
+                credential_id=shared_subscription.id,
+                credential_type=CredentialType.TRIGGER_SUBSCRIPTION,
+                account_id=user_id,
+                tenant_id=other_tenant_id,
+            )
+        )
+        sqlite_session.commit()
+
+        query = CredentialPermissionService.apply_visibility_filter_for_actor(
+            select(TriggerSubscription).where(TriggerSubscription.tenant_id == tenant_id),
+            tenant_id=tenant_id,
+            model_id_column=TriggerSubscription.id,
+            model_user_id_column=TriggerSubscription.user_id,
+            model_visibility_column=TriggerSubscription.visibility,
+            credential_type=CredentialType.TRIGGER_SUBSCRIPTION,
+            actor_id=user_id,
+        )
+
+        assert sqlite_session.scalars(query).all() == []

api/tests/unit_tests/services/test_data_source_oauth_service.py

@@ -1,182 +0,0 @@
-from unittest.mock import MagicMock, call
-
-import pytest
-
-from machinery.context import RequestContext
-from services.data_source_oauth_service import (
-    DataSourceOAuthConfigurationError,
-    DataSourceOAuthError,
-    DataSourceOAuthService,
-    InvalidDataSourceOAuthCodeError,
-)
-from services.entities.data_source_oauth_entities import (
-    DataSourceOAuthAuthorization,
-    DataSourceOAuthBindingRecord,
-    DataSourceOAuthCallback,
-)
-
-
-@pytest.fixture
-def request_context() -> RequestContext:
-    return RequestContext(
-        request_id="request-1",
-        trace_id="trace-1",
-        account_id="account-1",
-        active_workspace_id="workspace-1",
-    )
-
-
-def test_start_authorization_delegates_to_external_provider(request_context: RequestContext) -> None:
-    provider = MagicMock()
-    provider.get_authorization_url.return_value = "https://notion.example/authorize"
-    bindings = MagicMock()
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=provider,
-        bindings=bindings,
-    )
-
-    result = service.start_authorization(request_context)
-
-    assert result == "https://notion.example/authorize"
-    provider.get_authorization_url.assert_called_once_with()
-    bindings.upsert_authorization.assert_not_called()
-
-
-def test_start_authorization_persists_internal_authorization(request_context: RequestContext) -> None:
-    authorization = DataSourceOAuthAuthorization(access_token="secret", source_info={"pages": []})
-    provider = MagicMock()
-    provider.authorize_internal.return_value = authorization
-    bindings = MagicMock()
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=provider,
-        bindings=bindings,
-        is_internal_provider=True,
-        internal_access_token="secret",
-    )
-
-    result = service.start_authorization(request_context)
-
-    assert result == "internal"
-    provider.authorize_internal.assert_called_once_with("secret", "workspace-1")
-    bindings.upsert_authorization.assert_called_once_with(
-        workspace_id="workspace-1",
-        provider="notion",
-        authorization=authorization,
-    )
-
-
-def test_start_authorization_rejects_missing_internal_secret(request_context: RequestContext) -> None:
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=MagicMock(),
-        bindings=MagicMock(),
-        is_internal_provider=True,
-    )
-
-    with pytest.raises(DataSourceOAuthConfigurationError, match="Internal secret is not set"):
-        service.start_authorization(request_context)
-
-
-@pytest.mark.parametrize(
-    ("code", "error", "expected_error"),
-    [
-        ("auth/code", None, None),
-        (None, "access denied", "access denied"),
-        (None, None, "Access denied"),
-    ],
-)
-def test_complete_callback_returns_framework_neutral_result(
-    code: str | None,
-    error: str | None,
-    expected_error: str | None,
-) -> None:
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=MagicMock(),
-        bindings=MagicMock(),
-    )
-
-    assert service.complete_callback(code=code, error=error) == DataSourceOAuthCallback(
-        provider="notion",
-        code=code,
-        error=expected_error,
-    )
-
-
-def test_bind_authorizes_before_persisting(request_context: RequestContext) -> None:
-    events = MagicMock()
-    authorization = DataSourceOAuthAuthorization(access_token="token", source_info={"pages": []})
-    provider = MagicMock()
-    provider.authorize.side_effect = lambda code: (events.authorize(code), authorization)[1]
-    bindings = MagicMock()
-    bindings.upsert_authorization.side_effect = lambda **kwargs: events.persist(**kwargs)
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=provider,
-        bindings=bindings,
-    )
-
-    service.bind(request_context, code="code-1")
-
-    assert events.mock_calls == [
-        call.authorize("code-1"),
-        call.persist(workspace_id="workspace-1", provider="notion", authorization=authorization),
-    ]
-
-
-def test_bind_rejects_empty_code(request_context: RequestContext) -> None:
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=MagicMock(),
-        bindings=MagicMock(),
-    )
-
-    with pytest.raises(InvalidDataSourceOAuthCodeError):
-        service.bind(request_context, code="")
-
-
-def test_sync_reads_before_provider_io_and_persists_afterward(request_context: RequestContext) -> None:
-    events = MagicMock()
-    binding = DataSourceOAuthBindingRecord(id="binding-1", access_token="token", source_info={"pages": []})
-    refreshed_source_info = {"pages": [{"page_id": "page-1"}]}
-    bindings = MagicMock()
-    bindings.get_enabled.side_effect = lambda **kwargs: (events.read(**kwargs), binding)[1]
-    bindings.update_source_info.side_effect = lambda **kwargs: (events.update(**kwargs), True)[1]
-    provider = MagicMock()
-    provider.refresh.side_effect = lambda *args: (events.refresh(*args), refreshed_source_info)[1]
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=provider,
-        bindings=bindings,
-    )
-
-    service.sync(request_context, binding_id="binding-1")
-
-    assert events.mock_calls == [
-        call.read(workspace_id="workspace-1", provider="notion", binding_id="binding-1"),
-        call.refresh("token", {"pages": []}),
-        call.update(
-            workspace_id="workspace-1",
-            provider="notion",
-            binding_id="binding-1",
-            source_info=refreshed_source_info,
-        ),
-    ]
-
-
-def test_sync_rejects_missing_binding_without_provider_io(request_context: RequestContext) -> None:
-    provider = MagicMock()
-    bindings = MagicMock()
-    bindings.get_enabled.return_value = None
-    service = DataSourceOAuthService(
-        provider_name="notion",
-        provider_gateway=provider,
-        bindings=bindings,
-    )
-
-    with pytest.raises(DataSourceOAuthError, match="Data source binding not found"):
-        service.sync(request_context, binding_id="missing")
-
-    provider.refresh.assert_not_called()

api/tests/unit_tests/services/test_dataset_service_document.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from sqlalchemy import event, select
 from sqlalchemy.orm import Session
 
+from core.model_manager import ModelInstance
 from models.account import Tenant
 from models.dataset import Dataset, DatasetCollectionBinding, DocumentSegment
 from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
@@ -39,7 +40,6 @@ from .dataset_service_test_helpers import (
     RetrievalModel,
     Rule,
     Segmentation,
-    SimpleNamespace,
     WebsiteInfo,
     _make_features,
     _make_lock_context,
@@ -933,13 +933,22 @@ class TestDocumentServiceCreateValidation:
     """Unit tests for document creation validation helpers."""
 
     def test_document_create_args_validate_requires_data_source_or_process_rule(self):
-        knowledge_config = SimpleNamespace(data_source=None, process_rule=None)
+        knowledge_config = KnowledgeConfig(indexing_technique="economy")
 
         with pytest.raises(ValueError, match="Data source or Process rule is required"):
             DocumentService.document_create_args_validate(knowledge_config)
 
     def test_document_create_args_validate_delegates_to_sub_validators(self):
-        knowledge_config = SimpleNamespace(data_source=object(), process_rule=object())
+        knowledge_config = KnowledgeConfig(
+            indexing_technique="economy",
+            data_source=DataSource(
+                info_list=InfoList(
+                    data_source_type="upload_file",
+                    file_info_list=FileInfo(file_ids=["file-1"]),
+                )
+            ),
+            process_rule=ProcessRule(mode="automatic"),
+        )
 
         with (
             patch.object(DocumentService, "data_source_args_validate") as validate_data_source,
@@ -950,39 +959,20 @@ class TestDocumentServiceCreateValidation:
         validate_data_source.assert_called_once_with(knowledge_config)
         validate_process_rule.assert_called_once_with(knowledge_config)
 
-    def test_data_source_args_validate_rejects_invalid_type(self):
-        knowledge_config = SimpleNamespace(
-            data_source=SimpleNamespace(
-                info_list=SimpleNamespace(
-                    data_source_type="bad-source",
-                    file_info_list=None,
-                    notion_info_list=None,
-                    website_info_list=None,
-                )
-            )
-        )
-
-        with pytest.raises(ValueError, match="Data source type is invalid"):
-            DocumentService.data_source_args_validate(knowledge_config)
-
     @pytest.mark.parametrize(
-        ("data_source_type", "field_name", "message"),
+        ("data_source_type", "message"),
         [
-            ("upload_file", "file_info_list", "File source info is required"),
-            ("notion_import", "notion_info_list", "Notion source info is required"),
-            ("website_crawl", "website_info_list", "Website source info is required"),
+            ("upload_file", "File source info is required"),
+            ("notion_import", "Notion source info is required"),
+            ("website_crawl", "Website source info is required"),
         ],
     )
-    def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, field_name, message):
-        info_values = {
-            "data_source_type": data_source_type,
-            "file_info_list": object(),
-            "notion_info_list": object(),
-            "website_info_list": object(),
-        }
-        info_values[field_name] = None
-        info_list = SimpleNamespace(**info_values)
-        knowledge_config = SimpleNamespace(data_source=SimpleNamespace(info_list=info_list))
+    def test_data_source_args_validate_requires_source_specific_info(self, data_source_type, message):
+        info_list = InfoList.model_validate({"data_source_type": data_source_type})
+        knowledge_config = KnowledgeConfig(
+            indexing_technique="economy",
+            data_source=DataSource(info_list=info_list),
+        )
 
         with pytest.raises(ValueError, match=message):
             DocumentService.data_source_args_validate(knowledge_config)
@@ -1203,44 +1193,6 @@ class TestDocumentServiceSaveDocumentWithDatasetId:
                     session=sqlite_session,
                 )
 
-    def test_save_document_with_dataset_id_rejects_invalid_indexing_technique(
-        self, account_context, unbound_session: Session
-    ):
-        dataset = _dataset_row(indexing_technique=None)
-        knowledge_config = SimpleNamespace(
-            doc_form=IndexStructureType.PARAGRAPH_INDEX,
-            original_document_id=None,
-            data_source=None,
-            indexing_technique="broken-technique",
-        )
-
-        with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)):
-            with pytest.raises(ValueError, match="Indexing technique is invalid"):
-                DocumentService.save_document_with_dataset_id(
-                    dataset,
-                    knowledge_config,
-                    account_context,
-                    session=unbound_session,
-                )
-
-    def test_save_document_with_dataset_id_returns_empty_for_invalid_process_rule_mode(
-        self, account_context, unbound_session: Session
-    ):
-        dataset = _dataset_row()
-        knowledge_config = _make_upload_knowledge_config(file_ids=["file-1"])
-        knowledge_config.process_rule = SimpleNamespace(mode="unsupported-mode", rules=None)
-
-        with patch("services.dataset_service.FeatureService.get_features", return_value=_make_features(enabled=False)):
-            documents, batch = DocumentService.save_document_with_dataset_id(
-                dataset,
-                knowledge_config,
-                account_context,
-                session=unbound_session,
-            )
-
-        assert documents == []
-        assert batch == ""
-
     def test_save_document_with_dataset_id_upload_file_creates_and_reindexes_documents(
         self, account_context, sqlite_session: Session
     ):
@@ -1798,13 +1750,9 @@ class TestDocumentServiceSaveWithoutDatasetBilling:
 
 
 class TestDocumentServiceEstimateValidation:
-    """Unit tests for estimate_args_validate branches."""
+    """Keep only the compatibility contract owned by the legacy service."""
 
-    def test_estimate_args_validate_rejects_missing_info_list(self):
-        with pytest.raises(ValueError, match="Field required"):
-            DocumentService.estimate_args_validate({})
-
-    def test_estimate_args_validate_sets_empty_rules_for_automatic_mode(self):
+    def test_estimate_args_validate_delegates_to_shared_normalizer(self):
         args = {
             "info_list": {"data_source_type": "upload_file"},
             "process_rule": {"mode": "automatic", "rules": {"ignored": True}},
@@ -1814,114 +1762,6 @@ class TestDocumentServiceEstimateValidation:
 
         assert args["process_rule"]["rules"] == {}
 
-    def test_estimate_args_validate_rejects_unknown_pre_processing_rule_id(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "custom",
-                "rules": {
-                    "pre_processing_rules": [{"id": "unknown", "enabled": True}],
-                    "segmentation": {"separator": "\n", "max_tokens": 128},
-                },
-            },
-        }
-
-        with pytest.raises(ValueError, match="pre_processing_rules id is invalid"):
-            DocumentService.estimate_args_validate(args)
-
-    def test_estimate_args_validate_deduplicates_rules_for_custom_mode(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "custom",
-                "rules": {
-                    "pre_processing_rules": [
-                        {"id": "remove_stopwords", "enabled": True},
-                        {"id": "remove_stopwords", "enabled": False},
-                    ],
-                    "segmentation": {"separator": "\n", "max_tokens": 128},
-                },
-            },
-        }
-
-        DocumentService.estimate_args_validate(args)
-
-        assert args["process_rule"]["rules"]["pre_processing_rules"] == [{"id": "remove_stopwords", "enabled": False}]
-
-    def test_estimate_args_validate_custom_mode_drops_hierarchical_fields(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "custom",
-                "rules": {
-                    "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
-                    "segmentation": {"separator": "\n", "max_tokens": 128},
-                    "parent_mode": "full-doc",
-                    "subchunk_segmentation": {"separator": "###", "max_tokens": 64},
-                },
-            },
-        }
-
-        DocumentService.estimate_args_validate(args)
-
-        assert args["process_rule"]["rules"] == {
-            "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
-            "segmentation": {"separator": "\n", "max_tokens": 128},
-        }
-
-    def test_estimate_args_validate_requires_summary_index_provider_name(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "custom",
-                "rules": {
-                    "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
-                    "segmentation": {"separator": "\n", "max_tokens": 128},
-                },
-                "summary_index_setting": {"enable": True, "model_name": "summary-model"},
-            },
-        }
-
-        with pytest.raises(ValueError, match="Field required"):
-            DocumentService.estimate_args_validate(args)
-
-    def test_estimate_args_validate_preserves_hierarchical_fields(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "hierarchical",
-                "rules": {
-                    "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
-                    "segmentation": {"separator": "\n", "max_tokens": 512},
-                    "parent_mode": "full-doc",
-                    "subchunk_segmentation": {"separator": "###", "max_tokens": 128},
-                },
-            },
-        }
-
-        DocumentService.estimate_args_validate(args)
-
-        assert args["process_rule"]["rules"]["parent_mode"] == "full-doc"
-        assert args["process_rule"]["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128}
-
-    def test_estimate_args_validate_hierarchical_defaults_parent_mode_to_paragraph(self):
-        args = {
-            "info_list": {"data_source_type": "upload_file"},
-            "process_rule": {
-                "mode": "hierarchical",
-                "rules": {
-                    "pre_processing_rules": [{"id": "remove_stopwords", "enabled": True}],
-                    "segmentation": {"separator": "\n", "max_tokens": 512},
-                    "subchunk_segmentation": {"separator": "###", "max_tokens": 128},
-                },
-            },
-        }
-
-        DocumentService.estimate_args_validate(args)
-
-        assert args["process_rule"]["rules"]["parent_mode"] == "paragraph"
-        assert args["process_rule"]["rules"]["subchunk_segmentation"] == {"separator": "###", "max_tokens": 128}
-
 
 class TestDocumentServiceSaveDocumentAdditionalBranches:
     """Additional unit tests for dataset bootstrap and process-rule branches."""
@@ -1963,10 +1803,10 @@ class TestDocumentServiceSaveDocumentAdditionalBranches:
             ) as get_binding,
             patch.object(DocumentService, "update_document_with_dataset_id", return_value=updated_document),
         ):
-            model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = SimpleNamespace(
-                model_name="default-embedding",
-                provider="default-provider",
-            )
+            default_model = object.__new__(ModelInstance)
+            default_model.model_name = "default-embedding"
+            default_model.provider = "default-provider"
+            model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = default_model
 
             documents, batch = DocumentService.save_document_with_dataset_id(
                 dataset,

api/tests/unit_tests/services/test_dataset_service_segment.py

@@ -2,7 +2,15 @@
 
 from collections.abc import Callable
 
-from services.dataset_ref_service import DatasetRef, DatasetRefService, DocumentRef, SegmentRef
+from flask_sqlalchemy.pagination import SelectPagination
+from sqlalchemy import literal, select
+from sqlalchemy.orm import Session
+
+from core.rag.entities.dataset_reference import DatasetRef, DocumentRef, SegmentRef
+from models.account import Tenant
+from models.dataset import DatasetProcessRule, DocumentSegmentSummary
+from models.enums import ProcessRuleMode
+from services.dataset_ref_service import DatasetRefService
 
 from .dataset_service_test_helpers import (
     Account,
@@ -16,7 +24,6 @@ from .dataset_service_test_helpers import (
     ModelType,
     SegmentService,
     SegmentUpdateArgs,
-    SimpleNamespace,
     _make_child_chunk,
     _make_dataset,
     _make_document,
@@ -28,6 +35,35 @@ from .dataset_service_test_helpers import (
 )
 
 
+def _account() -> Account:
+    account = Account(name="User", email="user-1@example.com")
+    account.id = "user-1"
+    tenant = Tenant(name="Tenant")
+    tenant.id = "tenant-1"
+    account._current_tenant = tenant
+    return account
+
+
+def _process_rule(rule_id: str) -> DatasetProcessRule:
+    process_rule = DatasetProcessRule(
+        dataset_id="dataset-1",
+        mode=ProcessRuleMode.AUTOMATIC,
+        rules=None,
+        created_by="user-1",
+    )
+    process_rule.id = rule_id
+    return process_rule
+
+
+def _summary(content: str) -> DocumentSegmentSummary:
+    return DocumentSegmentSummary(
+        dataset_id="dataset-1",
+        document_id="doc-1",
+        chunk_id="segment-1",
+        summary_content=content,
+    )
+
+
 def _make_segment_ref(segment_id: str = "segment-1"):
     dataset = _make_dataset()
     document = _make_document(dataset_id=dataset.id, tenant_id=dataset.tenant_id)
@@ -40,13 +76,11 @@ def _make_segment_ref(segment_id: str = "segment-1"):
 class TestDatasetRefService:
     """Unit tests for typed dataset resource refs."""
 
-    def test_dataset_ref_is_plain_named_tuple(self):
+    def test_create_dataset_ref_carries_workspace_owner(self):
         dataset_ref = DatasetRefService.create_dataset_ref(_make_dataset())
 
-        assert isinstance(dataset_ref, DatasetRef)
         assert dataset_ref.tenant_id == "tenant-1"
         assert dataset_ref.dataset_id == "dataset-1"
-        assert tuple(dataset_ref) == ("tenant-1", "dataset-1")
 
     @pytest.mark.parametrize(
         ("document_dataset_id", "document_tenant_id"),
@@ -105,7 +139,7 @@ class TestSegmentServiceChildChunks:
 
     def test_create_child_chunk_assigns_next_position_and_commits(self, account_context):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         document = _make_document()
         segment = _make_segment()
 
@@ -137,7 +171,7 @@ class TestSegmentServiceChildChunks:
 
     def test_create_child_chunk_rolls_back_and_raises_on_vector_failure(self, account_context):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         document = _make_document()
         segment = _make_segment()
 
@@ -159,7 +193,7 @@ class TestSegmentServiceChildChunks:
 
     def test_update_child_chunks_updates_deletes_and_creates_records(self, account_context):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         document = _make_document()
         segment = _make_segment()
         existing_a = ChildChunk(
@@ -220,7 +254,7 @@ class TestSegmentServiceChildChunks:
 
     def test_update_child_chunks_rolls_back_on_vector_failure(self, account_context):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         document = _make_document()
         segment = _make_segment()
         existing_chunk = _make_child_chunk()
@@ -245,11 +279,11 @@ class TestSegmentServiceChildChunks:
 
     def test_update_child_chunk_updates_vector_and_commits(self, account_context):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         child_chunk = _make_child_chunk()
 
         with (
-            patch("services.dataset_service.current_user", SimpleNamespace(id="user-1")),
+            patch("services.dataset_service.current_user", _account()),
             patch("services.dataset_service.naive_utc_now", return_value="now"),
             patch("services.dataset_service.VectorService") as vector_service,
         ):
@@ -270,7 +304,7 @@ class TestSegmentServiceChildChunks:
 
     def test_delete_child_chunk_raises_delete_index_error_on_vector_failure(self):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         child_chunk = _make_child_chunk()
 
         with (
@@ -297,9 +331,18 @@ class TestSegmentServiceQueries:
         with patch("services.dataset_service.current_user", account):
             yield account
 
-    def test_get_child_chunks_applies_keyword_filter_and_paginate(self, account_context):
+    def test_get_child_chunks_applies_keyword_filter_and_paginate(
+        self,
+        account_context,
+        sqlite_session: Session,
+    ):
         session = MagicMock()
-        paginated = SimpleNamespace(items=["chunk"], total=1)
+        paginated = SelectPagination(
+            select=select(literal("chunk")),
+            session=sqlite_session,
+            page=1,
+            per_page=20,
+        )
 
         with (
             patch("services.dataset_service.paginate_query") as mock_paginate,
@@ -330,7 +373,7 @@ class TestSegmentServiceQueries:
 
         assert result is child_chunk
 
-        session.scalar.return_value = SimpleNamespace()
+        session.scalar.return_value = object()
         result = SegmentService.get_child_chunk_by_id("child-a", "tenant-1", session)
 
         assert result is None
@@ -354,9 +397,14 @@ class TestSegmentServiceQueries:
         assert "child_chunks.document_id = 'doc-1'" in sql
         assert "child_chunks.segment_id = 'segment-1'" in sql
 
-    def test_get_segments_uses_status_and_keyword_filters(self):
+    def test_get_segments_uses_status_and_keyword_filters(self, sqlite_session: Session):
         session = MagicMock()
-        paginated = SimpleNamespace(items=["segment"], total=1)
+        paginated = SelectPagination(
+            select=select(literal("segment")),
+            session=sqlite_session,
+            page=1,
+            per_page=20,
+        )
 
         with (
             patch("services.dataset_service.paginate_query") as mock_paginate,
@@ -397,7 +445,7 @@ class TestSegmentServiceQueries:
 
         assert result is segment
 
-        session.scalar.return_value = SimpleNamespace()
+        session.scalar.return_value = object()
         result = SegmentService.get_segment_by_id("segment-1", "tenant-1", session)
 
         assert result is None
@@ -509,7 +557,7 @@ class TestSegmentServiceMutations:
             doc_form=IndexStructureType.QA_INDEX,
             word_count=0,
         )
-        refreshed_segment = SimpleNamespace(id="segment-1")
+        refreshed_segment = _make_segment()
         args = {
             "content": "question",
             "answer": "answer",
@@ -662,7 +710,7 @@ class TestSegmentServiceMutations:
         segment = _make_segment(content="same content", keywords=["old"])
         document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=20)
         dataset = _make_dataset()
-        refreshed_segment = SimpleNamespace(id=segment.id)
+        refreshed_segment = _make_segment(segment_id=segment.id)
         args = SegmentUpdateArgs(content="same content", keywords=["new"])
 
         with (
@@ -687,9 +735,9 @@ class TestSegmentServiceMutations:
             word_count=20,
         )
         dataset = _make_dataset(indexing_technique="high_quality")
-        refreshed_segment = SimpleNamespace(id=segment.id)
-        processing_rule = SimpleNamespace(id=document.dataset_process_rule_id)
-        existing_summary = SimpleNamespace(summary_content="old summary")
+        refreshed_segment = _make_segment(segment_id=segment.id)
+        processing_rule = _process_rule(document.dataset_process_rule_id)
+        existing_summary = _summary("old summary")
         embedding_model_instance = object()
         args = SegmentUpdateArgs(
             content="same content",
@@ -732,8 +780,8 @@ class TestSegmentServiceMutations:
         document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=10)
         dataset = _make_dataset(indexing_technique="high_quality")
         dataset.summary_index_setting = {"enable": True}
-        refreshed_segment = SimpleNamespace(id=segment.id)
-        existing_summary = SimpleNamespace(summary_content="old summary")
+        refreshed_segment = _make_segment(segment_id=segment.id)
+        existing_summary = _summary("old summary")
         embedding_model = MagicMock()
         embedding_model.get_text_embedding_num_tokens.return_value = [9]
         args = SegmentUpdateArgs(content="new content", keywords=["kw-1"])
@@ -771,8 +819,8 @@ class TestSegmentServiceMutations:
         document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=10)
         dataset = _make_dataset(indexing_technique="high_quality")
         dataset.summary_index_setting = {"enable": True}
-        refreshed_segment = SimpleNamespace(id=segment.id)
-        existing_summary = SimpleNamespace(summary_content="same summary")
+        refreshed_segment = _make_segment(segment_id=segment.id)
+        existing_summary = _summary("same summary")
         embedding_model = MagicMock()
         embedding_model.get_text_embedding_num_tokens.return_value = [7]
         args = SegmentUpdateArgs(content="new text", summary="same summary")
@@ -844,7 +892,7 @@ class TestSegmentServiceMutations:
         session = MagicMock()
         dataset = _make_dataset()
         document = _make_document(word_count=3)
-        current_user = SimpleNamespace(current_tenant_id="tenant-1")
+        current_user = _account()
 
         with (
             patch("services.dataset_service.current_user", current_user),
@@ -886,7 +934,7 @@ class TestSegmentServiceMutations:
         document = _make_document()
         segment_a = _make_segment(segment_id="segment-a", enabled=False)
         segment_b = _make_segment(segment_id="segment-b", enabled=False)
-        current_user = SimpleNamespace(id="user-1", current_tenant_id="tenant-1")
+        current_user = _account()
 
         with (
             patch("services.dataset_service.current_user", current_user),
@@ -913,7 +961,7 @@ class TestSegmentServiceMutations:
         document = _make_document()
         segment_a = _make_segment(segment_id="segment-a", enabled=True)
         segment_b = _make_segment(segment_id="segment-b", enabled=True)
-        current_user = SimpleNamespace(id="user-1", current_tenant_id="tenant-1")
+        current_user = _account()
 
         with (
             patch("services.dataset_service.current_user", current_user),
@@ -940,11 +988,11 @@ class TestSegmentServiceChildChunkTailHelpers:
 
     def test_update_child_chunk_rolls_back_on_vector_failure(self):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         child_chunk = _make_child_chunk()
 
         with (
-            patch("services.dataset_service.current_user", SimpleNamespace(id="user-1")),
+            patch("services.dataset_service.current_user", _account()),
             patch("services.dataset_service.naive_utc_now", return_value="now"),
             patch("services.dataset_service.VectorService") as vector_service,
         ):
@@ -952,7 +1000,12 @@ class TestSegmentServiceChildChunkTailHelpers:
 
             with pytest.raises(ChildChunkIndexingError, match="vector failed"):
                 SegmentService.update_child_chunk(
-                    "new content", child_chunk, SimpleNamespace(), SimpleNamespace(), dataset, session
+                    "new content",
+                    child_chunk,
+                    _make_segment(),
+                    _make_document(),
+                    dataset,
+                    session,
                 )
 
         session.rollback.assert_called_once()
@@ -960,7 +1013,7 @@ class TestSegmentServiceChildChunkTailHelpers:
 
     def test_delete_child_chunk_commits_after_successful_vector_delete(self):
         session = MagicMock()
-        dataset = SimpleNamespace(id="dataset-1")
+        dataset = _make_dataset()
         child_chunk = _make_child_chunk()
 
         with (
@@ -990,7 +1043,7 @@ class TestSegmentServiceAdditionalRegenerationBranches:
         segment = _make_segment(content="question", word_count=8)
         document = _make_document(doc_form=IndexStructureType.QA_INDEX, word_count=20)
         dataset = _make_dataset()
-        refreshed_segment = SimpleNamespace(id=segment.id)
+        refreshed_segment = _make_segment(segment_id=segment.id)
 
         with (
             patch("services.dataset_service.redis_client") as mock_redis,
@@ -1019,7 +1072,7 @@ class TestSegmentServiceAdditionalRegenerationBranches:
         segment = _make_segment(content="old", word_count=3)
         document = _make_document(doc_form=IndexStructureType.QA_INDEX, word_count=10)
         dataset = _make_dataset(indexing_technique="high_quality")
-        refreshed_segment = SimpleNamespace(id=segment.id)
+        refreshed_segment = _make_segment(segment_id=segment.id)
         embedding_model = MagicMock()
         embedding_model.get_text_embedding_num_tokens.return_value = [21]
 
@@ -1062,9 +1115,9 @@ class TestSegmentServiceAdditionalRegenerationBranches:
         )
         dataset = _make_dataset(indexing_technique="high_quality")
         dataset.embedding_model_provider = None
-        refreshed_segment = SimpleNamespace(id=segment.id)
-        processing_rule = SimpleNamespace(id=document.dataset_process_rule_id)
-        existing_summary = SimpleNamespace(summary_content="old summary")
+        refreshed_segment = _make_segment(segment_id=segment.id)
+        processing_rule = _process_rule(document.dataset_process_rule_id)
+        existing_summary = _summary("old summary")
         embedding_model_instance = object()
 
         with (
@@ -1119,7 +1172,7 @@ class TestSegmentServiceAdditionalRegenerationBranches:
             word_count=20,
         )
         dataset = _make_dataset(indexing_technique="economy")
-        refreshed_segment = SimpleNamespace(id=segment.id)
+        refreshed_segment = _make_segment(segment_id=segment.id)
 
         with (
             patch("services.dataset_service.redis_client") as mock_redis,

api/tests/unit_tests/services/test_notion_data_source_gateway.py

@@ -1,156 +0,0 @@
-import json
-from unittest.mock import MagicMock
-
-import httpx
-import pytest
-
-from core.helper.ssrf_proxy import SSRFProxy
-from core.tools.errors import ToolSSRFError
-from services.data_source_oauth_service import DataSourceOAuthError
-from services.notion_data_source_gateway import NotionDataSourceGateway
-
-
-def _gateway(handler: httpx.MockTransport) -> tuple[NotionDataSourceGateway, httpx.Client]:
-    client = httpx.Client(transport=handler)
-    return (
-        NotionDataSourceGateway(
-            client_id="client-id",
-            client_secret="client-secret",
-            redirect_uri="https://api.example/oauth/data-source/callback/notion",
-            http_client=client,
-        ),
-        client,
-    )
-
-
-def test_get_authorization_url_encodes_configured_values() -> None:
-    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(500)))
-    try:
-        result = gateway.get_authorization_url()
-    finally:
-        client.close()
-
-    assert result == (
-        "https://api.notion.com/v1/oauth/authorize?client_id=client-id&response_type=code&"
-        "redirect_uri=https%3A%2F%2Fapi.example%2Foauth%2Fdata-source%2Fcallback%2Fnotion&owner=user"
-    )
-
-
-def test_authorize_exchanges_code_and_builds_source_info_without_persistence() -> None:
-    search_types: list[str] = []
-
-    def handler(request: httpx.Request) -> httpx.Response:
-        if request.url.path == "/v1/oauth/token":
-            assert request.headers["authorization"].startswith("Basic ")
-            return httpx.Response(
-                200,
-                json={
-                    "access_token": "notion-token",
-                    "workspace_name": "Notion Workspace",
-                    "workspace_icon": "https://example/icon.png",
-                    "workspace_id": "notion-workspace",
-                },
-            )
-        if request.url.path == "/v1/search":
-            payload = json.loads(request.content)
-            search_types.append(payload["filter"]["value"])
-            return httpx.Response(200, json={"results": [], "has_more": False})
-        raise AssertionError(f"Unexpected Notion request: {request.url}")
-
-    gateway, client = _gateway(httpx.MockTransport(handler))
-    try:
-        authorization = gateway.authorize("authorization-code")
-    finally:
-        client.close()
-
-    assert authorization.access_token == "notion-token"
-    assert authorization.source_info == {
-        "workspace_name": "Notion Workspace",
-        "workspace_icon": "https://example/icon.png",
-        "workspace_id": "notion-workspace",
-        "pages": [],
-        "total": 0,
-    }
-    assert search_types == ["page", "database"]
-
-
-def test_refresh_preserves_workspace_metadata() -> None:
-    gateway, client = _gateway(
-        httpx.MockTransport(
-            lambda request: (
-                httpx.Response(200, json={"results": [], "has_more": False})
-                if request.url.path == "/v1/search"
-                else httpx.Response(404)
-            )
-        )
-    )
-    try:
-        source_info = gateway.refresh(
-            "notion-token",
-            {
-                "workspace_name": "Workspace",
-                "workspace_icon": None,
-                "workspace_id": "notion-workspace",
-                "pages": [
-                    {
-                        "page_id": "old",
-                        "page_name": "Old page",
-                        "page_icon": None,
-                        "parent_id": "root",
-                        "type": "page",
-                    }
-                ],
-                "total": 1,
-            },
-        )
-    finally:
-        client.close()
-
-    assert source_info == {
-        "workspace_name": "Workspace",
-        "workspace_icon": None,
-        "workspace_id": "notion-workspace",
-        "pages": [],
-        "total": 0,
-    }
-
-
-def test_search_rejects_missing_cursor_when_more_results_are_declared() -> None:
-    requests = 0
-
-    def handler(_request: httpx.Request) -> httpx.Response:
-        nonlocal requests
-        requests += 1
-        return httpx.Response(200, json={"results": [], "has_more": True, "next_cursor": None})
-
-    gateway, client = _gateway(httpx.MockTransport(handler))
-    try:
-        with pytest.raises(DataSourceOAuthError, match="valid next_cursor"):
-            gateway.notion_page_search("notion-token")
-    finally:
-        client.close()
-
-    assert requests == 1
-
-
-def test_provider_http_failure_is_translated_to_application_boundary_error() -> None:
-    gateway, client = _gateway(httpx.MockTransport(lambda _request: httpx.Response(503, text="unavailable")))
-    try:
-        with pytest.raises(DataSourceOAuthError, match="unsuccessful response"):
-            gateway.authorize("authorization-code")
-    finally:
-        client.close()
-
-
-def test_ssrf_rejection_is_translated_to_application_boundary_error() -> None:
-    client = MagicMock(spec=SSRFProxy)
-    client.post.side_effect = ToolSSRFError("blocked")
-    gateway = NotionDataSourceGateway(
-        client_id="client-id",
-        client_secret="client-secret",
-        redirect_uri="https://api.example/oauth/data-source/callback/notion",
-        http_client=client,
-    )
-
-    with pytest.raises(DataSourceOAuthError, match="Notion request failed"):
-        gateway.authorize("authorization-code")

packages/contracts/generated/api/console/datasets/types.gen.ts

@@ -931,7 +931,7 @@ export type WeightModel = {
 
 export type PreProcessingRule = {
   enabled: boolean
-  id: 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails'
+  id: PreProcessingRuleKey
 }
 
 export type Segmentation = {
@@ -1064,6 +1064,8 @@ export type WeightVectorSetting = {
   vector_weight: number
 }
 
+export type PreProcessingRuleKey = 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails'
+
 export type HitTestingDocument = {
   data_source_type: string
   doc_metadata: unknown | null

packages/contracts/generated/api/console/datasets/zod.gen.ts

@@ -828,14 +828,6 @@ export const zRetrievalMethod = z.enum([
   'semantic_search',
 ])
 
-/**
- * PreProcessingRule
- */
-export const zPreProcessingRule = z.object({
-  enabled: z.boolean(),
-  id: z.enum(['remove_extra_spaces', 'remove_stopwords', 'remove_urls_emails']),
-})
-
 /**
  * Segmentation
  */
@@ -845,33 +837,6 @@ export const zSegmentation = z.object({
   separator: z.string().optional().default('\n'),
 })
 
-/**
- * Rule
- */
-export const zRule = z.object({
-  parent_mode: z.enum(['full-doc', 'paragraph']).nullish(),
-  pre_processing_rules: z.array(zPreProcessingRule).nullish(),
-  segmentation: zSegmentation.nullish(),
-  subchunk_segmentation: zSegmentation.nullish(),
-})
-
-/**
- * ProcessRuleResponse
- */
-export const zProcessRuleResponse = z.object({
-  limits: z.record(z.string(), z.unknown()),
-  mode: zProcessRuleMode,
-  rules: zRule.nullish(),
-})
-
-/**
- * ProcessRule
- */
-export const zProcessRule = z.object({
-  mode: zProcessRuleMode,
-  rules: zRule.nullish(),
-})
-
 /**
  * MetadataDetail
  */
@@ -1279,6 +1244,47 @@ export const zHitTestingPayload = z.object({
   retrieval_model: zRetrievalModel.nullish(),
 })
 
+export const zPreProcessingRuleKey = z.enum([
+  'remove_extra_spaces',
+  'remove_stopwords',
+  'remove_urls_emails',
+])
+
+/**
+ * PreProcessingRule
+ */
+export const zPreProcessingRule = z.object({
+  enabled: z.boolean(),
+  id: zPreProcessingRuleKey,
+})
+
+/**
+ * Rule
+ */
+export const zRule = z.object({
+  parent_mode: z.enum(['full-doc', 'paragraph']).nullish(),
+  pre_processing_rules: z.array(zPreProcessingRule).nullish(),
+  segmentation: zSegmentation.nullish(),
+  subchunk_segmentation: zSegmentation.nullish(),
+})
+
+/**
+ * ProcessRuleResponse
+ */
+export const zProcessRuleResponse = z.object({
+  limits: z.record(z.string(), z.unknown()),
+  mode: zProcessRuleMode,
+  rules: zRule.nullish(),
+})
+
+/**
+ * ProcessRule
+ */
+export const zProcessRule = z.object({
+  mode: zProcessRuleMode,
+  rules: zRule.nullish(),
+})
+
 /**
  * HitTestingDocument
  */

packages/contracts/generated/api/service/types.gen.ts

@@ -1779,9 +1779,11 @@ export type PipelineUploadFileResponse = {
 
 export type PreProcessingRule = {
   enabled: boolean
-  id: 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails'
+  id: PreProcessingRuleKey
 }
 
+export type PreProcessingRuleKey = 'remove_extra_spaces' | 'remove_stopwords' | 'remove_urls_emails'
+
 export type ProcessRule = {
   mode: ProcessRuleMode
   rules?: Rule | null

packages/contracts/generated/api/service/zod.gen.ts

@@ -1670,12 +1670,18 @@ export const zPipelineUploadFileResponse = z.object({
   size: z.int(),
 })
 
+export const zPreProcessingRuleKey = z.enum([
+  'remove_extra_spaces',
+  'remove_stopwords',
+  'remove_urls_emails',
+])
+
 /**
  * PreProcessingRule
  */
 export const zPreProcessingRule = z.object({
   enabled: z.boolean(),
-  id: z.enum(['remove_extra_spaces', 'remove_stopwords', 'remove_urls_emails']),
+  id: zPreProcessingRuleKey,
 })
 
 /**