From 43359ac0a94e4061ac6e1ea936f62b284fbd70e8 Mon Sep 17 00:00:00 2001 From: TINGyu123644 <2663180454@qq.com> Date: Sun, 1 Nov 2026 18:44:30 +0800 Subject: [PATCH 1/2] fix(segments): reconcile segment indexes on retried status changes Restores the original persist-then-schedule ordering in change_segment_status after review feedback: scheduling the index task before the DB write introduced a reverse partial-failure window where a successful publish followed by a failed commit could leave the retrieval index out of sync with the persisted state. The fix for the original issue (#45238) is kept via reconciliation: a persisted status that already matches the request does not prove that a previous dispatch succeeded, so explicit retries dispatch the index change again for every non-indexing segment instead of being skipped as a no-op. The enable/disable tasks select segments by id and are idempotent, so repeated dispatches are safe. New tests cover the dispatch-failure-after-commit path (state stays persisted, error translated) and the reconciliation retry (no second write, index dispatch repeated); the ordering test now asserts persist-then-dispatch. --- .../knowledge/segments/application.py | 24 +++- .../knowledge/segments/test_application.py | 109 ++++++++++++++++++ 2 files changed, 127 insertions(+), 6 deletions(-) diff --git a/api/services/knowledge/segments/application.py b/api/services/knowledge/segments/application.py index 6b8ce7c5d5a83a..5d848ab03d39a2 100644 --- a/api/services/knowledge/segments/application.py +++ b/api/services/knowledge/segments/application.py @@ -378,22 +378,34 @@ def change_segment_status( raise SegmentDocumentIndexingError("Document is being indexed, please try again later") try: segments = self._store.get_segments(scope.document.ref, segment_ids) - ids = [] + changed_ids = [] + matching_ids = [] for segment in segments: - if segment.enabled == (action == "enable") or self._indexing_state.is_segment_indexing(segment.id): + if self._indexing_state.is_segment_indexing(segment.id): continue - ids.append(segment.id) - if ids: + if segment.enabled != (action == "enable"): + changed_ids.append(segment.id) + else: + matching_ids.append(segment.id) + 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 changed_ids or matching_ids: + # A persisted status that already matches the request does not prove that a + # previous dispatch succeeded: when the DB write commits and the task + # publication then fails, the state is stored while the segment is still + # present in (or missing from) the retrieval index. Explicit retries therefore + # dispatch again for every non-indexing segment instead of being skipped as a + # no-op; the enable/disable tasks select segments by id and are idempotent, so + # repeated dispatches are safe. + self._index.change_status(scope.document.ref, changed_ids + matching_ids, action) except Exception as error: raise SegmentStatusUpdateError(str(error)) from error diff --git a/api/tests/unit_tests/services/knowledge/segments/test_application.py b/api/tests/unit_tests/services/knowledge/segments/test_application.py index b671d228d1f64f..05813199ad3999 100644 --- a/api/tests/unit_tests/services/knowledge/segments/test_application.py +++ b/api/tests/unit_tests/services/knowledge/segments/test_application.py @@ -32,6 +32,7 @@ SegmentEmbeddingModelUnavailableError, SegmentIndex, SegmentIndexingState, + SegmentIndexTarget, SegmentInvalidFileTypeError, SegmentListFilter, SegmentModelGuard, @@ -134,6 +135,32 @@ def _service(scope: SegmentScope | None = None, *, uploads: SegmentUploadCatalog return service, dataset_access, store, scopes, model_guard, indexing_state, dispatcher +def _service_with_index() -> tuple[DatasetSegmentApplicationService, Mock, Mock]: + dataset_access = create_autospec(DatasetAccess, instance=True, spec_set=True) + dataset_access.check_access.return_value = AccessibleDataset(id="dataset-1", workspace_id="workspace-1") + store = create_autospec(SegmentStore, instance=True, spec_set=True) + scopes = create_autospec(SegmentScopeReader, instance=True, spec_set=True) + scopes.get_segment_scope.return_value = _scope() + indexing_state = create_autospec(SegmentIndexingState, instance=True, spec_set=True) + indexing_state.is_document_indexing.return_value = False + indexing_state.is_segment_indexing.return_value = False + index = create_autospec(SegmentIndex, instance=True, spec_set=True) + service = DatasetSegmentApplicationService( + dataset_access=dataset_access, + store=store, + limits=SegmentLimitsStub(), + text_hash=generate_text_hash, + index=index, + uploads=create_autospec(SegmentUploadCatalog, instance=True, spec_set=True), + scopes=scopes, + model_guard=create_autospec(SegmentModelGuard, instance=True, spec_set=True), + indexing_state=indexing_state, + batch_dispatcher=create_autospec(SegmentBatchImportDispatcher, instance=True, spec_set=True), + job_id_factory=lambda: "job-1", + ) + return service, store, index + + def test_list_segments_uses_complete_owner_scope() -> None: service, dataset_access, store, scopes, _, _, _ = _service() expected = SegmentPage(items=(), total=0, total_pages=0, page=1, limit=20) @@ -277,6 +304,88 @@ def test_status_change_translates_operation_failure() -> None: ) +def test_status_change_persists_state_before_dispatching() -> None: + service, store, index = _service_with_index() + store.get_segments.return_value = (SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=True),) + calls: list[str] = [] + index.change_status.side_effect = lambda *args, **kwargs: calls.append("dispatch") + store.set_segments.side_effect = lambda *args, **kwargs: calls.append("persist") + + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="disable", + ) + + assert calls == ["persist", "dispatch"] + store.set_segments.assert_called_once() + persisted_ids, persisted_values = store.set_segments.call_args[0][1:3] + assert persisted_ids == ["segment-1"] + assert persisted_values["enabled"] is False + index.change_status.assert_called_once() + + +def test_status_change_persists_state_when_dispatch_fails() -> None: + """With the restored persist-then-schedule order the remaining failure window is a + dispatch failure after the state was committed: the desired status stays persisted + and the error is translated for the caller, so an explicit retry can reconcile.""" + service, store, index = _service_with_index() + store.get_segments.return_value = (SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=False),) + index.change_status.side_effect = RuntimeError("broker unavailable") + + with pytest.raises(SegmentStatusUpdateError, match="broker unavailable"): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="enable", + ) + + store.set_segments.assert_called_once() + persisted_ids, persisted_values = store.set_segments.call_args[0][1:3] + assert persisted_ids == ["segment-1"] + assert persisted_values["enabled"] is True + index.change_status.assert_called_once() + + +def test_status_change_retry_after_dispatch_failure_reconciles_index() -> None: + """A retry whose persisted status already matches still re-dispatches the index + change so a previous failed dispatch cannot leave DB state and the retrieval + index permanently out of sync.""" + service, store, index = _service_with_index() + store.get_segments.return_value = (SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=False),) + index.change_status.side_effect = RuntimeError("broker unavailable") + + with pytest.raises(SegmentStatusUpdateError): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="enable", + ) + assert store.set_segments.called + assert index.change_status.call_count == 1 + + index.change_status.side_effect = None + store.get_segments.return_value = (SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=True),) + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="enable", + ) + + assert index.change_status.call_count == 2 + assert index.change_status.call_args[0][1] == ["segment-1"] + assert index.change_status.call_args[0][2] == "enable" + assert store.set_segments.call_count == 1 + + def test_update_and_delete_segment_map_missing_segment() -> None: service, _, store, scopes, _, _, _ = _service() store.get_segment_update_state.return_value = None From 84b557e5ebed1654588bf1a445e73f99b838a40d Mon Sep 17 00:00:00 2001 From: TINGyu123644 <2663180454@qq.com> Date: Sun, 1 Nov 2026 20:25:14 +0800 Subject: [PATCH 2/2] test(segments): broaden regression coverage for retried status changes Adds the regression coverage suggested in review on top of the reconciliation fix:/n/n- parametrized enable/disable x single/mixed-batch matrix asserting that persistence covers only actual status changes while the dispatch covers every non-indexing segment in scope, with out-of-scope ids dropped; - broker-unavailable injection for a batch whose segments already match the requested status, asserting the failure is translated instead of being swallowed as a no-op; - end-to-end retry over a real SQLite session through the SQLAlchemy repository: a failed dispatch leaves the desired status persisted, the retry re-enqueues the index task for every non-indexing segment and never rewrites existing disabled_at/disabled_by audit metadata; - a task-level test proving repeated batch enable dispatches stay idempotent for an already enabled completed segment, which the reconciliation retry relies on. --- .../knowledge/segments/test_application.py | 210 ++++++++++++++++++ .../test_enable_segments_to_index_task.py | 37 +++ 2 files changed, 247 insertions(+) diff --git a/api/tests/unit_tests/services/knowledge/segments/test_application.py b/api/tests/unit_tests/services/knowledge/segments/test_application.py index 05813199ad3999..316d23cc5bcbd8 100644 --- a/api/tests/unit_tests/services/knowledge/segments/test_application.py +++ b/api/tests/unit_tests/services/knowledge/segments/test_application.py @@ -4,9 +4,13 @@ from unittest.mock import Mock, create_autospec import pytest +from sqlalchemy.orm import Session, sessionmaker from libs.helper import generate_text_hash from machinery.context import RequestContext +from models.dataset import DocumentSegment +from models.enums import SegmentStatus +from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository from services.knowledge.dataset_access import ( AccessibleDataset, DatasetAccess, @@ -274,6 +278,63 @@ def invoke() -> object: invoke() +def _service_with_real_store( + sqlite_session_factory: sessionmaker[Session], +) -> tuple[DatasetSegmentApplicationService, Mock]: + """Wire the application service to the real SQLAlchemy store over a SQLite session.""" + dataset_access = create_autospec(DatasetAccess, instance=True, spec_set=True) + dataset_access.check_access.return_value = AccessibleDataset(id="dataset-1", workspace_id="workspace-1") + store = SQLAlchemySegmentRepository(session_factory=sqlite_session_factory) + scopes = create_autospec(SegmentScopeReader, instance=True, spec_set=True) + scopes.get_segment_scope.return_value = _scope() + indexing_state = create_autospec(SegmentIndexingState, instance=True, spec_set=True) + indexing_state.is_document_indexing.return_value = False + indexing_state.is_segment_indexing.return_value = False + index = create_autospec(SegmentIndex, instance=True, spec_set=True) + service = DatasetSegmentApplicationService( + dataset_access=dataset_access, + store=store, + limits=SegmentLimitsStub(), + text_hash=generate_text_hash, + index=index, + uploads=create_autospec(SegmentUploadCatalog, instance=True, spec_set=True), + scopes=scopes, + model_guard=create_autospec(SegmentModelGuard, instance=True, spec_set=True), + indexing_state=indexing_state, + batch_dispatcher=create_autospec(SegmentBatchImportDispatcher, instance=True, spec_set=True), + job_id_factory=lambda: "job-1", + ) + return service, index + + +def _sqlite_segment( + segment_id: str, + *, + workspace_id: str = "workspace-1", + enabled: bool = True, + disabled_at: datetime | None = None, + disabled_by: str | None = None, +) -> DocumentSegment: + segment = DocumentSegment( + tenant_id=workspace_id, + dataset_id="dataset-1", + document_id="document-1", + position=1, + content="content", + word_count=7, + tokens=2, + created_by="account-1", + index_node_id="node-1", + index_node_hash="hash-1", + status=SegmentStatus.COMPLETED, + enabled=enabled, + disabled_at=disabled_at, + disabled_by=disabled_by, + ) + segment.id = segment_id + return segment + + def test_status_change_is_fenced_by_document_indexing_state() -> None: service, _, store, scopes, _, state, _ = _service() state.is_document_indexing.return_value = True @@ -386,6 +447,155 @@ def test_status_change_retry_after_dispatch_failure_reconciles_index() -> None: assert store.set_segments.call_count == 1 +def test_status_change_broker_failure_for_matching_batch_is_translated() -> None: + """A dispatch failure is surfaced even when every requested segment already matches + the requested status, instead of being swallowed as a no-op.""" + service, store, index = _service_with_index() + store.get_segments.return_value = (SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=True),) + index.change_status.side_effect = RuntimeError("broker unavailable") + + with pytest.raises(SegmentStatusUpdateError, match="broker unavailable"): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=["segment-1"], + action="enable", + ) + + store.set_segments.assert_not_called() + index.change_status.assert_called_once_with(_document().ref, ["segment-1"], "enable") + + +@pytest.mark.parametrize("action", ["enable", "disable"]) +@pytest.mark.parametrize("batch", ["single", "mixed"], ids=["single", "mixed-batch"]) +def test_status_change_persists_changes_and_dispatches_full_scope( + action: Literal["enable", "disable"], + batch: Literal["single", "mixed"], +) -> None: + """Persistence covers only actual status changes while the dispatch covers every + non-indexing segment in scope; foreign ids never reach either call.""" + service, store, index = _service_with_index() + targets = [ + SegmentIndexTarget(id="segment-1", index_node_id="node-1", enabled=action != "enable"), + SegmentIndexTarget(id="already-target", index_node_id="node-2", enabled=action == "enable"), + ] + store.get_segments.return_value = tuple(targets) if batch == "mixed" else tuple(targets[:1]) + requested_ids = ["segment-1", "already-target", "foreign"] if batch == "mixed" else ["segment-1", "foreign"] + + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=requested_ids, + action=action, + ) + + store.set_segments.assert_called_once() + persisted_ids, persisted_values = store.set_segments.call_args[0][1:3] + assert persisted_ids == ["segment-1"] + assert persisted_values["enabled"] == (action == "enable") + if action == "disable": + assert persisted_values["disabled_by"] == "account-1" + assert persisted_values["disabled_at"] is not None + else: + assert persisted_values["disabled_at"] is None + assert persisted_values["disabled_by"] is None + + index.change_status.assert_called_once() + dispatched_ref, dispatched_ids, dispatched_action = index.change_status.call_args[0] + assert dispatched_ref == _document().ref + assert dispatched_action == action + assert set(dispatched_ids) == {"segment-1"} | ({"already-target"} if batch == "mixed" else set()) + assert "foreign" not in dispatched_ids + + +@pytest.mark.parametrize("action", ["enable", "disable"]) +@pytest.mark.parametrize("batch", ["single", "mixed"], ids=["single", "mixed-batch"]) +def test_status_retry_reconciles_index_with_sqlite_persistence( + sqlite_session_factory: sessionmaker[Session], + action: Literal["enable", "disable"], + batch: Literal["single", "mixed"], +) -> None: + """End-to-end retry over a real SQLite session: a failed dispatch leaves the desired + status persisted, the retry re-enqueues the index task for every non-indexing + segment in scope (skipping out-of-scope ids) and never rewrites audit metadata.""" + service, index = _service_with_real_store(sqlite_session_factory) + previous_disabled_at = datetime(2025, 1, 1) + previous_disabled_by = "previous-editor" + with sqlite_session_factory.begin() as session: + session.add( + _sqlite_segment( + "segment-1", + enabled=action != "enable", + disabled_at=previous_disabled_at, + disabled_by=previous_disabled_by, + ) + ) + session.add(_sqlite_segment("foreign", workspace_id="workspace-2", enabled=action != "enable")) + if batch == "mixed": + session.add( + _sqlite_segment( + "already-target", + enabled=action == "enable", + disabled_at=previous_disabled_at, + disabled_by=previous_disabled_by, + ) + ) + + requested_ids = ["segment-1", "foreign"] + (["already-target"] if batch == "mixed" else []) + expected_ids = {"segment-1"} | ({"already-target"} if batch == "mixed" else set()) + index.change_status.side_effect = RuntimeError("broker unavailable") + + with pytest.raises(SegmentStatusUpdateError, match="broker unavailable"): + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=requested_ids, + action=action, + ) + + assert index.change_status.call_count == 1 + assert set(index.change_status.call_args[0][1]) == expected_ids + with sqlite_session_factory() as session: + target = session.get(DocumentSegment, "segment-1") + assert target is not None + assert target.enabled == (action == "enable") + if action == "disable": + assert target.disabled_by == "account-1" + assert target.disabled_at is not None + else: + assert (target.disabled_at, target.disabled_by) == (None, None) + target_audit = (target.disabled_at, target.disabled_by) + + index.change_status.side_effect = None + service.change_segment_status( + _context(), + dataset_id="dataset-1", + document_id="document-1", + segment_ids=requested_ids, + action=action, + ) + + assert index.change_status.call_count == 2 + assert index.change_status.call_args[0][0] == _document().ref + assert set(index.change_status.call_args[0][1]) == expected_ids + assert index.change_status.call_args[0][2] == action + with sqlite_session_factory() as session: + target = session.get(DocumentSegment, "segment-1") + assert target is not None + assert (target.disabled_at, target.disabled_by) == target_audit + foreign = session.get(DocumentSegment, "foreign") + assert foreign is not None + assert foreign.enabled == (action != "enable") + if batch == "mixed": + matched = session.get(DocumentSegment, "already-target") + assert matched is not None + assert matched.enabled == (action == "enable") + assert (matched.disabled_at, matched.disabled_by) == (previous_disabled_at, previous_disabled_by) + + def test_update_and_delete_segment_map_missing_segment() -> None: service, _, store, scopes, _, _, _ = _service() store.get_segment_update_state.return_value = 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..cb6eab7e7485d3 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,40 @@ 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() + + +def test_repeated_batch_enable_dispatches_stay_idempotent( + sqlite_session: Session, + indexed_segment: tuple[Dataset, Document, DocumentSegment], +) -> None: + """Reconciliation retries re-dispatch the same batch: repeated task runs keep an + already enabled completed segment stable instead of corrupting its state.""" + dataset, document, segment = indexed_segment + loads: list[list[str]] = [] + index_processor = MagicMock() + index_processor.load.side_effect = lambda _dataset, documents, **_kwargs: loads.append( + [item.metadata["doc_id"] for item in documents] + ) + + with ( + patch("tasks.enable_segments_to_index_task.IndexProcessorFactory") as processor_factory, + patch( + "services.knowledge.summaries.adapters.SummaryIndexAdapter.enable_summaries_for_segments" + ) as enable_summaries, + patch("tasks.enable_segments_to_index_task.redis_client.delete"), + ): + processor_factory.return_value.init_index_processor.return_value = index_processor + enable_segments_to_index_task.run([segment.id], dataset.id, document.id) + enable_segments_to_index_task.run([segment.id], dataset.id, document.id) + + assert loads == [["node-1"], ["node-1"]] + assert enable_summaries.call_count == 2 + + sqlite_session.expire_all() + persisted_segment = sqlite_session.get(DocumentSegment, segment.id) + assert persisted_segment is not None + assert persisted_segment.status == SegmentStatus.COMPLETED + assert persisted_segment.enabled is True + assert persisted_segment.index_node_id == "node-1" + assert persisted_segment.error is None + assert persisted_segment.disabled_at is None