opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-20 09:38:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): bind scoped lookups to nested owner refs (#38470)

Browse Source
This commit is contained in:
WH-2099 2026-07-15 09:31:09 +00:00 committed by GitHub
parent b721e7a32d
commit 6a511da325
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
28 changed files with 708 additions and 215 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
api/AGENTS.md
View File

@ -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:

4
api/controllers/console/app/mcp_server.py
View File

@ -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)
)

63
api/core/app/apps/pipeline/pipeline_runner.py
View File

@ -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)

48
api/core/rag/index_processor/index_processor.py
View File

@ -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

23
api/services/annotation_service.py
View File

@ -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())

18
api/services/app_ref_service.py
View File

@ -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)

5
api/services/audio_service.py
View File

@ -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:

43
api/services/dataset_ref_service.py
View File

@ -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,
)
)

12
api/services/dataset_service.py
View File

@ -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)
)

27
api/services/rag_pipeline/pipeline_generate_service.py
View File

@ -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)

24
api/services/rag_pipeline/rag_pipeline.py
View File

@ -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

7
api/tests/test_containers_integration_tests/services/test_annotation_service.py
View File

@ -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,
)

18
api/tests/test_containers_integration_tests/services/test_audio_service_db.py
View File

@ -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,
),

6
api/tests/test_containers_integration_tests/services/test_dataset_service_document.py
View File

@ -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(

6
api/tests/unit_tests/controllers/console/app/test_annotation_api.py
View File

@ -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)

4
api/tests/unit_tests/controllers/console/app/test_audio.py
View File

@ -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:

5
api/tests/unit_tests/controllers/console/explore/test_audio.py
View File

@ -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",
)

5
api/tests/unit_tests/controllers/console/explore/test_trial.py
View File

@ -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",
)

4
api/tests/unit_tests/controllers/service_api/app/test_audio.py
View File

@ -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(

5
api/tests/unit_tests/controllers/web/test_audio.py
View File

@ -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",
)

164
api/tests/unit_tests/core/app/apps/pipeline/test_pipeline_runner.py
View File

@ -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 = []

110
api/tests/unit_tests/core/rag/indexing/test_index_processor.py
View File

@ -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] = []

119
api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py
View File

@ -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()

127
api/tests/unit_tests/services/rag_pipeline/test_rag_pipeline_service.py
View File

@ -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(

4
api/tests/unit_tests/services/test_annotation_service.py
View File

@ -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)

14
api/tests/unit_tests/services/test_audio_service.py
View File

@ -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,

4
api/tests/unit_tests/services/test_dataset_service_document.py
View File

@ -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"],

42
api/tests/unit_tests/services/test_dataset_service_segment.py
View File

@ -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."""