fix(knowledge): reconcile segment indexes on repeated status requests

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
 

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:

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"