diff --git a/api/AGENTS.md b/api/AGENTS.md index 984322590b3..0b7082113c7 100644 --- a/api/AGENTS.md +++ b/api/AGENTS.md @@ -109,6 +109,18 @@ class Example: - Reuse existing helpers in `core/`, `services/`, and `libs/` before creating new abstractions. - Optimise for observability: deterministic control flow, clear logging, actionable errors. +### Owner-Bound Resource References + +- Resolve and validate the outer owner before binding a nested resource ID. +- For stable single-parent chains, use immutable nested `NamedTuple` refs. +- Root refs carry tenant plus root ID; child refs carry the parent ref. +- In production, construct refs through the domain ref service. +- Python allowing direct construction does not grant authorization. +- Scope every consuming query with complete owner predicates; refs are not security tokens. +- Keep polymorphic owners flat until explicit nominal owner types exist. +- Do not add generic ref bases or compatibility fields only for uniformity. +- Reconstruct internal refs from validated database state after payload or async boundaries. + ### Logging & Errors - Never use `print`; use a module-level logger: diff --git a/api/controllers/console/app/mcp_server.py b/api/controllers/console/app/mcp_server.py index 6042b13d831..f5ab8b103ad 100644 --- a/api/controllers/console/app/mcp_server.py +++ b/api/controllers/console/app/mcp_server.py @@ -153,8 +153,8 @@ class AppMCPServerController(Resource): select(AppMCPServer) .where( AppMCPServer.id == server_ref.server_id, - AppMCPServer.tenant_id == server_ref.tenant_id, - AppMCPServer.app_id == server_ref.app_id, + AppMCPServer.tenant_id == server_ref.app.tenant_id, + AppMCPServer.app_id == server_ref.app.app_id, ) .limit(1) ) diff --git a/api/core/app/apps/pipeline/pipeline_runner.py b/api/core/app/apps/pipeline/pipeline_runner.py index 3ad0990cbb4..39bd4c6980c 100644 --- a/api/core/app/apps/pipeline/pipeline_runner.py +++ b/api/core/app/apps/pipeline/pipeline_runner.py @@ -27,9 +27,10 @@ from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent from graphon.runtime import GraphRuntimeState, VariablePool from graphon.variable_loader import VariableLoader from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput -from models.dataset import Document, Pipeline +from models.dataset import Pipeline from models.model import EndUser from models.workflow import Workflow +from services.dataset_ref_service import DatasetRefService, DocumentRef logger = logging.getLogger(__name__) @@ -93,9 +94,38 @@ class PipelineRunner(WorkflowBasedAppRunner): user_id = self.application_generate_entity.user_id pipeline = session.get(Pipeline, app_config.app_id) - if not pipeline: + if not pipeline or pipeline.tenant_id != app_config.tenant_id: raise ValueError("Pipeline not found") + dataset = pipeline.retrieve_dataset(session) + if ( + not dataset + or dataset.tenant_id != pipeline.tenant_id + or dataset.id != self.application_generate_entity.dataset_id + ): + raise ValueError("Pipeline dataset not found") + + document_id = self.application_generate_entity.document_id + original_document_id = self.application_generate_entity.original_document_id + dataset_ref = DatasetRefService.create_dataset_ref(dataset) + document_ref = ( + DatasetRefService.create_document_ref_from_id( + dataset_ref, + document_id, + ) + if document_id + else None + ) + if document_ref and DatasetRefService.get_document_by_ref(document_ref, session=session) is None: + raise ValueError("Pipeline document not found") + if original_document_id and original_document_id != document_id: + original_document_ref = DatasetRefService.create_document_ref_from_id( + dataset_ref, + original_document_id, + ) + if DatasetRefService.get_document_by_ref(original_document_ref, session=session) is None: + raise ValueError("Pipeline original document not found") + workflow = self.get_workflow(session=session, pipeline=pipeline, workflow_id=app_config.workflow_id) if not workflow: raise ValueError("Workflow not initialized") @@ -206,9 +236,7 @@ class PipelineRunner(WorkflowBasedAppRunner): generator = workflow_entry.run() for event in generator: - self._update_document_status( - event, self.application_generate_entity.document_id, self.application_generate_entity.dataset_id - ) + self._update_document_status(event, document_ref) self._handle_event(workflow_entry, event) def get_workflow(self, session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None: @@ -295,17 +323,14 @@ class PipelineRunner(WorkflowBasedAppRunner): return graph - def _update_document_status(self, event: GraphEngineEvent, document_id: str | None, dataset_id: str | None) -> None: - """ - Update document status - """ - if isinstance(event, GraphRunFailedEvent): - if document_id and dataset_id: - with create_session() as session, session.begin(): - document = session.scalar( - select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1) - ) - if document: - document.indexing_status = "error" - document.error = event.error or "Unknown error" - session.add(document) + def _update_document_status(self, event: GraphEngineEvent, document_ref: DocumentRef | None) -> None: + """Set an owner-bound document to error after a failed graph run, if it exists.""" + if not isinstance(event, GraphRunFailedEvent) or document_ref is None: + return + + with create_session() as session, session.begin(): + document = DatasetRefService.get_document_by_ref(document_ref, session=session) + if document: + document.indexing_status = "error" + document.error = event.error or "Unknown error" + session.add(document) diff --git a/api/core/rag/index_processor/index_processor.py b/api/core/rag/index_processor/index_processor.py index 39254bb0ac7..bf6eb0a1262 100644 --- a/api/core/rag/index_processor/index_processor.py +++ b/api/core/rag/index_processor/index_processor.py @@ -65,14 +65,22 @@ class IndexProcessor: *, session: Session, ) -> IndexingResultDict: - document = session.scalar(select(Document).where(Document.id == document_id).limit(1)) - if not document: - raise KnowledgeIndexNodeError(f"Document {document_id} not found.") - dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) if not dataset: raise KnowledgeIndexNodeError(f"Dataset {dataset_id} not found.") + document = session.scalar( + select(Document) + .where( + Document.id == document_id, + Document.dataset_id == dataset.id, + Document.tenant_id == dataset.tenant_id, + ) + .limit(1) + ) + if not document: + raise KnowledgeIndexNodeError(f"Document {document_id} not found.") + dataset_name_value = dataset.name document_name_value = document.name created_at_value = document.created_at @@ -83,7 +91,11 @@ class IndexProcessor: index_processor = IndexProcessorFactory(dataset.chunk_structure).init_index_processor() if original_document_id: segments = session.scalars( - select(DocumentSegment).where(DocumentSegment.document_id == original_document_id) + select(DocumentSegment).where( + DocumentSegment.document_id == original_document_id, + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.tenant_id == dataset.tenant_id, + ) ).all() if segments: index_node_ids = [segment.index_node_id for segment in segments if segment.index_node_id] @@ -97,7 +109,11 @@ class IndexProcessor: dataset, index_node_ids, with_keywords=True, delete_child_chunks=True, session=session ) session.commit() - segment_delete_stmt = delete(DocumentSegment).where(DocumentSegment.document_id == original_document_id) + segment_delete_stmt = delete(DocumentSegment).where( + DocumentSegment.document_id == original_document_id, + DocumentSegment.dataset_id == dataset.id, + DocumentSegment.tenant_id == dataset.tenant_id, + ) session.execute(segment_delete_stmt) session.commit() @@ -113,6 +129,7 @@ class IndexProcessor: select(func.sum(DocumentSegment.word_count)).where( DocumentSegment.document_id == document_id, DocumentSegment.dataset_id == dataset_id, + DocumentSegment.tenant_id == dataset.tenant_id, ) ) ) or 0 @@ -125,6 +142,7 @@ class IndexProcessor: .where( DocumentSegment.document_id == document_id, DocumentSegment.dataset_id == dataset_id, + DocumentSegment.tenant_id == dataset.tenant_id, ) .values( status="completed", @@ -156,15 +174,23 @@ class IndexProcessor: session: Session, ) -> Preview: doc_language = None - if document_id: - document = session.scalar(select(Document).where(Document.id == document_id).limit(1)) - else: - document = None - dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1)) if not dataset: raise KnowledgeIndexNodeError(f"Dataset {dataset_id} not found.") + if document_id: + document = session.scalar( + select(Document) + .where( + Document.id == document_id, + Document.dataset_id == dataset.id, + Document.tenant_id == dataset.tenant_id, + ) + .limit(1) + ) + else: + document = None + if summary_index_setting is None: summary_index_setting = dataset.summary_index_setting diff --git a/api/services/annotation_service.py b/api/services/annotation_service.py index c36db8d4b51..947c35fc0a0 100644 --- a/api/services/annotation_service.py +++ b/api/services/annotation_service.py @@ -96,7 +96,7 @@ class AppAnnotationService: select(MessageAnnotation) .where( MessageAnnotation.id == annotation_ref.annotation_id, - MessageAnnotation.app_id == annotation_ref.app_id, + MessageAnnotation.app_id == annotation_ref.app.app_id, ) .limit(1) ) @@ -343,15 +343,15 @@ class AppAnnotationService: session.commit() # if annotation reply is enabled , add annotation to index app_annotation_setting = session.scalar( - select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == annotation_ref.app_id).limit(1) + select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == annotation_ref.app.app_id).limit(1) ) if app_annotation_setting: update_annotation_to_index_task.delay( annotation.id, annotation.question_text, - annotation_ref.tenant_id, - annotation_ref.app_id, + annotation_ref.app.tenant_id, + annotation_ref.app.app_id, app_annotation_setting.collection_binding_id, ) @@ -368,7 +368,7 @@ class AppAnnotationService: annotation_hit_histories = session.scalars( select(AppAnnotationHitHistory).where( - AppAnnotationHitHistory.app_id == annotation_ref.app_id, + AppAnnotationHitHistory.app_id == annotation_ref.app.app_id, AppAnnotationHitHistory.annotation_id == annotation_ref.annotation_id, ) ).all() @@ -379,14 +379,14 @@ class AppAnnotationService: session.commit() # if annotation reply is enabled , delete annotation index app_annotation_setting = session.scalar( - select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == annotation_ref.app_id).limit(1) + select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == annotation_ref.app.app_id).limit(1) ) if app_annotation_setting: delete_annotation_index_task.delay( annotation.id, - annotation_ref.app_id, - annotation_ref.tenant_id, + annotation_ref.app.app_id, + annotation_ref.app.tenant_id, app_annotation_setting.collection_binding_id, ) @@ -396,7 +396,10 @@ class AppAnnotationService: annotations_to_delete = session.execute( select(MessageAnnotation, AppAnnotationSetting) .outerjoin(AppAnnotationSetting, MessageAnnotation.app_id == AppAnnotationSetting.app_id) - .where(MessageAnnotation.id.in_(annotation_ids), MessageAnnotation.app_id == app_ref.app_id) + .where( + MessageAnnotation.id.in_(annotation_ids), + MessageAnnotation.app_id == app_ref.app_id, + ) ).all() if not annotations_to_delete: @@ -576,7 +579,7 @@ class AppAnnotationService: stmt = ( select(AppAnnotationHitHistory) .where( - AppAnnotationHitHistory.app_id == annotation_ref.app_id, + AppAnnotationHitHistory.app_id == annotation_ref.app.app_id, AppAnnotationHitHistory.annotation_id == annotation_ref.annotation_id, ) .order_by(AppAnnotationHitHistory.created_at.desc()) diff --git a/api/services/app_ref_service.py b/api/services/app_ref_service.py index 5f794491811..0fa9a7f9aee 100644 --- a/api/services/app_ref_service.py +++ b/api/services/app_ref_service.py @@ -15,8 +15,7 @@ class AppRef(NamedTuple): class MessageRef(NamedTuple): """Message identifiers used to scope downstream resource lookups.""" - tenant_id: str - app_id: str + app: AppRef message_id: str end_user_id: str | None = None account_id: str | None = None @@ -25,16 +24,14 @@ class MessageRef(NamedTuple): class AnnotationRef(NamedTuple): """Annotation identifiers used to scope downstream resource lookups.""" - tenant_id: str - app_id: str + app: AppRef annotation_id: str class AppMCPServerRef(NamedTuple): """MCP server identifiers used to scope downstream resource lookups.""" - tenant_id: str - app_id: str + app: AppRef server_id: str @@ -54,8 +51,7 @@ class AppRefService: account_id: str | None = None, ) -> MessageRef: return MessageRef( - tenant_id=app_ref.tenant_id, - app_id=app_ref.app_id, + app=app_ref, message_id=message_id, end_user_id=end_user_id, account_id=account_id, @@ -63,8 +59,10 @@ class AppRefService: @staticmethod def create_annotation_ref(app_ref: AppRef, annotation_id: str) -> AnnotationRef: - return AnnotationRef(tenant_id=app_ref.tenant_id, app_id=app_ref.app_id, annotation_id=annotation_id) + """Bind a candidate annotation ID; ownership is enforced when the ref is consumed.""" + return AnnotationRef(app=app_ref, annotation_id=annotation_id) @staticmethod def create_mcp_server_ref(app_ref: AppRef, server_id: str) -> AppMCPServerRef: - return AppMCPServerRef(tenant_id=app_ref.tenant_id, app_id=app_ref.app_id, server_id=server_id) + """Bind a candidate MCP server ID; ownership is enforced when the ref is consumed.""" + return AppMCPServerRef(app=app_ref, server_id=server_id) diff --git a/api/services/audio_service.py b/api/services/audio_service.py index 0d267e2277c..ea085116ce5 100644 --- a/api/services/audio_service.py +++ b/api/services/audio_service.py @@ -37,7 +37,10 @@ logger = logging.getLogger(__name__) class AudioService: @staticmethod def _get_message_by_ref(session: Session, message_ref: MessageRef) -> Message | None: - stmt = select(Message).where(Message.id == message_ref.message_id, Message.app_id == message_ref.app_id) + stmt = select(Message).where( + Message.id == message_ref.message_id, + Message.app_id == message_ref.app.app_id, + ) if message_ref.end_user_id is not None: stmt = stmt.where(Message.from_end_user_id == message_ref.end_user_id) if message_ref.account_id is not None: diff --git a/api/services/dataset_ref_service.py b/api/services/dataset_ref_service.py index 049a8d05759..6bd5cf444ff 100644 --- a/api/services/dataset_ref_service.py +++ b/api/services/dataset_ref_service.py @@ -2,6 +2,9 @@ from typing import NamedTuple +from sqlalchemy import select +from sqlalchemy.orm import Session + from models.dataset import Dataset, Document @@ -13,44 +16,50 @@ class DatasetRef(NamedTuple): class DocumentRef(NamedTuple): - """Document identifiers used to scope downstream resource lookups.""" + """Owner-bound lookup coordinates, not proof that a document is authorized or exists.""" - tenant_id: str - dataset_id: str + dataset: DatasetRef document_id: str class SegmentRef(NamedTuple): """Segment identifiers used to scope downstream resource lookups.""" - tenant_id: str - dataset_id: str - document_id: str + document: DocumentRef segment_id: str class DatasetRefService: - """Factory helpers for dataset, document, and segment refs.""" + """Build child locators from validated dataset roots and resolve them with owner predicates.""" @staticmethod def create_dataset_ref(dataset: Dataset) -> DatasetRef: + """Create a root ref from a dataset already validated by the caller.""" return DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) @staticmethod def create_document_ref(dataset_ref: DatasetRef, document: Document) -> DocumentRef | None: if document.tenant_id != dataset_ref.tenant_id or document.dataset_id != dataset_ref.dataset_id: return None - return DocumentRef( - tenant_id=dataset_ref.tenant_id, - dataset_id=dataset_ref.dataset_id, - document_id=document.id, - ) + return DatasetRefService.create_document_ref_from_id(dataset_ref, document.id) + + @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) @staticmethod def create_segment_ref(document_ref: DocumentRef, segment_id: str) -> SegmentRef: - return SegmentRef( - tenant_id=document_ref.tenant_id, - dataset_id=document_ref.dataset_id, - document_id=document_ref.document_id, - segment_id=segment_id, + """Bind a candidate segment ID; ownership is enforced when the ref is consumed.""" + return SegmentRef(document=document_ref, segment_id=segment_id) + + @staticmethod + def get_document_by_ref(document_ref: DocumentRef, *, session: Session) -> Document | None: + """Resolve a document through its complete tenant and dataset ownership chain.""" + return 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, + ) ) diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index f595888ff9e..880166890e3 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -4163,9 +4163,9 @@ class SegmentService: select(ChildChunk) .where( ChildChunk.id == child_chunk_id, - ChildChunk.tenant_id == segment_ref.tenant_id, - ChildChunk.dataset_id == segment_ref.dataset_id, - ChildChunk.document_id == segment_ref.document_id, + ChildChunk.tenant_id == segment_ref.document.dataset.tenant_id, + ChildChunk.dataset_id == segment_ref.document.dataset.dataset_id, + ChildChunk.document_id == segment_ref.document.document_id, ChildChunk.segment_id == segment_ref.segment_id, ) .limit(1) @@ -4219,9 +4219,9 @@ class SegmentService: select(DocumentSegment) .where( DocumentSegment.id == segment_ref.segment_id, - DocumentSegment.tenant_id == segment_ref.tenant_id, - DocumentSegment.dataset_id == segment_ref.dataset_id, - DocumentSegment.document_id == segment_ref.document_id, + DocumentSegment.tenant_id == segment_ref.document.dataset.tenant_id, + DocumentSegment.dataset_id == segment_ref.document.dataset.dataset_id, + DocumentSegment.document_id == segment_ref.document.document_id, ) .limit(1) ) diff --git a/api/services/rag_pipeline/pipeline_generate_service.py b/api/services/rag_pipeline/pipeline_generate_service.py index d7ef2205380..182e3fb283e 100644 --- a/api/services/rag_pipeline/pipeline_generate_service.py +++ b/api/services/rag_pipeline/pipeline_generate_service.py @@ -6,10 +6,11 @@ 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 models.dataset import Document, Pipeline +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.rag_pipeline.rag_pipeline import RagPipelineService @@ -37,8 +38,12 @@ class PipelineGenerateService: try: workflow = cls._get_workflow(pipeline, invoke_from, session) if original_document_id := args.get("original_document_id"): - # update document status to waiting - cls.update_document_status(original_document_id, session=session) + dataset = pipeline.retrieve_dataset(session) + if dataset is None or dataset.tenant_id != pipeline.tenant_id: + raise ValueError("Pipeline dataset is required") + dataset_ref = DatasetRefService.create_dataset_ref(dataset) + document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, original_document_id) + cls.update_document_status(document_ref, session=session) return PipelineGenerator.convert_to_event_stream( PipelineGenerator().generate( session=session, @@ -123,12 +128,10 @@ class PipelineGenerateService: return workflow @classmethod - def update_document_status(cls, document_id: str, *, session: Session): - """ - Update document status to waiting - :param document_id: document id - """ - document = session.get(Document, document_id) - if document: - document.indexing_status = IndexingStatus.WAITING - session.add(document) + def update_document_status(cls, document_ref: DocumentRef, *, session: Session) -> None: + """Set a document in the owner-bound dataset to waiting.""" + document = DatasetRefService.get_document_by_ref(document_ref, session=session) + if document is None: + raise ValueError("Pipeline document not found") + document.indexing_status = IndexingStatus.WAITING + session.add(document) diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 3e8340c8a11..f700e0088e3 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -73,6 +73,7 @@ from models.workflow import ( WorkflowType, ) from repositories.factory import DifyAPIRepositoryFactory +from services.dataset_ref_service import DatasetRefService from services.datasource_provider_service import DatasourceProviderService from services.entities.knowledge_entities.rag_pipeline_entities import ( KnowledgeConfiguration, @@ -1013,23 +1014,24 @@ class RagPipelineService: dataset_id = get_system_segment(variable_pool, SystemVariableKey.DATASET_ID) pipeline_id = get_system_segment(variable_pool, SystemVariableKey.APP_ID) if document_id and dataset_id and pipeline_id: - document = self._session.scalar( - select(Document) - .join(Dataset, Dataset.id == Document.dataset_id) + dataset = self._session.scalar( + select(Dataset) .where( - Document.id == document_id.value, - Document.tenant_id == tenant_id, - Document.dataset_id == dataset_id.value, + Dataset.id == dataset_id.value, Dataset.tenant_id == tenant_id, Dataset.pipeline_id == pipeline_id.value, ) .limit(1) ) - if document: - document.indexing_status = IndexingStatus.ERROR - document.error = error - self._session.add(document) - self._session.commit() + if dataset: + dataset_ref = DatasetRefService.create_dataset_ref(dataset) + document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, document_id.value) + document = DatasetRefService.get_document_by_ref(document_ref, session=self._session) + if document: + document.indexing_status = IndexingStatus.ERROR + document.error = error + self._session.add(document) + self._session.commit() return workflow_node_execution diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/test_containers_integration_tests/services/test_annotation_service.py index 0aa6cc820a5..b416aee33fe 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_annotation_service.py @@ -9,7 +9,7 @@ from models import Account from models.enums import ConversationFromSource, InvokeFrom from models.model import MessageAnnotation from services.annotation_service import AppAnnotationService -from services.app_ref_service import AnnotationRef +from services.app_ref_service import AnnotationRef, AppRef from services.app_service import AppService, CreateAppParams from tests.test_containers_integration_tests.helpers import generate_valid_password @@ -122,7 +122,7 @@ class TestAnnotationService: @staticmethod def _annotation_ref(app, annotation_id: str) -> AnnotationRef: - return AnnotationRef(tenant_id=app.tenant_id, app_id=app.id, annotation_id=annotation_id) + return AnnotationRef(app=AppRef(tenant_id=app.tenant_id, app_id=app.id), annotation_id=annotation_id) def _create_test_conversation(self, db_session_with_containers: Session, app, account, fake): """ @@ -624,8 +624,7 @@ class TestAnnotationService: non_existent_app_id = fake.uuid4() annotation_id = fake.uuid4() app_ref = AnnotationRef( - tenant_id=fake.uuid4(), - app_id=non_existent_app_id, + app=AppRef(tenant_id=fake.uuid4(), app_id=non_existent_app_id), annotation_id=annotation_id, ) diff --git a/api/tests/test_containers_integration_tests/services/test_audio_service_db.py b/api/tests/test_containers_integration_tests/services/test_audio_service_db.py index d1a025c2067..0fe10029158 100644 --- a/api/tests/test_containers_integration_tests/services/test_audio_service_db.py +++ b/api/tests/test_containers_integration_tests/services/test_audio_service_db.py @@ -24,7 +24,7 @@ from core.app.entities.app_invoke_entities import InvokeFrom from models.account import TenantAccountJoin from models.enums import ConversationFromSource, MessageStatus from models.model import App, AppMode, Conversation, Message -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from tests.test_containers_integration_tests.controllers.console.helpers import ( create_console_account_and_tenant, @@ -161,8 +161,7 @@ class TestAudioServiceTranscriptTTSMessageLookup: app_model=app, session=db_session_with_containers, message_ref=MessageRef( - tenant_id=app.tenant_id, - app_id=app.id, + app=AppRef(tenant_id=app.tenant_id, app_id=app.id), message_id=message.id, account_id=account_id, ), @@ -182,7 +181,10 @@ class TestAudioServiceTranscriptTTSMessageLookup: result = AudioService.transcript_tts( app_model=app, session=db_session_with_containers, - message_ref=MessageRef(tenant_id=app.tenant_id, app_id=app.id, message_id="invalid-uuid"), + message_ref=MessageRef( + app=AppRef(tenant_id=app.tenant_id, app_id=app.id), + message_id="invalid-uuid", + ), ) assert result is None @@ -194,7 +196,10 @@ class TestAudioServiceTranscriptTTSMessageLookup: result = AudioService.transcript_tts( app_model=app, session=db_session_with_containers, - message_ref=MessageRef(tenant_id=app.tenant_id, app_id=app.id, message_id=str(uuid4())), + message_ref=MessageRef( + app=AppRef(tenant_id=app.tenant_id, app_id=app.id), + message_id=str(uuid4()), + ), ) assert result is None @@ -216,8 +221,7 @@ class TestAudioServiceTranscriptTTSMessageLookup: app_model=app, session=db_session_with_containers, message_ref=MessageRef( - tenant_id=app.tenant_id, - app_id=app.id, + app=AppRef(tenant_id=app.tenant_id, app_id=app.id), message_id=message.id, account_id=account_id, ), diff --git a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py index 7f4bd96b28f..e722f943820 100644 --- a/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py +++ b/api/tests/test_containers_integration_tests/services/test_dataset_service_document.py @@ -15,7 +15,7 @@ from models import Account from models.dataset import Dataset, Document from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import UploadFile -from services.dataset_ref_service import DatasetRef +from services.dataset_ref_service import DatasetRefService from services.dataset_service import DocumentService from services.errors.account import NoPermissionError @@ -616,7 +616,7 @@ def test_delete_document_emits_signal_and_commits(db_session_with_containers: Se def test_delete_documents_ignores_empty_input(db_session_with_containers: Session): dataset = DocumentServiceIntegrationFactory.create_dataset(db_session_with_containers) - dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + dataset_ref = DatasetRefService.create_dataset_ref(dataset) with patch("services.dataset_service.batch_clean_document_task.delay") as delay: DocumentService.delete_documents(dataset_ref, [], dataset.doc_form, session=db_session_with_containers) @@ -651,7 +651,7 @@ def test_delete_documents_deletes_rows_and_dispatches_cleanup_task(db_session_wi position=2, data_source_info={"upload_file_id": upload_file_b.id}, ) - dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + dataset_ref = DatasetRefService.create_dataset_ref(dataset) with patch("services.dataset_service.batch_clean_document_task.delay") as delay: DocumentService.delete_documents( diff --git a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py b/api/tests/unit_tests/controllers/console/app/test_annotation_api.py index b12a015217e..33f19ae3550 100644 --- a/api/tests/unit_tests/controllers/console/app/test_annotation_api.py +++ b/api/tests/unit_tests/controllers/console/app/test_annotation_api.py @@ -168,7 +168,7 @@ class TestConsoleAnnotationRefBoundaries: assert response["question"] == "q" update_mock.assert_called_once() - assert update_mock.call_args.args[1] == AnnotationRef("tenant-1", "app-1", "ann-1") + assert update_mock.call_args.args[1] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") assert update_mock.call_args.args[2] is session def test_delete_uses_annotation_ref(self, app: Flask): @@ -192,7 +192,7 @@ class TestConsoleAnnotationRefBoundaries: assert response == "" assert status == 204 delete_mock.assert_called_once() - assert delete_mock.call_args.args[0] == AnnotationRef("tenant-1", "app-1", "ann-1") + assert delete_mock.call_args.args[0] == AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1") assert delete_mock.call_args.args[1] is session def test_hit_history_uses_annotation_ref(self, app: Flask): @@ -223,4 +223,4 @@ class TestConsoleAnnotationRefBoundaries: response = handler(api, session, "app-1", "ann-1") assert response["total"] == 1 - hit_history_mock.assert_called_once_with(AnnotationRef("tenant-1", "app-1", "ann-1"), 2, 5, session) + hit_history_mock.assert_called_once_with(AnnotationRef(AppRef("tenant-1", "app-1"), "ann-1"), 2, 5, session) diff --git a/api/tests/unit_tests/controllers/console/app/test_audio.py b/api/tests/unit_tests/controllers/console/app/test_audio.py index 3a23ccf7f09..8b661aa2646 100644 --- a/api/tests/unit_tests/controllers/console/app/test_audio.py +++ b/api/tests/unit_tests/controllers/console/app/test_audio.py @@ -37,7 +37,7 @@ from models.agent import AgentConfigDraftType from models.agent_config_entities import AgentSoulConfig from services.agent.composer_service import AgentComposerService from services.agent.errors import AgentVersionNotFoundError -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from services.errors.app_model_config import AppModelConfigBrokenError from services.errors.audio import ( @@ -318,7 +318,7 @@ def test_console_text_api_builds_message_ref(app: Flask, monkeypatch: pytest.Mon response = handler(api, app_model=app_model) assert response == {"audio": "ok"} - assert calls["message_ref"] == MessageRef("tenant-1", "app-1", "message-1", account_id="account-1") + assert calls["message_ref"] == MessageRef(AppRef("tenant-1", "app-1"), "message-1", account_id="account-1") def test_console_text_api_error_mapping(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/controllers/console/explore/test_audio.py b/api/tests/unit_tests/controllers/console/explore/test_audio.py index e3c856926ca..b21e0e4b2c6 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_audio.py +++ b/api/tests/unit_tests/controllers/console/explore/test_audio.py @@ -22,7 +22,7 @@ from core.errors.error import ( QuotaExceededError, ) from graphon.model_runtime.errors.invoke import InvokeError -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.errors.audio import ( AudioTooLargeServiceError, NoAudioUploadedServiceError, @@ -277,8 +277,7 @@ class TestChatTextApi: assert resp == {"audio": "ok"} assert transcript_tts.call_args.kwargs["message_ref"] == MessageRef( - tenant_id="tenant-1", - app_id="app-1", + app=AppRef(tenant_id="tenant-1", app_id="app-1"), message_id="m1", account_id="account-1", ) diff --git a/api/tests/unit_tests/controllers/console/explore/test_trial.py b/api/tests/unit_tests/controllers/console/explore/test_trial.py index ed853952255..570e63a3b7b 100644 --- a/api/tests/unit_tests/controllers/console/explore/test_trial.py +++ b/api/tests/unit_tests/controllers/console/explore/test_trial.py @@ -36,7 +36,7 @@ from graphon.model_runtime.errors.invoke import InvokeError from models import Account from models.account import TenantStatus from models.model import AppMode -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.errors.audio import SpeechToTextDisabledServiceError from services.errors.conversation import ConversationNotExistsError from services.errors.llm import InvokeRateLimitError @@ -934,8 +934,7 @@ class TestTrialChatTextApi: assert result == {"audio": "base64_data"} assert transcript_tts.call_args.kwargs["message_ref"] == MessageRef( - "tenant-1", - "a-chat", + AppRef("tenant-1", "a-chat"), "message-1", account_id="u1", ) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_audio.py b/api/tests/unit_tests/controllers/service_api/app/test_audio.py index 28ebbb6f190..2019891d281 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_audio.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_audio.py @@ -34,7 +34,7 @@ from controllers.service_api.app.error import ( ) from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from services.errors.app_model_config import AppModelConfigBrokenError from services.errors.audio import ( @@ -285,7 +285,7 @@ class TestTextApi: response = handler(api, app_model=app_model, end_user=end_user) assert response == {"audio": "ok"} - assert calls["message_ref"] == MessageRef("tenant-1", "a1", "message-1", end_user_id="end-user-1") + assert calls["message_ref"] == MessageRef(AppRef("tenant-1", "a1"), "message-1", end_user_id="end-user-1") def test_error_mapping(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( diff --git a/api/tests/unit_tests/controllers/web/test_audio.py b/api/tests/unit_tests/controllers/web/test_audio.py index 4764a95940b..6bbde4df887 100644 --- a/api/tests/unit_tests/controllers/web/test_audio.py +++ b/api/tests/unit_tests/controllers/web/test_audio.py @@ -23,7 +23,7 @@ from controllers.web.error import ( ) from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError from graphon.model_runtime.errors.invoke import InvokeError -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.errors.audio import ( AudioTooLargeServiceError, NoAudioUploadedServiceError, @@ -147,8 +147,7 @@ class TestTextApi: assert result == "audio-bytes" assert mock_tts.call_args.kwargs["message_ref"] == MessageRef( - "tenant-1", - "app-1", + AppRef("tenant-1", "app-1"), message_id, end_user_id="eu-1", ) diff --git a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py index f24f799464f..f81620d6606 100644 --- a/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py +++ b/api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py @@ -140,15 +140,21 @@ def test_init_rag_pipeline_graph_not_found(mocker, runner): def test_update_document_status_on_failure(mocker, runner): document = MagicMock() + document_ref = MagicMock() session = MagicMock() - session.scalar.return_value = document _patch_create_session(mocker, session) + get_document_by_ref = mocker.patch.object( + module.DatasetRefService, + "get_document_by_ref", + return_value=document, + ) event = GraphRunFailedEvent(error="boom") - runner._update_document_status(event, document_id="doc", dataset_id="ds") + runner._update_document_status(event, document_ref) + get_document_by_ref.assert_called_once_with(document_ref, session=session) assert document.indexing_status == "error" assert document.error == "boom" session.add.assert_called_once_with(document) @@ -157,6 +163,30 @@ def test_update_document_status_on_failure(mocker, runner): session.begin.return_value.__exit__.assert_called_once() +def test_update_document_status_skips_when_document_not_found(mocker, runner): + document_ref = MagicMock() + session = MagicMock() + _patch_create_session(mocker, session) + get_document_by_ref = mocker.patch.object( + module.DatasetRefService, + "get_document_by_ref", + return_value=None, + ) + + runner._update_document_status(GraphRunFailedEvent(error="boom"), document_ref) + + get_document_by_ref.assert_called_once_with(document_ref, session=session) + session.add.assert_not_called() + + +def test_update_document_status_skips_without_document_ref(mocker, runner): + create_session = mocker.patch.object(module, "create_session") + + runner._update_document_status(GraphRunFailedEvent(error="boom"), None) + + create_session.assert_not_called() + + def test_run_pipeline_not_found(mocker: MockerFixture): app_generate_entity = _build_app_generate_entity() app_generate_entity.invoke_from = InvokeFrom.WEB_APP @@ -181,10 +211,108 @@ def test_run_pipeline_not_found(mocker: MockerFixture): runner.run() +def test_run_pipeline_from_other_tenant_is_not_found(mocker: MockerFixture, runner: PipelineRunner): + pipeline = MagicMock(id="pipe", tenant_id="other-tenant") + session = MagicMock() + session.get.side_effect = [None, pipeline] + _patch_create_session(mocker, session) + + with pytest.raises(ValueError, match="Pipeline not found"): + runner.run() + + pipeline.retrieve_dataset.assert_not_called() + + +@pytest.mark.parametrize( + "dataset", + [ + pytest.param(None, id="missing"), + pytest.param(SimpleNamespace(id="ds", tenant_id="other-tenant"), id="other-tenant"), + pytest.param(SimpleNamespace(id="other-dataset", tenant_id="tenant"), id="other-dataset"), + ], +) +def test_run_rejects_unowned_pipeline_dataset( + mocker: MockerFixture, + runner: PipelineRunner, + dataset: SimpleNamespace | None, +): + pipeline = MagicMock(id="pipe", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = dataset + session = MagicMock() + session.get.side_effect = [None, pipeline] + _patch_create_session(mocker, session) + runner.get_workflow = MagicMock() + + with pytest.raises(ValueError, match="Pipeline dataset not found"): + runner.run() + + pipeline.retrieve_dataset.assert_called_once_with(session) + runner.get_workflow.assert_not_called() + + +def test_run_rejects_document_outside_pipeline_dataset_after_async_boundary( + mocker: MockerFixture, + runner: PipelineRunner, +): + runner.application_generate_entity.document_id = "foreign-doc" + runner.application_generate_entity.original_document_id = "foreign-doc" + pipeline = MagicMock(id="pipe", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") + session = MagicMock() + session.get.side_effect = [None, pipeline] + _patch_create_session(mocker, session) + runner.get_workflow = MagicMock() + get_document_by_ref = mocker.patch.object( + module.DatasetRefService, + "get_document_by_ref", + return_value=None, + ) + + with pytest.raises(ValueError, match="Pipeline document not found"): + runner.run() + + document_ref = get_document_by_ref.call_args.args[0] + assert document_ref.dataset.tenant_id == "tenant" + assert document_ref.dataset.dataset_id == "ds" + assert document_ref.document_id == "foreign-doc" + get_document_by_ref.assert_called_once_with(document_ref, session=session) + runner.get_workflow.assert_not_called() + + +def test_run_rejects_original_document_outside_pipeline_dataset_after_async_boundary( + mocker: MockerFixture, + runner: PipelineRunner, +): + runner.application_generate_entity.document_id = "doc" + runner.application_generate_entity.original_document_id = "foreign-doc" + pipeline = MagicMock(id="pipe", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") + session = MagicMock() + session.get.side_effect = [None, pipeline] + _patch_create_session(mocker, session) + runner.get_workflow = MagicMock() + get_document_by_ref = mocker.patch.object( + module.DatasetRefService, + "get_document_by_ref", + side_effect=[SimpleNamespace(id="doc"), None], + ) + + with pytest.raises(ValueError, match="Pipeline original document not found"): + runner.run() + + document_refs = [call.args[0] for call in get_document_by_ref.call_args_list] + assert [document_ref.document_id for document_ref in document_refs] == ["doc", "foreign-doc"] + for document_ref in document_refs: + assert document_ref.dataset.tenant_id == "tenant" + assert document_ref.dataset.dataset_id == "ds" + runner.get_workflow.assert_not_called() + + def test_run_workflow_not_initialized(mocker: MockerFixture): app_generate_entity = _build_app_generate_entity() - pipeline = MagicMock(id="pipe") + pipeline = MagicMock(id="pipe", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") session = MagicMock() session.get.side_effect = [None, pipeline] @@ -209,11 +337,11 @@ def test_run_single_iteration_path(mocker: MockerFixture): app_generate_entity = _build_app_generate_entity() app_generate_entity.single_iteration_run = MagicMock() - pipeline = MagicMock(id="pipe") - end_user = MagicMock(session_id="sess") - + pipeline = MagicMock(id="pipe", tenant_id="tenant") + dataset = SimpleNamespace(id="ds", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = dataset session = MagicMock() - session.get.side_effect = [end_user, pipeline] + session.get.return_value = pipeline _patch_create_session(mocker, session) runner = PipelineRunner( @@ -241,23 +369,41 @@ def test_run_single_iteration_path(mocker: MockerFixture): runner._update_document_status = MagicMock() runner._handle_event = MagicMock() + dataset_ref = MagicMock() + document_ref = MagicMock() + create_dataset_ref = mocker.patch.object( + module.DatasetRefService, + "create_dataset_ref", + return_value=dataset_ref, + ) + create_document_ref_from_id = mocker.patch.object( + module.DatasetRefService, + "create_document_ref_from_id", + return_value=document_ref, + ) + + event = MagicMock() workflow_entry = MagicMock() workflow_entry.graph_engine = MagicMock() - workflow_entry.run.return_value = [MagicMock()] + workflow_entry.run.return_value = [event] mocker.patch.object(module, "WorkflowEntry", return_value=workflow_entry) mocker.patch.object(module, "WorkflowPersistenceLayer", return_value=MagicMock()) runner.run() + create_dataset_ref.assert_called_once_with(dataset) + create_document_ref_from_id.assert_called_once_with(dataset_ref, "doc") runner._prepare_single_node_execution.assert_called_once() + runner._update_document_status.assert_called_once_with(event, document_ref) runner._handle_event.assert_called() def test_run_normal_path_builds_graph(mocker: MockerFixture): app_generate_entity = _build_app_generate_entity() - pipeline = MagicMock(id="pipe") + pipeline = MagicMock(id="pipe", tenant_id="tenant") + pipeline.retrieve_dataset.return_value = SimpleNamespace(id="ds", tenant_id="tenant") end_user = MagicMock(session_id="sess") events = [] diff --git a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py index 59d0fb126d6..ed943fe1129 100644 --- a/api/tests/unit_tests/core/rag/indexing/test_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/test_index_processor.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch from core.rag.index_processor.constant.index_type import IndexTechniqueType from core.rag.index_processor.index_processor import IndexProcessor from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem +from models.dataset import Dataset, Document class TestIndexProcessor: @@ -34,12 +35,13 @@ class TestIndexProcessor: ) dataset = SimpleNamespace( id="dataset-1", + tenant_id="tenant-1", name="Dataset", chunk_structure="text_model", summary_index_setting=None, ) session = MagicMock() - session.scalar.side_effect = [document, dataset, 3] + session.scalar.side_effect = [dataset, document, 3] phase_events: list[str] = [] session.commit.side_effect = lambda: phase_events.append("commit") @@ -59,6 +61,112 @@ class TestIndexProcessor: assert phase_events == ["commit", "index", "commit"] + def test_index_and_clean_scopes_replacement_queries_to_dataset_owner(self) -> None: + dataset = SimpleNamespace( + id="dataset-1", + tenant_id="tenant-1", + name="Dataset", + summary_index_setting=None, + chunk_structure="text_model", + ) + document = SimpleNamespace( + id="doc-1", + tenant_id="tenant-1", + dataset_id="dataset-1", + name="Document", + created_at=datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC), + ) + segment = SimpleNamespace(index_node_id="node-1") + session = MagicMock() + + def resolve_owner(statement): + entity = statement.column_descriptions[0]["entity"] + if entity is Dataset: + return dataset + if entity is Document: + return document + return 3 + + session.scalar.side_effect = resolve_owner + session.scalars.return_value.all.return_value = [segment] + + with patch("core.rag.index_processor.index_processor.IndexProcessorFactory") as index_processor_factory: + index_backend = index_processor_factory.return_value.init_index_processor.return_value + IndexProcessor().index_and_clean( + dataset_id="dataset-1", + document_id="doc-1", + original_document_id="original-doc", + chunks={}, + batch="batch-1", + session=session, + ) + + document_statement = next( + call.args[0] + for call in session.scalar.call_args_list + if call.args[0].column_descriptions[0]["entity"] is Document + ) + segment_statement = session.scalars.call_args.args[0] + delete_statement = session.execute.call_args_list[0].args[0] + word_count_statement = session.scalar.call_args_list[-1].args[0] + segment_update_statement = session.execute.call_args_list[1].args[0] + document_owner = {"doc-1", "dataset-1", "tenant-1"} + original_document_owner = {"original-doc", "dataset-1", "tenant-1"} + + assert document_owner <= set(document_statement.compile().params.values()) + assert original_document_owner <= set(segment_statement.compile().params.values()) + assert original_document_owner <= set(delete_statement.compile().params.values()) + assert document_owner <= set(word_count_statement.compile().params.values()) + assert document_owner <= set(segment_update_statement.compile().params.values()) + index_backend.clean.assert_called_once_with( + dataset, + ["node-1"], + with_keywords=True, + delete_child_chunks=True, + session=session, + ) + index_backend.index.assert_called_once_with(dataset, document, {}, session) + + def test_get_preview_output_scopes_document_to_dataset_owner(self) -> None: + dataset = SimpleNamespace( + id="dataset-1", + tenant_id="tenant-1", + indexing_technique=IndexTechniqueType.ECONOMY, + summary_index_setting=None, + ) + document = SimpleNamespace(doc_language="English") + session = MagicMock() + + def resolve_owner(statement): + entity = statement.column_descriptions[0]["entity"] + if entity is Dataset: + return dataset + if entity is Document: + return document + raise AssertionError(f"Unexpected entity: {entity}") + + session.scalar.side_effect = resolve_owner + processor = IndexProcessor() + expected_preview = MagicMock() + + with patch.object(processor, "format_preview", return_value=expected_preview): + result = processor.get_preview_output( + chunks={}, + dataset_id="dataset-1", + document_id="doc-1", + chunk_structure="text_model", + summary_index_setting=None, + session=session, + ) + + document_statement = next( + call.args[0] + for call in session.scalar.call_args_list + if call.args[0].column_descriptions[0]["entity"] is Document + ) + assert {"doc-1", "dataset-1", "tenant-1"} <= set(document_statement.compile().params.values()) + assert result is expected_preview + def test_preview_summary_workers_use_independent_sessions(self) -> None: caller_session = MagicMock() phase_events: list[str] = [] diff --git a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py index d35ffcdb525..1ad4e070edd 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py @@ -5,8 +5,9 @@ import pytest from pytest_mock import MockerFixture from core.app.entities.app_invoke_entities import InvokeFrom -from models.dataset import Pipeline +from models.dataset import Dataset, Pipeline from models.model import Account, App, EndUser +from services.dataset_ref_service import DatasetRefService from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService @@ -58,7 +59,15 @@ def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_err def test_generate_updates_document_status_and_returns_event_stream(mocker: MockerFixture) -> None: - pipeline = cast(Pipeline, SimpleNamespace(id="pipeline-1")) + dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) + pipeline = cast( + Pipeline, + SimpleNamespace( + id="pipeline-1", + tenant_id="tenant-1", + retrieve_dataset=mocker.Mock(return_value=dataset), + ), + ) user = cast(Account | EndUser, SimpleNamespace(id="user-1")) args = {"original_document_id": "doc-1", "query": "hello"} session_mock = mocker.Mock() @@ -81,30 +90,124 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke ) assert result == "stream-events" - update_status_mock.assert_called_once_with("doc-1", session=session_mock) + document_ref = update_status_mock.call_args.args[0] + assert document_ref.dataset.tenant_id == "tenant-1" + assert document_ref.dataset.dataset_id == "dataset-1" + assert document_ref.document_id == "doc-1" + update_status_mock.assert_called_once_with(document_ref, session=session_mock) assert generator_instance.generate.call_args.kwargs["session"] is session_mock +def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture) -> None: + dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-2")) + pipeline = cast( + Pipeline, + SimpleNamespace( + id="pipeline-1", + tenant_id="tenant-1", + retrieve_dataset=mocker.Mock(return_value=dataset), + ), + ) + mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) + update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status") + + with pytest.raises(ValueError, match="Pipeline dataset is required"): + PipelineGenerateService.generate( + pipeline=pipeline, + user=cast(Account, SimpleNamespace(id="user-1")), + args={"original_document_id": "doc-1"}, + invoke_from=InvokeFrom.WEB_APP, + session=mocker.Mock(), + ) + + update_status_mock.assert_not_called() + + +def test_generate_rejects_original_document_outside_pipeline_dataset_before_dispatch( + mocker: MockerFixture, +) -> None: + dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) + pipeline = cast( + Pipeline, + SimpleNamespace( + id="pipeline-1", + tenant_id="tenant-1", + retrieve_dataset=mocker.Mock(return_value=dataset), + ), + ) + session = mocker.Mock() + session.scalar.return_value = None + mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) + generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") + + with pytest.raises(ValueError, match="Pipeline document not found"): + PipelineGenerateService.generate( + pipeline=pipeline, + user=cast(Account, SimpleNamespace(id="user-1")), + args={"original_document_id": "foreign-doc"}, + invoke_from=InvokeFrom.PUBLISHED_PIPELINE, + session=session, + ) + + statement = session.scalar.call_args.args[0] + assert {"foreign-doc", "dataset-1", "tenant-1"} <= set(statement.compile().params.values()) + generator_cls.assert_not_called() + + def test_update_document_status_updates_existing_document(mocker: MockerFixture) -> None: document = SimpleNamespace(indexing_status="completed") + dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) + dataset_ref = DatasetRefService.create_dataset_ref(dataset) + document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") session_mock = mocker.Mock() - session_mock.get.return_value = document + get_document_mock = mocker.patch.object(DatasetRefService, "get_document_by_ref", return_value=document) add_mock = session_mock.add - PipelineGenerateService.update_document_status("doc-1", session=session_mock) + PipelineGenerateService.update_document_status(document_ref, session=session_mock) assert document.indexing_status == "waiting" + get_document_mock.assert_called_once_with(document_ref, session=session_mock) add_mock.assert_called_once_with(document) -def test_update_document_status_skips_when_document_missing(mocker: MockerFixture) -> None: +@pytest.mark.parametrize( + ("document_tenant_id", "document_dataset_id"), + [ + pytest.param("other-tenant", "dataset-1", id="other-tenant"), + pytest.param("tenant-1", "other-dataset", id="other-dataset"), + ], +) +def test_update_document_status_rejects_document_outside_owner( + mocker: MockerFixture, + document_tenant_id: str, + document_dataset_id: str, +) -> None: + dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) + dataset_ref = DatasetRefService.create_dataset_ref(dataset) + document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") + outside_document = SimpleNamespace( + id="doc-1", + tenant_id=document_tenant_id, + dataset_id=document_dataset_id, + indexing_status="completed", + ) session_mock = mocker.Mock() - session_mock.get.return_value = None + + def resolve_document(statement): + params = set(statement.compile().params.values()) + outside_owner = {outside_document.id, outside_document.dataset_id, outside_document.tenant_id} + return outside_document if outside_owner <= params else None + + session_mock.scalar.side_effect = resolve_document add_mock = session_mock.add - PipelineGenerateService.update_document_status("missing", session=session_mock) + with pytest.raises(ValueError, match="Pipeline document not found"): + PipelineGenerateService.update_document_status(document_ref, session=session_mock) + statement = session_mock.scalar.call_args.args[0] + assert {"doc-1", "dataset-1", "tenant-1"} <= set(statement.compile().params.values()) + assert outside_document.indexing_status == "completed" add_mock.assert_not_called() diff --git a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py index 47f16c76b87..32908e4cb09 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py @@ -8,9 +8,14 @@ from unittest.mock import Mock import pytest from pytest_mock import MockerFixture +from core.app.entities.app_invoke_entities import InvokeFrom +from graphon.enums import WorkflowNodeExecutionStatus +from graphon.graph_events import NodeRunFailedEvent +from graphon.node_events.base import NodeRunResult from models import Account, Tenant from models.dataset import Dataset, Pipeline, PipelineCustomizedTemplate, PipelineRecommendedPlugin from models.workflow import Workflow +from services.dataset_ref_service import DatasetRefService from services.entities.knowledge_entities.rag_pipeline_entities import IconInfo, PipelineTemplateInfoEntity from services.rag_pipeline.rag_pipeline import RagPipelineService from services.workflow_ref_service import WorkflowRef @@ -122,6 +127,45 @@ def _make_dataset(*, dataset_id: str = "d1", pipeline_id: str = "p1", tenant_id: return dataset +def _make_failed_published_node_run() -> tuple[SimpleNamespace, NodeRunFailedEvent]: + class FakeVariablePool: + def __init__(self) -> None: + self._values = { + ("sys", "invoke_from"): SimpleNamespace(value=InvokeFrom.PUBLISHED_PIPELINE), + ("sys", "app_id"): SimpleNamespace(value="pipeline-1"), + ("sys", "dataset_id"): SimpleNamespace(value="dataset-1"), + ("sys", "document_id"): SimpleNamespace(value="doc-1"), + } + + def get(self, path: list[str]) -> SimpleNamespace | None: + return self._values.get(tuple(path)) + + node_instance = SimpleNamespace( + workflow_id="wf-1", + node_type="start", + title="Start", + graph_runtime_state=SimpleNamespace(variable_pool=FakeVariablePool()), + error_strategy=None, + ) + run_result = NodeRunResult( + status=WorkflowNodeExecutionStatus.FAILED, + error="boom", + error_type="runtime", + inputs={}, + outputs={}, + ) + event = NodeRunFailedEvent( + id="evt-1", + start_at=time.time(), + node_id="node-1", + node_type="start", + node_run_result=run_result, + error="boom", + route_node_id=None, + ) + return node_instance, event + + def _make_customized_template() -> PipelineCustomizedTemplate: return PipelineCustomizedTemplate( tenant_id="t1", @@ -1680,54 +1724,15 @@ def test_handle_node_run_result_raises_when_no_terminal_event( def test_handle_node_run_result_marks_document_error_for_published_invoke( mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext ) -> None: - from core.app.entities.app_invoke_entities import InvokeFrom - from graphon.enums import WorkflowNodeExecutionStatus - from graphon.graph_events import NodeRunFailedEvent - from graphon.node_events.base import NodeRunResult - - class FakeVariablePool: - def __init__(self): - self._values = { - ("sys", "invoke_from"): SimpleNamespace(value=InvokeFrom.PUBLISHED_PIPELINE), - ("sys", "app_id"): SimpleNamespace(value="pipeline-1"), - ("sys", "dataset_id"): SimpleNamespace(value="dataset-1"), - ("sys", "document_id"): SimpleNamespace(value="doc-1"), - } - - def get(self, path): - return self._values.get(tuple(path)) - - node_instance = SimpleNamespace( - workflow_id="wf-1", - node_type="start", - title="Start", - graph_runtime_state=SimpleNamespace(variable_pool=FakeVariablePool()), - error_strategy=None, - ) - run_result = NodeRunResult( - status=WorkflowNodeExecutionStatus.FAILED, - error="boom", - error_type="runtime", - inputs={}, - outputs={}, - ) - - def _event_generator(): - yield NodeRunFailedEvent( - id="evt-1", - start_at=time.time(), - node_id="node-1", - node_type="start", - node_run_result=run_result, - error="boom", - route_node_id=None, - ) - + node_instance, event = _make_failed_published_node_run() document = SimpleNamespace(indexing_status="waiting", error=None) - rag_pipeline_service.session.scalar.return_value = document + rag_pipeline_service.session.scalar.return_value = _make_dataset( + dataset_id="dataset-1", pipeline_id="pipeline-1", tenant_id="t1" + ) + get_document_by_ref = mocker.patch.object(DatasetRefService, "get_document_by_ref", return_value=document) result = rag_pipeline_service.service._handle_node_run_result( - getter=lambda: (node_instance, _event_generator()), + getter=lambda: (node_instance, iter([event])), start_at=time.perf_counter(), tenant_id="t1", node_id="node-1", @@ -1737,17 +1742,41 @@ def test_handle_node_run_result_marks_document_error_for_published_invoke( stmt = rag_pipeline_service.session.scalar.call_args.args[0] compiled = stmt.compile() statement = str(compiled) - assert "documents.id" in statement - assert "documents.tenant_id" in statement - assert "documents.dataset_id" in statement + assert "datasets.id" in statement assert "datasets.tenant_id" in statement assert "datasets.pipeline_id" in statement - assert "doc-1" in compiled.params.values() assert "t1" in compiled.params.values() assert "dataset-1" in compiled.params.values() assert "pipeline-1" in compiled.params.values() + document_ref = get_document_by_ref.call_args.args[0] + assert document_ref.dataset.tenant_id == "t1" + assert document_ref.dataset.dataset_id == "dataset-1" + assert document_ref.document_id == "doc-1" + get_document_by_ref.assert_called_once_with(document_ref, session=rag_pipeline_service.session) assert document.indexing_status == "error" assert document.error == "boom" + rag_pipeline_service.session.add.assert_called_once_with(document) + rag_pipeline_service.session.commit.assert_called_once_with() + + +def test_handle_node_run_result_does_not_write_when_pipeline_dataset_is_missing( + mocker: MockerFixture, rag_pipeline_service: RagPipelineServiceTestContext +) -> None: + node_instance, event = _make_failed_published_node_run() + rag_pipeline_service.session.scalar.return_value = None + get_document_by_ref = mocker.patch.object(DatasetRefService, "get_document_by_ref") + + result = rag_pipeline_service.service._handle_node_run_result( + getter=lambda: (node_instance, iter([event])), + start_at=time.perf_counter(), + tenant_id="t1", + node_id="node-1", + ) + + assert result.status == WorkflowNodeExecutionStatus.FAILED + get_document_by_ref.assert_not_called() + rag_pipeline_service.session.add.assert_not_called() + rag_pipeline_service.session.commit.assert_not_called() def test_run_datasource_node_preview_raises_for_unsupported_provider( diff --git a/api/tests/unit_tests/services/test_annotation_service.py b/api/tests/unit_tests/services/test_annotation_service.py index 87863697cfb..d729b1c5a80 100644 --- a/api/tests/unit_tests/services/test_annotation_service.py +++ b/api/tests/unit_tests/services/test_annotation_service.py @@ -32,7 +32,7 @@ def _make_app_ref(app: MagicMock) -> AppRef: def _make_annotation_ref(app: MagicMock, annotation_id: str = "ann-1") -> AnnotationRef: - return AnnotationRef(tenant_id=app.tenant_id, app_id=app.id, annotation_id=annotation_id) + return AnnotationRef(app=AppRef(tenant_id=app.tenant_id, app_id=app.id), annotation_id=annotation_id) def _make_user(user_id: str = "user-1") -> MagicMock: @@ -82,6 +82,7 @@ def _assert_statement_binds_annotation(stmt: Any, annotation_id: str, app_id: st statement = str(compiled) assert "message_annotations.id" in statement assert "message_annotations.app_id" in statement + assert "JOIN apps" not in statement assert annotation_id in compiled.params.values() assert app_id in compiled.params.values() @@ -731,6 +732,7 @@ class TestAppAnnotationServiceDirectManipulation: statement = str(compiled) assert "message_annotations.id IN" in statement assert "message_annotations.app_id" in statement + assert "JOIN apps" not in statement assert ["ann-1", "ann-2"] in compiled.params.values() assert app.id in compiled.params.values() mock_task.delay.assert_called_once_with(annotation1.id, app.id, tenant_id, setting.collection_binding_id) diff --git a/api/tests/unit_tests/services/test_audio_service.py b/api/tests/unit_tests/services/test_audio_service.py index b9f44bc212b..0722faad651 100644 --- a/api/tests/unit_tests/services/test_audio_service.py +++ b/api/tests/unit_tests/services/test_audio_service.py @@ -65,7 +65,7 @@ from models.agent_config_entities import AgentSoulConfig from models.enums import ConversationFromSource, MessageStatus from models.model import App, AppMode, AppModelConfig, Message from models.workflow import Workflow -from services.app_ref_service import MessageRef +from services.app_ref_service import AppRef, MessageRef from services.audio_service import AudioService from services.errors.audio import ( AudioTooLargeServiceError, @@ -661,8 +661,7 @@ class TestAudioServiceTTS: # Arrange app = factory.create_app_mock(app_id=APP_ID, tenant_id=TENANT_ID, mode=AppMode.CHAT) message_ref = MessageRef( - tenant_id=TENANT_ID, - app_id=APP_ID, + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), message_id=MESSAGE_ID, end_user_id=END_USER_ID, account_id=ACCOUNT_ID, @@ -678,22 +677,19 @@ class TestAudioServiceTTS: # Act for wrong_ref in ( MessageRef( - tenant_id=TENANT_ID, - app_id=OTHER_ID, + app=AppRef(tenant_id=TENANT_ID, app_id=OTHER_ID), message_id=MESSAGE_ID, end_user_id=END_USER_ID, account_id=ACCOUNT_ID, ), MessageRef( - tenant_id=TENANT_ID, - app_id=APP_ID, + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), message_id=MESSAGE_ID, end_user_id=OTHER_ID, account_id=ACCOUNT_ID, ), MessageRef( - tenant_id=TENANT_ID, - app_id=APP_ID, + app=AppRef(tenant_id=TENANT_ID, app_id=APP_ID), message_id=MESSAGE_ID, end_user_id=END_USER_ID, account_id=OTHER_ID, diff --git a/api/tests/unit_tests/services/test_dataset_service_document.py b/api/tests/unit_tests/services/test_dataset_service_document.py index 0c9969815e5..26a3ac08d5f 100644 --- a/api/tests/unit_tests/services/test_dataset_service_document.py +++ b/api/tests/unit_tests/services/test_dataset_service_document.py @@ -1,6 +1,6 @@ """Unit tests for DocumentService behaviors in dataset_service.""" -from services.dataset_ref_service import DatasetRef +from services.dataset_ref_service import DatasetRefService from .dataset_service_test_helpers import ( Account, @@ -130,7 +130,7 @@ class TestDocumentServiceMutations: ): session.scalars.return_value.all.return_value = [document] - dataset_ref = DatasetRef(tenant_id=dataset.tenant_id, dataset_id=dataset.id) + dataset_ref = DatasetRefService.create_dataset_ref(dataset) DocumentService.delete_documents( dataset_ref, ["doc-1", "other-doc"], diff --git a/api/tests/unit_tests/services/test_dataset_service_segment.py b/api/tests/unit_tests/services/test_dataset_service_segment.py index b53cd152f07..3e229d9a787 100644 --- a/api/tests/unit_tests/services/test_dataset_service_segment.py +++ b/api/tests/unit_tests/services/test_dataset_service_segment.py @@ -1,6 +1,6 @@ """Unit tests for SegmentService behaviors in dataset_service.""" -from services.dataset_ref_service import DatasetRef, DatasetRefService +from services.dataset_ref_service import DatasetRef, DatasetRefService, DocumentRef, SegmentRef from .dataset_service_test_helpers import ( Account, @@ -39,15 +39,24 @@ class TestDatasetRefService: """Unit tests for typed dataset resource refs.""" def test_dataset_ref_is_plain_named_tuple(self): - dataset_ref = DatasetRef("tenant-1", "dataset-1") + 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") - def test_create_document_ref_rejects_document_outside_dataset(self): + @pytest.mark.parametrize( + ("document_dataset_id", "document_tenant_id"), + [("other-dataset", "tenant-1"), ("dataset-1", "other-tenant")], + ) + def test_create_document_ref_rejects_document_outside_dataset(self, document_dataset_id, document_tenant_id): dataset = _make_dataset(dataset_id="dataset-1", tenant_id="tenant-1") - document = _make_document(document_id="doc-1", dataset_id="other-dataset", tenant_id="tenant-1") + document = _make_document( + document_id="doc-1", + dataset_id=document_dataset_id, + tenant_id=document_tenant_id, + ) dataset_ref = DatasetRefService.create_dataset_ref(dataset) assert DatasetRefService.create_document_ref(dataset_ref, document) is None @@ -55,11 +64,30 @@ class TestDatasetRefService: def test_create_segment_ref_carries_full_parent_chain(self): segment_ref = _make_segment_ref() - assert segment_ref.tenant_id == "tenant-1" - assert segment_ref.dataset_id == "dataset-1" - assert segment_ref.document_id == "doc-1" + assert isinstance(segment_ref, SegmentRef) + assert isinstance(segment_ref.document, DocumentRef) + assert isinstance(segment_ref.document.dataset, DatasetRef) + assert segment_ref.document.dataset.tenant_id == "tenant-1" + assert segment_ref.document.dataset.dataset_id == "dataset-1" + assert segment_ref.document.document_id == "doc-1" assert segment_ref.segment_id == "segment-1" + def test_get_document_by_ref_uses_full_ownership_chain(self): + dataset_ref = DatasetRefService.create_dataset_ref(_make_dataset()) + document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") + document = _make_document() + session = MagicMock() + session.scalar.return_value = document + + result = DatasetRefService.get_document_by_ref(document_ref, session=session) + + assert result is document + stmt = session.scalar.call_args.args[0] + sql = str(stmt.compile(compile_kwargs={"literal_binds": True})) + assert "documents.id = 'doc-1'" in sql + assert "documents.dataset_id = 'dataset-1'" in sql + assert "documents.tenant_id = 'tenant-1'" in sql + class TestSegmentServiceChildChunks: """Unit tests for child-chunk CRUD helpers."""