""" Tests for vector store reconciliation. This module tests the vector reconciliation system that syncs documents and message embeddings to the vector store, handling failures and retries. """ import datetime from contextlib import asynccontextmanager from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from nanoid import generate as generate_nanoid from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.reconciler.sync_vectors import ( MAX_SYNC_ATTEMPTS, ReconciliationMetrics, _get_documents_needing_sync, # pyright: ignore[reportPrivateUsage] _get_message_embeddings_needing_sync, # pyright: ignore[reportPrivateUsage] _reconcile_documents_batch, # pyright: ignore[reportPrivateUsage] _reconcile_message_embeddings_batch, # pyright: ignore[reportPrivateUsage] _sync_documents, # pyright: ignore[reportPrivateUsage] _sync_message_embeddings, # pyright: ignore[reportPrivateUsage] build_message_vector_record, compute_chunk_positions, run_vector_reconciliation_cycle, ) from src.vector_store import ( VectorRecord, VectorStore, _hash_namespace_components, ) @pytest.mark.asyncio class TestStateTransitions: """Test document sync_state transitions.""" async def test_pending_to_synced_on_success( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test documents transition from pending → synced on successful sync.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create documents in pending state with embeddings docs = [ models.Document( content=f"doc_{i}", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=0, embedding=[float(i)] * 1536, # Mock embedding ) for i in range(3) ] db_session.add_all(docs) await db_session.commit() for doc in docs: await db_session.refresh(doc) # Mock vector store to succeed mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" ) mock_vector_store.upsert_many = AsyncMock(return_value=None) # Run sync synced, failed = await _sync_documents(db_session, docs, mock_vector_store) # Verify results assert synced == 3 assert failed == 0 # Check state transitions for doc in docs: await db_session.refresh(doc) assert doc.sync_state == "synced" assert doc.sync_attempts == 0 assert doc.last_sync_at is not None async def test_pending_to_pending_with_incremented_attempts( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test documents remain pending with incremented attempts on partial failure.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create document in pending state doc = models.Document( content="test doc", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=2, # Already failed twice embedding=[1.0] * 1536, ) db_session.add(doc) await db_session.commit() await db_session.refresh(doc) # Mock vector store to fail with exception mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" ) mock_vector_store.upsert_many = AsyncMock( side_effect=Exception("Vector store failed") ) # Run sync synced, failed = await _sync_documents(db_session, [doc], mock_vector_store) # Verify failure recorded assert synced == 0 assert failed == 1 # Check sync_attempts incremented await db_session.refresh(doc) assert doc.sync_state == "pending" # Still pending assert doc.sync_attempts == 3 # Incremented assert doc.last_sync_at is not None # Updated timestamp async def test_pending_to_failed_after_max_attempts( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test documents transition to failed after MAX_SYNC_ATTEMPTS failures.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create document at max attempts - 1 doc = models.Document( content="failing doc", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=MAX_SYNC_ATTEMPTS - 1, # One more attempt will hit limit embedding=[1.0] * 1536, ) db_session.add(doc) await db_session.commit() await db_session.refresh(doc) # Mock vector store to fail with exception mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" ) mock_vector_store.upsert_many = AsyncMock( side_effect=Exception("Vector store failed") ) # Run sync - this should be the final attempt synced, failed = await _sync_documents(db_session, [doc], mock_vector_store) # Verify marked as failed assert synced == 0 assert failed == 1 await db_session.refresh(doc) assert doc.sync_state == "failed" # Permanently failed assert doc.sync_attempts == MAX_SYNC_ATTEMPTS @pytest.mark.asyncio class TestBatchProcessing: """Test batch processing and namespace grouping.""" async def test_documents_grouped_by_namespace( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test documents from different collections are grouped by namespace.""" workspace, peer1 = sample_data # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create another peer for different observer/observed combinations peer2 = models.Peer( name="peer2", workspace_name=workspace.name, ) db_session.add(peer2) await db_session.commit() # Create collections for both observer/observed combinations collection1 = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) collection2 = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, ) db_session.add_all([collection1, collection2]) await db_session.commit() # Create documents for different namespaces docs = [ # Namespace 1: peer1 → peer1 models.Document( content="doc1", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=[1.0] * 1536, ), # Namespace 2: peer1 → peer2 models.Document( content="doc2", workspace_name=workspace.name, observer=peer1.name, observed=peer2.name, session_name=session.name, sync_state="pending", embedding=[2.0] * 1536, ), # Namespace 1 again: peer1 → peer1 models.Document( content="doc3", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=[3.0] * 1536, ), ] db_session.add_all(docs) await db_session.commit() # Mock vector store to track calls by namespace mock_vector_store = MagicMock(spec=VectorStore) namespace_calls: dict[str, list[VectorRecord]] = {} def mock_get_namespace( _namespace_type: str, workspace: str, observer: str, observed: str ) -> str: return f"honcho.doc.{_hash_namespace_components(workspace, observer, observed)}" async def mock_upsert(namespace: str, vectors: list[VectorRecord]) -> None: if namespace not in namespace_calls: namespace_calls[namespace] = [] namespace_calls[namespace].extend(vectors) return mock_vector_store.get_vector_namespace = mock_get_namespace mock_vector_store.upsert_many = mock_upsert # Run sync synced, failed = await _sync_documents(db_session, docs, mock_vector_store) # Verify all synced assert synced == 3 assert failed == 0 # Verify namespaces expected_ns1 = f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" expected_ns2 = f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer2.name)}" assert expected_ns1 in namespace_calls assert expected_ns2 in namespace_calls assert len(namespace_calls[expected_ns1]) == 2 # doc1 and doc3 assert len(namespace_calls[expected_ns2]) == 1 # doc2 async def test_batch_size_respected( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test that batch size limits are respected when fetching pending documents.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create more documents than batch size docs = [ models.Document( content=f"doc_{i}", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=[float(i)] * 1536, ) for i in range(150) # More than RECONCILIATION_BATCH_SIZE (100) ] db_session.add_all(docs) await db_session.commit() # Fetch documents with batch size limit batch = await _get_documents_needing_sync(db_session, batch_size=100) # Verify batch size respected assert len(batch) == 100 async def test_documents_respect_retry_backoff( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Pending documents should only be fetched once their backoff has elapsed.""" workspace, peer1 = sample_data collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() now = datetime.datetime.now(datetime.timezone.utc) ineligible_doc = models.Document( content="too soon", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=1, last_sync_at=now - datetime.timedelta(minutes=9, seconds=59), embedding=[1.0] * 1536, ) eligible_doc = models.Document( content="ready", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=1, last_sync_at=now - datetime.timedelta(minutes=10, seconds=1), embedding=[2.0] * 1536, ) db_session.add_all([ineligible_doc, eligible_doc]) await db_session.commit() pending = await _get_documents_needing_sync(db_session) pending_ids = {doc.id for doc in pending} assert eligible_doc.id in pending_ids assert ineligible_doc.id not in pending_ids @pytest.mark.asyncio class TestReEmbedding: """Test re-embedding logic for documents with NULL embeddings.""" async def test_documents_without_embeddings_are_reembedded( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test documents with NULL embeddings are re-embedded during reconciliation.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create documents without embeddings docs = [ models.Document( content=f"doc_{i}", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=None, # NULL embedding ) for i in range(3) ] db_session.add_all(docs) await db_session.commit() for doc in docs: await db_session.refresh(doc) # Mock embedding client with patch("src.reconciler.sync_vectors.embedding_client") as mock_embed_client: mock_embed_client.simple_batch_embed = AsyncMock( return_value=[[float(i)] * 1536 for i in range(3)] ) # Mock vector store mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" ) mock_vector_store.upsert_many = AsyncMock(return_value=None) # Run sync synced, failed = await _sync_documents(db_session, docs, mock_vector_store) # Verify embedding was called mock_embed_client.simple_batch_embed.assert_called_once() # Verify documents were synced assert synced == 3 assert failed == 0 async def test_large_documents_embedded_in_single_batch( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test that documents are embedded in a single batch (no sub-batching).""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create documents without embeddings docs = [ models.Document( content=f"doc_{i}", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=None, # Will be re-embedded ) for i in range(3) ] db_session.add_all(docs) await db_session.commit() for doc in docs: await db_session.refresh(doc) # Mock embedding client to track batch calls batch_call_count = 0 async def track_batch_embed(contents: list[str]) -> list[list[float]]: nonlocal batch_call_count batch_call_count += 1 return [[1.0] * 1536 for _ in contents] with patch("src.reconciler.sync_vectors.embedding_client") as mock_embed_client: mock_embed_client.simple_batch_embed = track_batch_embed # Mock vector store mock_vector_store = MagicMock(spec=VectorStore) mock_vector_store.get_vector_namespace = MagicMock( return_value=f"honcho.doc.{_hash_namespace_components(workspace.name, peer1.name, peer1.name)}" ) mock_vector_store.upsert_many = AsyncMock(return_value=None) # Run sync await _sync_documents(db_session, docs, mock_vector_store) # Verify single batch call (no sub-batching) assert batch_call_count == 1 @pytest.mark.asyncio class TestSoftDeleteCleanup: """Test soft delete cleanup functionality.""" async def test_cleanup_respects_grace_period( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test cleanup only processes documents deleted_at older than threshold.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create recently soft-deleted document (within grace period) recent_doc = models.Document( content="recent", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, deleted_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=2), # Only 2 minutes ago ) db_session.add(recent_doc) # Create old soft-deleted document (outside grace period) old_doc = models.Document( content="old", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, deleted_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10), # 10 minutes ago ) db_session.add(old_doc) await db_session.commit() # Query for documents ready for cleanup (older than 5 minutes) cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( minutes=5 ) stmt = ( select(models.Document) .where(models.Document.deleted_at.is_not(None)) .where(models.Document.deleted_at < cutoff) ) result = await db_session.execute(stmt) ready_for_cleanup = result.scalars().all() # Only old_doc should be ready assert len(ready_for_cleanup) == 1 assert ready_for_cleanup[0].id == old_doc.id @pytest.mark.asyncio class TestMetricsTracking: """Test reconciliation metrics collection.""" async def test_metrics_track_sync_results( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Test ReconciliationMetrics tracks synced/failed counts.""" workspace, peer1 = sample_data # Create collection (required for documents) collection = models.Collection( workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, ) db_session.add(collection) await db_session.commit() # Create session session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # Create mix of documents that will succeed and fail success_doc = models.Document( content="success", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", embedding=[1.0] * 1536, ) fail_doc = models.Document( content="fail", workspace_name=workspace.name, observer=peer1.name, observed=peer1.name, session_name=session.name, sync_state="pending", sync_attempts=MAX_SYNC_ATTEMPTS - 1, # Will fail on next attempt embedding=[2.0] * 1536, ) db_session.add_all([success_doc, fail_doc]) await db_session.commit() # Note: Since both docs have same namespace, they'll be in one batch # This test structure needs adjustment for the actual grouping logic # For simplicity, let's test metrics at the function level metrics = ReconciliationMetrics() metrics.documents_synced = 5 metrics.documents_failed = 2 metrics.message_embeddings_synced = 3 assert metrics.total_synced == 8 assert metrics.total_failed == 2 @pytest.mark.asyncio class TestMessageEmbeddings: """Test message embedding reconciliation paths.""" async def _create_pending_message_embedding( self, db_session: AsyncSession, workspace: models.Workspace, peer: models.Peer, ) -> models.MessageEmbedding: """Helper to create a pending message embedding with no stored vector.""" session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() message = models.Message( public_id=str(generate_nanoid()), session_name=session.name, workspace_name=workspace.name, peer_name=peer.name, content="hello world", seq_in_session=1, ) db_session.add(message) await db_session.commit() emb = models.MessageEmbedding( content=message.content, message_id=message.public_id, workspace_name=workspace.name, session_name=session.name, peer_name=peer.name, sync_state="pending", embedding=None, ) db_session.add(emb) await db_session.commit() await db_session.refresh(emb) return emb async def test_pending_embeddings_are_selected_without_vectors( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Pending rows with NULL embeddings should still be reconciled.""" workspace, peer = sample_data pending_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) pending = await _get_message_embeddings_needing_sync(db_session) assert any(emb.id == pending_emb.id for emb in pending) async def test_message_embeddings_respect_retry_backoff( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Pending embeddings should only be fetched once their backoff has elapsed.""" workspace, peer = sample_data ineligible_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) eligible_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) now = datetime.datetime.now(datetime.timezone.utc) ineligible_emb.sync_attempts = 1 ineligible_emb.last_sync_at = now - datetime.timedelta(minutes=9, seconds=59) eligible_emb.sync_attempts = 1 eligible_emb.last_sync_at = now - datetime.timedelta(minutes=10, seconds=1) await db_session.commit() pending = await _get_message_embeddings_needing_sync(db_session) pending_ids = {emb.id for emb in pending} assert eligible_emb.id in pending_ids assert ineligible_emb.id not in pending_ids async def test_missing_embeddings_reembedded_and_synced( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], mock_vector_store: VectorStore, ) -> None: """Reconciliation should re-embed missing payloads and mark them synced.""" workspace, peer = sample_data pending_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) synced, failed = await _sync_message_embeddings( db_session, [pending_emb], mock_vector_store ) await db_session.refresh(pending_emb) assert synced == 1 assert failed == 0 assert pending_emb.sync_state == "synced" assert pending_emb.sync_attempts == 0 # Ensure vector upsert was attempted with a populated embedding upsert_mock: AsyncMock = cast(AsyncMock, mock_vector_store.upsert_many) await_args = upsert_mock.await_args assert await_args is not None args = await_args.args assert len(args) == 2 vector_records: list[VectorRecord] = args[1] assert vector_records assert vector_records[0].embedding async def test_upsert_failure_marks_attempt_and_continues( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], mock_vector_store: VectorStore, ) -> None: """Failures during upsert should bump attempts and keep row pending/failed.""" workspace, peer = sample_data pending_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) upsert_mock: AsyncMock = cast(AsyncMock, mock_vector_store.upsert_many) upsert_mock.side_effect = Exception("boom") synced, failed = await _sync_message_embeddings( db_session, [pending_emb], mock_vector_store ) await db_session.refresh(pending_emb) assert synced == 0 assert failed == 1 assert pending_emb.sync_state in {"pending", "failed"} assert pending_emb.sync_attempts == 1 async def test_unexpected_batch_exception_does_not_bump_unattempted_rows( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], mock_vector_store: VectorStore, ) -> None: """Unexpected wrapper-level failures should not penalize the whole batch.""" workspace, peer = sample_data pending_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) metrics = ReconciliationMetrics() with ( patch( "src.reconciler.sync_vectors._sync_message_embeddings", side_effect=RuntimeError("unexpected"), ), pytest.raises(RuntimeError, match="unexpected"), ): await _reconcile_message_embeddings_batch(mock_vector_store, metrics) await db_session.refresh(pending_emb) assert pending_emb.sync_state == "pending" assert pending_emb.sync_attempts == 0 assert pending_emb.last_sync_at is None async def test_pgvector_only_mode_embeds_and_marks_synced( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """In pgvector-only mode, the reconciler must still embed pending rows.""" workspace, peer = sample_data pending_emb = await self._create_pending_message_embedding( db_session, workspace, peer ) # external_vector_store=None == pgvector-only mode. The reconciler should # re-embed the pending row, write the vector to postgres, and mark synced. synced, failed = await _sync_message_embeddings(db_session, [pending_emb], None) await db_session.commit() await db_session.refresh(pending_emb) assert synced == 1 assert failed == 0 assert pending_emb.sync_state == "synced" assert pending_emb.sync_attempts == 0 assert pending_emb.embedding is not None async def test_all_chunks_of_a_message_claimed_together( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """A single message's chunks must always be claimed in one batch. Selecting by message_id (not row) keeps `{message_id}_{chunk_index}` vector IDs stable across reconciler cycles. """ workspace, peer = sample_data # Create one message with 5 chunks. session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() message_id = str(generate_nanoid()) message = models.Message( public_id=message_id, session_name=session.name, workspace_name=workspace.name, peer_name=peer.name, content="full message content", seq_in_session=1, ) db_session.add(message) await db_session.commit() chunk_count = 5 for i in range(chunk_count): db_session.add( models.MessageEmbedding( content=f"chunk-{i}", message_id=message_id, workspace_name=workspace.name, session_name=session.name, peer_name=peer.name, sync_state="pending", embedding=None, ) ) await db_session.commit() # Even with batch_size=1, all 5 chunks for the message should be claimed # together because the query selects by distinct message_id first. claimed = await _get_message_embeddings_needing_sync(db_session, batch_size=1) assert len(claimed) == chunk_count assert all(emb.message_id == message_id for emb in claimed) @pytest.mark.asyncio class TestEndToEndReconciliation: """Test full reconciliation cycle.""" async def test_reconciliation_cycle_completes_when_no_work_found(self) -> None: """Reconciliation should run each stage once and stop when no work is found.""" mock_vector_store = MagicMock(spec=VectorStore) with ( patch( "src.reconciler.sync_vectors.get_external_vector_store", return_value=mock_vector_store, ), patch( "src.reconciler.sync_vectors._reconcile_documents_batch", new_callable=AsyncMock, ) as mock_reconcile_docs, patch( "src.reconciler.sync_vectors._reconcile_message_embeddings_batch", new_callable=AsyncMock, ) as mock_reconcile_embs, patch( "src.reconciler.sync_vectors._cleanup_documents_batch", new_callable=AsyncMock, ) as mock_cleanup_docs, ): mock_reconcile_docs.return_value = False mock_reconcile_embs.return_value = False mock_cleanup_docs.return_value = False metrics = await run_vector_reconciliation_cycle() assert isinstance(metrics, ReconciliationMetrics) assert metrics.total_synced == 0 assert metrics.total_failed == 0 assert metrics.total_cleaned == 0 mock_reconcile_docs.assert_awaited_once() mock_reconcile_embs.assert_awaited_once() mock_cleanup_docs.assert_awaited_once() @pytest.mark.asyncio class TestReconcilerTracing: """A Sentry transaction is started only when a sync batch finds real work. Reconciler tasks poll on a fixed interval and usually find nothing; an idle cycle must create zero transactions so it doesn't drain Sentry quota. """ @staticmethod def _fake_tracked_db(db: AsyncMock): @asynccontextmanager async def _cm(*_args: object, **_kwargs: object): yield db return _cm async def test_no_transaction_when_no_embeddings_to_sync(self) -> None: """The no-work path returns before starting a transaction.""" metrics = ReconciliationMetrics() with ( patch( "src.reconciler.sync_vectors.tracked_db", self._fake_tracked_db(AsyncMock()), ), patch( "src.reconciler.sync_vectors._get_message_embeddings_needing_sync", new_callable=AsyncMock, return_value=[], ), patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, ): worked = await _reconcile_message_embeddings_batch(None, metrics) assert worked is False txn.assert_not_called() async def test_transaction_started_when_embeddings_present(self) -> None: """A batch with real work starts its own named transaction.""" metrics = ReconciliationMetrics() with ( patch( "src.reconciler.sync_vectors.tracked_db", self._fake_tracked_db(AsyncMock()), ), patch( "src.reconciler.sync_vectors._get_message_embeddings_needing_sync", new_callable=AsyncMock, return_value=[MagicMock()], ), patch( "src.reconciler.sync_vectors._sync_message_embeddings", new_callable=AsyncMock, return_value=(1, 0), ), patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, ): worked = await _reconcile_message_embeddings_batch(None, metrics) assert worked is True assert metrics.message_embeddings_synced == 1 txn.assert_called_once() assert txn.call_args.kwargs.get("name") == "reconcile_message_embeddings_batch" async def test_no_transaction_when_no_documents_to_sync(self) -> None: """The document batch also skips tracing when there is nothing to sync.""" metrics = ReconciliationMetrics() with ( patch( "src.reconciler.sync_vectors.tracked_db", self._fake_tracked_db(AsyncMock()), ), patch( "src.reconciler.sync_vectors._get_documents_needing_sync", new_callable=AsyncMock, return_value=[], ), patch("src.reconciler.sync_vectors.sentry_sdk.start_transaction") as txn, ): worked = await _reconcile_documents_batch( MagicMock(spec=VectorStore), metrics ) assert worked is False txn.assert_not_called() def test_build_message_vector_record() -> None: """The shared vector-id/metadata builder: id is {message_id}_{position}, embeddings are coerced to float, metadata shape is fixed.""" record = build_message_vector_record( message_id="msg_abc", chunk_position=2, session_name="sess", peer_name="peer", embedding=[1, 2, 3], # ints, must be coerced ) assert record.id == "msg_abc_2" assert record.embedding == [1.0, 2.0, 3.0] assert all(isinstance(x, float) for x in record.embedding) assert record.metadata == { "message_id": "msg_abc", "session_name": "sess", "peer_name": "peer", } @pytest.mark.asyncio class TestComputeChunkPositions: """Direct coverage for compute_chunk_positions, the source of truth for {message_id}_{position} vector ids shared by the reconciler and embed_now.""" async def test_empty_input_returns_empty(self, db_session: AsyncSession) -> None: assert await compute_chunk_positions(db_session, []) == {} async def test_positions_are_per_message_zero_indexed( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ) -> None: """Each message's rows are numbered from 0 in (message_id, id) order, independent of how rows from other messages interleave.""" workspace, peer = sample_data session = models.Session( name=str(generate_nanoid()), workspace_name=workspace.name ) db_session.add(session) await db_session.commit() # msg_a has 2 chunks, msg_b has 1 chunk. msg_a = str(generate_nanoid()) msg_b = str(generate_nanoid()) for seq, mid in enumerate((msg_a, msg_b), start=1): db_session.add( models.Message( public_id=mid, session_name=session.name, workspace_name=workspace.name, peer_name=peer.name, content="content", seq_in_session=seq, ) ) await db_session.commit() rows = [ models.MessageEmbedding( content=content, message_id=mid, workspace_name=workspace.name, session_name=session.name, peer_name=peer.name, sync_state="pending", embedding=None, ) for mid, content in ( (msg_a, "a0"), (msg_a, "a1"), (msg_b, "b0"), ) ] db_session.add_all(rows) await db_session.commit() for row in rows: await db_session.refresh(row) a0, a1, b0 = (row.id for row in rows) positions = await compute_chunk_positions(db_session, [msg_a, msg_b]) assert positions[a0] == 0 assert positions[a1] == 1 assert positions[b0] == 0