diff --git a/api/services/knowledge/segments/application.py b/api/services/knowledge/segments/application.py index 6b8ce7c5d5a83a..485034678a0cd5 100644 --- a/api/services/knowledge/segments/application.py +++ b/api/services/knowledge/segments/application.py @@ -371,29 +371,38 @@ def change_segment_status( segment_ids: Sequence[str], action: Literal["enable", "disable"], ) -> None: + """Apply the requested status and dispatch index reconciliation. + + Requires an accessible document that is not being indexed. Segments + currently being indexed are excluded. A matching persisted status does + not prove that a previous dispatch succeeded, so explicit retries must + dispatch again. On dispatch failure the desired status remains persisted + and the caller receives SegmentStatusUpdateError and can retry. + """ scope = self._require_scope(context, dataset_id=dataset_id, document_id=document_id) self._check_dataset_model(scope.dataset) self._check_embedding_model(scope.dataset) if self._indexing_state.is_document_indexing(scope.document.id): raise SegmentDocumentIndexingError("Document is being indexed, please try again later") try: - segments = self._store.get_segments(scope.document.ref, segment_ids) - ids = [] - for segment in segments: - if segment.enabled == (action == "enable") or self._indexing_state.is_segment_indexing(segment.id): - continue - ids.append(segment.id) - if ids: + segments = [ + segment + for segment in self._store.get_segments(scope.document.ref, segment_ids) + if not self._indexing_state.is_segment_indexing(segment.id) + ] + changed_ids = [segment.id for segment in segments if segment.enabled != (action == "enable")] + if changed_ids: self._store.set_segments( scope.document.ref, - ids, + changed_ids, { "enabled": action == "enable", "disabled_at": naive_utc_now() if action == "disable" else None, "disabled_by": context.account_id if action == "disable" else None, }, ) - self._index.change_status(scope.document.ref, ids, action) + if segments: + self._index.change_status(scope.document.ref, [segment.id for segment in segments], action) except Exception as error: raise SegmentStatusUpdateError(str(error)) from error diff --git a/api/tests/unit_tests/services/knowledge/segments/test_adapters.py b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py index 67e457c65a2b49..c4fd4240604a25 100644 --- a/api/tests/unit_tests/services/knowledge/segments/test_adapters.py +++ b/api/tests/unit_tests/services/knowledge/segments/test_adapters.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from datetime import datetime from pathlib import Path -from typing import Protocol +from typing import Literal, Protocol from unittest.mock import create_autospec, patch import pytest @@ -53,6 +53,7 @@ SegmentListFilter, SegmentModelGuard, SegmentModelProviderError, + SegmentStatusUpdateError, ) @@ -454,6 +455,7 @@ def lock(self, name: str, timeout: int) -> nullcontext[None]: class VectorProbe: active_transactions: set[Connection] fail: bool = False + dispatch_unavailable: bool = False writes: int = 0 deleted_ids: list[str] = field(default_factory=list) tasks: list[tuple[str, tuple[object, ...]]] = field(default_factory=list) @@ -530,6 +532,8 @@ def get_model_instance(self, **kwargs: object) -> TokenModel: def dispatch(kind: str, *args: object) -> None: assert not active + if probe.dispatch_unavailable: + raise RuntimeError("broker unavailable") probe.tasks.append((kind, args)) gateway = indexing_module.SegmentIndexingGateway( @@ -981,6 +985,78 @@ def test_child_query_escapes_wildcards_and_enforces_complete_scope( assert [child.content for child in page.items] == ["50%_literal"] +@pytest.mark.parametrize("action", ["enable", "disable"]) +@pytest.mark.parametrize("mixed_batch", [False, True], ids=["single", "mixed-batch"]) +def test_status_retry_after_dispatch_failure_reaches_index( + indexing_probe: IndexingProbe, + sqlite_session_factory: sessionmaker[Session], + action: Literal["enable", "disable"], + mixed_batch: bool, +) -> None: + """Retry the same scoped batch after failed dispatch without rewriting existing audit metadata.""" + app, _, probe = indexing_probe + context = RequestContext("request", None, "author", "workspace-1") + requested_ids = ["segment-1", "foreign"] + expected_ids = {"segment-1"} + previous_disabled_at = datetime(2025, 1, 1) if action == "disable" else None + previous_disabled_by = "previous-editor" if action == "disable" else None + with sqlite_session_factory.begin() as session: + segment = _segment("segment-1", "workspace-1", "dataset-1", "document-1") + segment.enabled = action != "enable" + session.add(segment) + session.add(_segment("foreign", "workspace-2", "dataset-1", "document-1")) + if mixed_batch: + already_target = _segment("already-target", "workspace-1", "dataset-1", "document-1") + already_target.enabled = action == "enable" + already_target.disabled_at = previous_disabled_at + already_target.disabled_by = previous_disabled_by + session.add(already_target) + requested_ids.append(already_target.id) + expected_ids.add(already_target.id) + + probe.dispatch_unavailable = True + with pytest.raises(SegmentStatusUpdateError, match="broker unavailable"): + app.change_segment_status( + context, + dataset_id="dataset-1", + document_id="document-1", + segment_ids=requested_ids, + action=action, + ) + with sqlite_session_factory() as session: + segment = session.get(DocumentSegment, "segment-1") + assert segment is not None + assert segment.enabled == (action == "enable") + disabled_at = segment.disabled_at + disabled_by = segment.disabled_by + assert probe.tasks == [] + + probe.dispatch_unavailable = False + app.change_segment_status( + context, dataset_id="dataset-1", document_id="document-1", segment_ids=requested_ids, action=action + ) + assert len(probe.tasks) == 1 + kind, task_args = probe.tasks[0] + assert kind == action + assert isinstance(task_args[0], list) + assert len(task_args[0]) == len(expected_ids) + assert set(task_args[0]) == expected_ids + assert task_args[1:] == ("dataset-1", "document-1") + with sqlite_session_factory() as session: + segment = session.get(DocumentSegment, "segment-1") + assert segment is not None + assert (segment.disabled_at, segment.disabled_by) == (disabled_at, disabled_by) + + if mixed_batch: + already_target = session.get(DocumentSegment, "already-target") + assert already_target is not None + assert already_target.enabled == (action == "enable") + assert (already_target.disabled_at, already_target.disabled_by) == ( + previous_disabled_at, + previous_disabled_by, + ) + + def test_batch_status_and_delete_only_modify_scoped_segments( indexing_probe: IndexingProbe, sqlite_session_factory: sessionmaker[Session] ) -> None: diff --git a/api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py b/api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py index 8653d53244602d..dc084342722390 100644 --- a/api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py +++ b/api/tests/unit_tests/tasks/test_enable_segments_to_index_task.py @@ -234,3 +234,25 @@ def fail_load(*_args: object, **_kwargs: object) -> None: assert persisted_segment.disabled_at is not None assert phase_events == ["load", "rollback", "commit"] enable_summaries.assert_not_called() + + +@pytest.mark.parametrize("document_state", ["disabled", "archived", "incomplete"]) +def test_repeated_batch_enable_preserves_document_eligibility( + sqlite_session: Session, + indexed_segment: tuple[Dataset, Document, DocumentSegment], + document_state: str, +) -> None: + """Repeated tasks for an ineligible document leave its completed segment untouched.""" + dataset, document, segment = indexed_segment + document.enabled = document_state != "disabled" + document.archived = document_state == "archived" + document.indexing_status = IndexingStatus.INDEXING if document_state == "incomplete" else IndexingStatus.COMPLETED + sqlite_session.commit() + + for _ in range(2): + enable_segments_to_index_task.run([segment.id], dataset.id, document.id) + + sqlite_session.refresh(segment) + assert segment.status == SegmentStatus.COMPLETED + assert segment.error is None + assert segment.index_node_id == "node-1"