fix: remap chunk results after skipped messages

src/typeagent/knowpro/conversation_base.py

@@ -243,6 +243,16 @@ async def _commit_batch_from_chunk_results(
         if not messages_batch:
             return AddMessagesResult()
 
+        # Chunk locations use source-stream ordinals. Those ordinals can have
+        # gaps when skip_failed_messages omits a message, so preserve their
+        # order for remapping to consecutive storage ordinals below.
+        source_message_ids: list[MessageOrdinal] = []
+        source_message_id: MessageOrdinal | None = None
+        for result in chunk_results:
+            if result.chunk_id.message_ordinal != source_message_id:
+                source_message_ids.append(result.chunk_id.message_ordinal)
+                source_message_id = result.chunk_id.message_ordinal
+
         # Process chunk results first to collect embeddings and knowledge items
         knowledge_items: list[tuple[MessageOrdinal, int, kplib.KnowledgeResponse]] = []
         fuzzy_terms: list[str] = []
@@ -290,7 +300,7 @@ async def _commit_batch_from_chunk_results(
                 )
             fuzzy_terms.extend(result.related_terms)
             fuzzy_term_embeddings.extend(result.related_term_embeddings)
-            # Store embedding for later retrieval in correct message/chunk order
+            # Store embedding by source message/chunk location until commit remaps it.
             chunk_embedding_map[
                 (result.chunk_id.message_ordinal, result.chunk_id.chunk_ordinal)
             ] = result.chunk_embedding
@@ -301,20 +311,53 @@ async def _commit_batch_from_chunk_results(
                 semref_count=await self.semantic_refs.size(),
             )
 
-            # Build chunk_embeddings in the correct order (matching message/chunk iteration)
+            source_to_storage_message_id: dict[MessageOrdinal, MessageOrdinal] = {}
             chunk_embeddings: list[NormalizedEmbedding] = []
-            for msg_ord, message in enumerate(
+            result_group_index = 0
+            for storage_message_id, message in enumerate(
                 messages_batch, start_points.message_count
             ):
-                for chunk_ord in range(len(message.text_chunks)):
-                    embedding = chunk_embedding_map.get((msg_ord, chunk_ord))
+                if not message.text_chunks:
+                    continue
+                if result_group_index >= len(source_message_ids):
+                    raise ValueError(
+                        "Missing chunk results for staged message: "
+                        f"message={storage_message_id}"
+                    )
+                source_message_id = source_message_ids[result_group_index]
+                result_group_index += 1
+                source_to_storage_message_id[source_message_id] = storage_message_id
+                for chunk_ordinal in range(len(message.text_chunks)):
+                    embedding = chunk_embedding_map.get(
+                        (source_message_id, chunk_ordinal)
+                    )
                     if embedding is None:
                         raise ValueError(
                             "Missing chunk embedding for staged message chunk: "
-                            f"message={msg_ord}, chunk={chunk_ord}"
+                            f"message={storage_message_id}, chunk={chunk_ordinal}"
                         )
                     chunk_embeddings.append(embedding)
 
+            if result_group_index != len(source_message_ids):
+                raise ValueError("Chunk results exceed staged messages with chunks")
+
+            remapped_knowledge_items: list[
+                tuple[MessageOrdinal, int, kplib.KnowledgeResponse]
+            ] = []
+            for source_message_id, chunk_ordinal, knowledge in knowledge_items:
+                storage_message_id = source_to_storage_message_id.get(
+                    source_message_id
+                )
+                if storage_message_id is None:
+                    raise ValueError(
+                        "No storage message id for chunk result: "
+                        f"source_message={source_message_id}, chunk={chunk_ordinal}"
+                    )
+                remapped_knowledge_items.append(
+                    (storage_message_id, chunk_ordinal, knowledge)
+                )
+            knowledge_items = remapped_knowledge_items
+
             # Use precomputed embeddings to avoid redundant embedding work
             await self.messages.extend(
                 messages_batch, chunk_embeddings=chunk_embeddings

tests/test_add_messages_pipeline.py

@@ -445,6 +445,62 @@ async def test_dispatcher_stops_on_sentinel_and_emits_result_sentinel() -> None:
     assert items[0].error is None
 
 
+@pytest.mark.asyncio
+async def test_shutdown_finishes_message_already_admitted_to_full_queue() -> None:
+    chunk_queue: asyncio.Queue[ChunkWorkItem[_Message] | None] = asyncio.Queue(
+        maxsize=1
+    )
+    result_queue = asyncio.Queue()
+    stop_state = PipelineStopState()
+    producer_state = ProducerState(next_message_id=0)
+    shutdown_event = asyncio.Event()
+    first_message = _Message(["first chunk", "second chunk"])
+    second_message = _Message(["not admitted"])
+
+    async def _iter_messages() -> AsyncIterator[_Message]:
+        yield first_message
+        yield second_message
+
+    producer_task = asyncio.create_task(
+        _producer_task(
+            _iter_messages(),
+            chunk_queue,
+            stop_state,
+            producer_state,
+            result_queue,
+            shutdown_event,
+        )
+    )
+    await asyncio.sleep(0)
+    assert chunk_queue.full()
+    shutdown_event.set()
+
+    await _dispatcher_task(
+        chunk_queue,
+        result_queue,
+        stop_state,
+        _SequenceExtractor(
+            [
+                typechat.Success(_empty_knowledge()),
+                typechat.Success(_empty_knowledge()),
+            ]
+        ),
+        _StubEmbeddingModel(),
+        concurrency=1,
+        skip_failed_messages=False,
+    )
+    await asyncio.wait_for(producer_task, timeout=1)
+
+    items = await _drain_result_queue(result_queue)
+    results = [item for item in items if item is not None]
+    assert [result.chunk_id for result in results] == [
+        TextLocation(0, 0),
+        TextLocation(0, 1),
+    ]
+    assert producer_state.produced_messages == 1
+    assert producer_state.produced_chunks == 2
+
+
 @pytest.mark.asyncio
 async def test_dispatcher_extraction_failure_lowers_stop() -> None:
     """A Failure from the extractor sets error and lowers stop_at_message_id."""

tests/test_add_messages_streaming.py

@@ -15,7 +15,7 @@
 from typeagent.knowpro import knowledge_schema as kplib
 from typeagent.knowpro.add_messages import add_messages_streaming
 from typeagent.knowpro.convsettings import ConversationSettings
-from typeagent.knowpro.interfaces_core import IKnowledgeExtractor
+from typeagent.knowpro.interfaces_core import AddMessagesResult, IKnowledgeExtractor
 from typeagent.storage.sqlite.provider import SqliteStorageProvider
 from typeagent.transcripts.transcript import (
     Transcript,
@@ -97,16 +97,22 @@ class ControlledExtractor:
     ``fail_on`` is a set of 0-based call indices for which the extractor
     returns a Failure instead of a Success.
     ``raise_on`` is a set of call indices that raise an exception.
+    ``fail_on_text`` is a set of chunk texts that always fail, regardless of
+    call order -- chunks are extracted concurrently, so call indices are not
+    deterministic when the failing chunk must be a specific one.  It is a
+    mutable set so a test can "repair" the extractor between runs.
     """
 
     def __init__(
         self,
         *,
         fail_on: set[int] | None = None,
         raise_on: set[int] | None = None,
+        fail_on_text: set[str] | None = None,
     ) -> None:
         self.fail_on = fail_on or set()
         self.raise_on = raise_on or set()
+        self.fail_on_text = fail_on_text or set()
         self.call_count = 0
 
     async def extract(self, message: str) -> typechat.Result[kplib.KnowledgeResponse]:
@@ -116,6 +122,8 @@ async def extract(self, message: str) -> typechat.Result[kplib.KnowledgeResponse
             raise RuntimeError(f"Systemic failure at call {idx}")
         if idx in self.fail_on:
             return typechat.Failure(f"Extraction failed for call {idx}")
+        if message in self.fail_on_text:
+            return typechat.Failure(f"Extraction failed for chunk {message!r}")
         return typechat.Success(_EMPTY_RESPONSE)
 
 
@@ -200,6 +208,36 @@ async def test_streaming_extraction_failure_stops_at_failing_message() -> None:
         await storage.close()
 
 
+@pytest.mark.asyncio
+async def test_streaming_skips_failed_message_and_commits_next_message() -> None:
+    """Skipped messages do not leave gaps in persisted message ordinals."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "test.db")
+        extractor = ControlledExtractor(fail_on={0})
+        transcript, storage = await _create_transcript(
+            db_path, auto_extract=True, knowledge_extractor=extractor
+        )
+        messages = [
+            _make_message("failed", source_id="failed-source"),
+            _make_message("succeeds", source_id="successful-source"),
+        ]
+
+        result = await add_messages_streaming(
+            transcript,
+            _async_iter(messages),
+            skip_failed_messages=True,
+        )
+
+        assert result.messages_added == 1
+        assert result.messages_skipped == 1
+        assert result.chunks_added == 1
+        assert await transcript.messages.get_slice(0, 1) == [messages[1]]
+        assert not await storage.is_source_ingested("failed-source")
+        assert await storage.is_source_ingested("successful-source")
+
+        await storage.close()
+
+
 @pytest.mark.asyncio
 async def test_streaming_exception_stops_run() -> None:
     """A raised exception stops processing; committed batches survive."""
@@ -706,3 +744,112 @@ async def test_streaming_extraction_returns_none_for_empty_chunks() -> None:
         assert extractor.call_count == 0
 
         await storage.close()
+
+
+# ---------------------------------------------------------------------------
+# Repair-and-reimport (replay) behavior
+# ---------------------------------------------------------------------------
+
+
+async def _message_texts(transcript: Transcript) -> list[str]:
+    """Return every stored chunk text, in storage order."""
+    size = await transcript.messages.size()
+    messages = await transcript.messages.get_slice(0, size)
+    return [chunk for message in messages for chunk in message.text_chunks]
+
+
+async def _replay(
+    transcript: Transcript,
+    storage: SqliteStorageProvider,
+    messages: list[TranscriptMessage],
+) -> AddMessagesResult:
+    """Re-submit ``messages``, pre-filtering sources the DB already has.
+
+    This mirrors what importers do (see ``tools/ingest_email.py``): ask the
+    storage provider which source IDs are already ingested and only stream the
+    remainder.  Re-running it must be idempotent.
+    """
+    source_ids = [m.source_id for m in messages if m.source_id is not None]
+    already_ingested = await storage.are_sources_ingested(source_ids)
+    pending = [m for m in messages if m.source_id not in already_ingested]
+    return await add_messages_streaming(
+        transcript,
+        _async_iter(pending),
+        skip_failed_messages=True,
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_replay_after_repair_ingests_each_source_once() -> None:
+    """fail -> repair -> replay -> replay stores each source exactly once.
+
+    The first run skips a message whose extraction fails, so its source is
+    never marked ingested.  After the extractor is repaired, a replay picks up
+    only that message, and a second replay is a no-op -- no duplicates.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "test.db")
+        extractor = ControlledExtractor(fail_on_text={"msg-1"})
+        transcript, storage = await _create_transcript(
+            db_path, auto_extract=True, knowledge_extractor=extractor
+        )
+        messages = [_make_message(f"msg-{i}", source_id=f"s-{i}") for i in range(3)]
+
+        # 1. Fail: msg-1 fails extraction and is skipped; the others commit.
+        first = await _replay(transcript, storage, messages)
+        assert first.messages_added == 2
+        assert first.messages_skipped == 1
+        assert await _message_texts(transcript) == ["msg-0", "msg-2"]
+        assert not await storage.is_source_ingested("s-1")
+        assert _ingested_count(storage) == 2
+
+        # 2. Repair: the extractor no longer fails on that chunk.
+        extractor.fail_on_text.clear()
+
+        # 3. Replay: only the previously failed source is re-submitted.
+        second = await _replay(transcript, storage, messages)
+        assert second.messages_added == 1
+        assert second.messages_skipped == 0
+        assert await storage.is_source_ingested("s-1")
+
+        # 4. Replay again: everything is ingested, so nothing is streamed.
+        third = await _replay(transcript, storage, messages)
+        assert third.messages_added == 0
+        assert third.messages_skipped == 0
+        assert third.chunks_added == 0
+
+        texts = await _message_texts(transcript)
+        assert len(texts) == len(set(texts))  # no duplicated messages
+        assert sorted(texts) == ["msg-0", "msg-1", "msg-2"]
+        assert await transcript.messages.size() == 3
+        assert _ingested_count(storage) == 3
+
+        await storage.close()
+
+
+@pytest.mark.asyncio
+async def test_streaming_replay_without_repair_does_not_duplicate() -> None:
+    """Replaying while the failure persists re-skips instead of duplicating."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        db_path = os.path.join(tmpdir, "test.db")
+        extractor = ControlledExtractor(fail_on_text={"msg-1"})
+        transcript, storage = await _create_transcript(
+            db_path, auto_extract=True, knowledge_extractor=extractor
+        )
+        messages = [_make_message(f"msg-{i}", source_id=f"s-{i}") for i in range(3)]
+
+        first = await _replay(transcript, storage, messages)
+        assert first.messages_added == 2
+        assert first.messages_skipped == 1
+
+        # Replay with the failure still in place: the two good sources are
+        # filtered out, and msg-1 fails and is skipped again.
+        second = await _replay(transcript, storage, messages)
+        assert second.messages_added == 0
+        assert second.messages_skipped == 1
+
+        texts = await _message_texts(transcript)
+        assert texts == ["msg-0", "msg-2"]
+        assert _ingested_count(storage) == 2
+
+        await storage.close()