diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index fe0f3ab1..8ad0b1ad 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -435,33 +435,29 @@ async def enqueue_dream( observer: str, observed: str, dream_type: schemas.DreamType, - document_count: int, session_name: str | None = None, ) -> None: """ Enqueue a dream task for immediate processing by the deriver. + Does not touch collection.internal_metadata["dream"]. Both guard fields + (last_dream_at and last_dream_document_count) are written atomically in + process_dream on successful completion, so failed dreams retry against + the same corpus and a queued dream already acts as a stampede latch via + check_and_schedule_dream's in-flight check. + Deduplication: If a dream with the same work_unit_key is already in-progress - (has an ActiveQueueSession), the enqueue is skipped to prevent running - multiple dreams concurrently for the same collection. + (has an ActiveQueueSession) or pending in the queue, the enqueue is skipped. Args: workspace_name: Name of the workspace observer: Name of the observer peer observed: Name of the observed peer dream_type: Type of dream to execute - document_count: Count of explicit-level documents only. Written as - last_dream_document_count and used as the baseline that - dream_scheduler.check_and_schedule_dream subtracts from a future - explicit-filtered count to compute documents_since_last_dream. - Callers that include non-explicit levels (deductive, inductive, - contradiction) will inflate the baseline and suppress the next - scheduled dream. session_name: Name of the session to scope the dream to if specified """ async with tracked_db("dream_enqueue") as db_session: try: - # Create the dream queue record dream_record = create_dream_record( workspace_name, observer=observer, @@ -472,11 +468,6 @@ async def enqueue_dream( work_unit_key = dream_record["work_unit_key"] - # Check if a dream with this work_unit_key is currently in progress - # (has an ActiveQueueSession, meaning a worker is processing it) - # We only block on in-progress dreams, not pending ones - if there's - # a pending dream, we don't need to add another one anyway since - # the queue processor will pick it up. in_progress_check = select( exists( select(models.ActiveQueueSession.id).where( @@ -496,7 +487,6 @@ async def enqueue_dream( ) return - # Check if there's already a pending dream with the same work_unit_key pending_check = select( exists( select(QueueItem.id).where( @@ -517,34 +507,9 @@ async def enqueue_dream( ) return - # Insert into queue stmt = insert(QueueItem).returning(QueueItem) await db_session.execute(stmt, [dream_record]) - - # Update collection metadata (CRUD handles cache invalidation). - # last_dream_at is written at completion in process_dream, not here. - # Read-modify-write: update_collection_internal_metadata uses a - # top-level JSONB || merge, so passing {"dream": {"last_dream_document_count": ...}} - # alone would replace the whole "dream" subkey and drop sibling - # last_dream_at (written by process_dream on the prior completion). - # Row-lock the collection during read-modify-write to prevent concurrent enqueue/completion from clobbering each other's dream metadata writes. - collection = await crud.get_collection( - db_session, - workspace_name, - observer=observer, - observed=observed, - with_for_update=True, - ) - dream_meta = dict(collection.internal_metadata.get("dream", {})) - dream_meta["last_dream_document_count"] = document_count - await crud.update_collection_internal_metadata( - db_session, - workspace_name, - observer, - observed, - update_data={"dream": dream_meta}, - ) - # update_collection_internal_metadata commits already + await db_session.commit() logger.info( "Enqueued dream task for %s/%s/%s (type: %s)", diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index 5e3a7340..d7f0aedf 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone from logging import getLogger import sentry_sdk -from sqlalchemy import func, select +from sqlalchemy import exists, func, select from sqlalchemy.ext.asyncio import AsyncSession from src import models @@ -160,16 +160,11 @@ class DreamScheduler: observer: str, observed: str, ) -> None: - """Execute the dream by enqueueing it and updating collection metadata.""" - # Import here to avoid circular dependency + """Execute the dream by enqueueing it.""" from src import crud from src.deriver.enqueue import enqueue_dream from src.utils.config_helpers import get_configuration - # Find the most recent explicit-level session and get current document count. - # Explicit-only, matching the count query below and check_and_schedule_dream — - # if the session is picked from a derived doc while the count filters to - # explicit, the two can disagree on which document set the dream is over. async with tracked_db("dream_session_lookup") as db: stmt = ( select(models.Document.session_name) @@ -190,16 +185,6 @@ class DreamScheduler: ) return - # Get current document count at execution time (not stale from scheduling). - # Explicit-only, matching check_and_schedule_dream — baseline must be symmetric or the delta goes negative. - count_stmt = select(func.count(models.Document.id)).where( - models.Document.workspace_name == workspace_name, - models.Document.observer == observer, - models.Document.observed == observed, - models.Document.level == "explicit", - ) - current_explicit_count = int(await db.scalar(count_stmt) or 0) - session = await crud.get_session( db, workspace_name=workspace_name, session_name=session_name ) @@ -218,7 +203,6 @@ class DreamScheduler: observer=observer, observed=observed, dream_type=dream_type, - document_count=current_explicit_count, session_name=session_name, ) @@ -237,13 +221,18 @@ async def check_and_schedule_dream( collection: models.Collection, ) -> bool: """ + From the moment a dream is scheduled until it completes or fails, no second + dream may be enqueued for the same (workspace, observer, observed) — and the + baseline count advances only when consolidation actually happened. + Check if a collection has reached the explicit-observation threshold and schedule a timer-based dream. This function only schedules a timer-based dream if: 1. Dreams are enabled 2. Explicit-observation threshold is reached (dreamer output does not count) 3. Minimum hours between dreams have passed - 4. No dream is already scheduled for this collection + 4. No dream is already pending in the queue for this collection (in-flight check) + 5. No dream is already scheduled for this collection Args: db: Database session @@ -255,7 +244,6 @@ async def check_and_schedule_dream( if not settings.DREAM.ENABLED: return False - # Get dream metadata from internal_metadata dream_metadata = collection.internal_metadata.get("dream", {}) last_dream_document_count = dream_metadata.get("last_dream_document_count", 0) last_dream_at = dream_metadata.get("last_dream_at") @@ -270,7 +258,6 @@ async def check_and_schedule_dream( ) current_explicit_count = int(await db.scalar(count_stmt) or 0) - # Calculate documents added since last dream documents_since_last_dream = current_explicit_count - last_dream_document_count logger.debug( @@ -286,9 +273,7 @@ async def check_and_schedule_dream( }, ) - # Only schedule timer if document threshold is reached if documents_since_last_dream >= settings.DREAM.DOCUMENT_THRESHOLD: - # Check if we're within minimum hours between dreams if last_dream_at: try: last_dream_time = datetime.fromisoformat(last_dream_at) @@ -307,11 +292,46 @@ async def check_and_schedule_dream( f"Invalid last_dream_at timestamp: {last_dream_at}, error: {e}" ) + # In-flight stampede defense: the guard fields don't advance until a + # dream completes, so a queued-but-unprocessed dream leaves the time + # and count guards stale. Query the queue directly (source of truth, + # matches uq_queue_dream_pending_work_unit_key) to block a second + # schedule for any enabled dream_type on this collection. + enabled_dream_types = settings.DREAM.ENABLED_TYPES + pending_keys = [ + construct_work_unit_key( + collection.workspace_name, + { + "task_type": "dream", + "observer": collection.observer, + "observed": collection.observed, + "dream_type": dream_type, + }, + ) + for dream_type in enabled_dream_types + ] + pending_exists = await db.scalar( + select( + exists( + select(models.QueueItem.id).where( + models.QueueItem.task_type == "dream", + models.QueueItem.processed == False, # noqa: E712 + models.QueueItem.work_unit_key.in_(pending_keys), + ) + ) + ) + ) + if pending_exists: + logger.info( + "Skipping dream schedule for %s/%s: pending dream already in queue", + collection.observer, + collection.observed, + ) + return False + dream_scheduler = get_dream_scheduler() if dream_scheduler: - enabled_dream_types = settings.DREAM.ENABLED_TYPES for dream_type in enabled_dream_types: - # Include dream_type in key so each dream type can be tracked independently dream_work_unit_key = construct_work_unit_key( collection.workspace_name, { diff --git a/src/dreamer/orchestrator.py b/src/dreamer/orchestrator.py index 763eca4a..5e8cda3b 100644 --- a/src/dreamer/orchestrator.py +++ b/src/dreamer/orchestrator.py @@ -21,8 +21,9 @@ from datetime import datetime, timezone from typing import Any import sentry_sdk +from sqlalchemy import func, select -from src import crud +from src import crud, models from src.config import settings from src.dependencies import tracked_db from src.dreamer.specialists import SPECIALISTS, SpecialistResult @@ -324,15 +325,12 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p + f"duration={result.total_duration_ms:.0f}ms" ) - # Write last_dream_at at completion (not enqueue) so duplicate - # enqueues can't reset the 8h guard. Lenient: any non-null result. - # Read-modify-write: update_collection_internal_metadata uses a - # top-level JSONB || merge, so passing {"dream": {"last_dream_at": ...}} - # alone would replace the whole "dream" subkey and drop - # last_dream_document_count (written by enqueue_dream). - # Row-lock the collection during read-modify-write to prevent concurrent enqueue/completion from clobbering each other's dream metadata writes. + # Atomic pair write: both guard fields advance together only + # when consolidation actually happened. Splitting the writes + # across enqueue + completion creates an in-flight window + # where the pair tells contradictory stories. now_iso = datetime.now(timezone.utc).isoformat() - async with tracked_db("dream.last_dream_at_write") as db: + async with tracked_db("dream.guard_pair_write") as db: collection = await crud.get_collection( db, workspace_name, @@ -340,8 +338,16 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p observed=payload.observed, with_for_update=True, ) + count_stmt = select(func.count(models.Document.id)).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == payload.observer, + models.Document.observed == payload.observed, + models.Document.level == "explicit", + ) + current_explicit_count = int(await db.scalar(count_stmt) or 0) dream_meta = dict(collection.internal_metadata.get("dream", {})) dream_meta["last_dream_at"] = now_iso + dream_meta["last_dream_document_count"] = current_explicit_count await crud.update_collection_internal_metadata( db, workspace_name, diff --git a/src/routers/workspaces.py b/src/routers/workspaces.py index df3eecea..2ce7cdda 100644 --- a/src/routers/workspaces.py +++ b/src/routers/workspaces.py @@ -3,10 +3,9 @@ import logging from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Response from fastapi_pagination import Page from fastapi_pagination.ext.sqlalchemy import apaginate -from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession -from src import crud, models, schemas +from src import crud, schemas from src.config import settings from src.dependencies import db from src.deriver.enqueue import enqueue_deletion, enqueue_dream @@ -201,7 +200,6 @@ async def schedule_dream( request: schemas.ScheduleDreamRequest = Body( ..., description="Dream scheduling parameters" ), - db: AsyncSession = db, ): """ Manually schedule a dream task for a specific collection. @@ -224,23 +222,11 @@ async def schedule_dream( observed = request.observed if request.observed is not None else request.observer dream_type = request.dream_type - # Explicit-only count, matching dream_scheduler.check_and_schedule_dream and - # execute_dream — document_count becomes last_dream_document_count (the baseline). - count_stmt = select(func.count(models.Document.id)).where( - models.Document.workspace_name == workspace_id, - models.Document.observer == observer, - models.Document.observed == observed, - models.Document.level == "explicit", - ) - document_count = int(await db.scalar(count_stmt) or 0) - - # Enqueue the dream task for immediate processing await enqueue_dream( workspace_id, observer=observer, observed=observed, dream_type=dream_type, - document_count=document_count, session_name=request.session_id, ) diff --git a/tests/deriver/test_enqueue_dream.py b/tests/deriver/test_enqueue_dream.py index f8884170..bea5c256 100644 --- a/tests/deriver/test_enqueue_dream.py +++ b/tests/deriver/test_enqueue_dream.py @@ -1,34 +1,24 @@ """Regression tests for `enqueue_dream` metadata write shape. -Finding 3 (code-level) moves the `last_dream_at` timestamp write from -enqueue time to dream-completion time (in `process_dream`). These tests -verify that `enqueue_dream` no longer writes `last_dream_at` and still -writes `last_dream_document_count`. +Loop 4 (PR #573): `enqueue_dream` no longer touches collection.internal_metadata +at all. Both guard fields (last_dream_at and last_dream_document_count) are +written atomically in `process_dream` on successful completion — this preserves +the invariant that the baseline advances only when consolidation actually +happened, and prevents the in-flight stampede from false-advancing a guard. """ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest -from src import models, schemas +from src import schemas from src.deriver.enqueue import enqueue_dream class TestEnqueueDreamMetadataShape: - """Regression tests for Finding 3 code-level: `last_dream_at` relocation.""" - @pytest.mark.asyncio - async def test_update_data_omits_last_dream_at(self): - """`enqueue_dream` must NOT write `last_dream_at`. - - Moved to completion (in `process_dream`) so duplicate enqueues can't - reset the 8-hour guard and failed dreams don't falsely advance it. - """ - # Mock a Collection with empty dream metadata so the read-modify-write - # in enqueue_dream has something to merge into. - mock_collection = MagicMock(spec=models.Collection) - mock_collection.internal_metadata = {} - + async def test_enqueue_does_not_touch_collection_metadata(self): + """`enqueue_dream` must not call update_collection_internal_metadata.""" with ( patch( "src.deriver.enqueue.crud.update_collection_internal_metadata", @@ -37,11 +27,7 @@ class TestEnqueueDreamMetadataShape: patch( "src.deriver.enqueue.crud.get_collection", new_callable=AsyncMock, - return_value=mock_collection, - ), - # Short-circuit the dedup / insert paths so we only exercise the - # metadata update call. db_session.scalar returns False for both - # the in-progress and pending checks; db_session.execute is a no-op. + ) as mock_get_collection, patch( "src.deriver.enqueue.tracked_db", ) as mock_db_ctx, @@ -49,6 +35,7 @@ class TestEnqueueDreamMetadataShape: mock_session = AsyncMock() mock_session.scalar = AsyncMock(return_value=False) mock_session.execute = AsyncMock() + mock_session.commit = AsyncMock() mock_db_ctx.return_value.__aenter__.return_value = mock_session await enqueue_dream( @@ -56,23 +43,17 @@ class TestEnqueueDreamMetadataShape: observer="alice", observed="bob", dream_type=schemas.DreamType.OMNI, - document_count=42, session_name=None, ) - assert ( - mock_update.called - ), "update_collection_internal_metadata must be called" - call_kwargs = mock_update.call_args.kwargs - update_data = call_kwargs["update_data"] - dream_metadata = update_data["dream"] - - assert ( - "last_dream_document_count" in dream_metadata - ), "last_dream_document_count should still be written at enqueue" - assert dream_metadata["last_dream_document_count"] == 42 - - assert "last_dream_at" not in dream_metadata, ( - "last_dream_at must NOT be written at enqueue — it now writes " - "at dream completion in process_dream (orchestrator.py)." + assert not mock_update.called, ( + "enqueue_dream must not write to collection.internal_metadata; " + "guard fields advance atomically in process_dream on success." ) + assert not mock_get_collection.called, ( + "enqueue_dream must not need to load the collection — it no " + "longer touches dream metadata." + ) + assert ( + mock_session.execute.called + ), "enqueue_dream must still insert the QueueItem row." diff --git a/tests/dreamer/test_dream_scheduler.py b/tests/dreamer/test_dream_scheduler.py index f9b3c50b..4db81d74 100644 --- a/tests/dreamer/test_dream_scheduler.py +++ b/tests/dreamer/test_dream_scheduler.py @@ -285,138 +285,6 @@ class TestCancelDreamsForObserved: assert key_ws2 in dream_scheduler.pending_dreams -class TestDocumentCountAtExecutionTime: - """Regression tests for Bug #2: Stale document count used in metadata update. - - Previously, the document count was captured when the dream was scheduled - (at check_and_schedule_dream time), then used 60 minutes later when the - dream actually executed. This caused incorrect metadata if documents were - added during the wait period. - - Now, execute_dream queries the current document count at execution time. - """ - - @pytest.mark.asyncio - async def test_execute_dream_queries_document_count_at_execution( - self, dream_scheduler: DreamScheduler - ): - """execute_dream should query current document count, not use a stale value. - - This test verifies that execute_dream fetches the document count fresh - from the database at execution time rather than using a pre-captured value. - - The key architectural change was: - - OLD: schedule_dream(document_count) -> _delayed_dream(document_count) -> execute_dream(document_count) - - NEW: schedule_dream() -> _delayed_dream() -> execute_dream() queries count internally - - We verify this by mocking the database to return a specific count and - checking that enqueue_dream receives that count. - """ - from contextlib import asynccontextmanager - from unittest.mock import MagicMock - - from src import models - from src.schemas import ( - ResolvedConfiguration, - ResolvedDreamConfiguration, - ResolvedPeerCardConfiguration, - ResolvedReasoningConfiguration, - ResolvedSummaryConfiguration, - ) - - workspace_name = "test_workspace" - observer = "bob" - observed = "bob" - session_name = "test_session" - - # The document count that the database will return - CURRENT_DOC_COUNT = 42 - - # Track what document_count is passed to enqueue_dream - captured_document_count: int | None = None - - async def capture_enqueue_dream( - _ws_name: str, - observer: str, # pyright: ignore[reportUnusedParameter] - observed: str, # pyright: ignore[reportUnusedParameter] - dream_type: Any, # pyright: ignore[reportUnusedParameter] - document_count: int, - session_name: str, # pyright: ignore[reportUnusedParameter] - ) -> None: - nonlocal captured_document_count - captured_document_count = document_count - - # Create mock database session that returns our test data - mock_session = MagicMock() - mock_workspace = MagicMock(spec=models.Workspace) - mock_db_session = MagicMock(spec=models.Session) - - # Mock scalar to return session_name for first call, document count for second - scalar_call_count = 0 - - async def mock_scalar(_stmt: Any) -> str | int: - nonlocal scalar_call_count - scalar_call_count += 1 - if scalar_call_count == 1: - return session_name # First call gets session_name from documents - else: - return CURRENT_DOC_COUNT # Second call gets document count - - mock_session.scalar = mock_scalar - - @asynccontextmanager - async def mock_tracked_db(_: str | None = None): - yield mock_session - - with ( - patch( - "src.dreamer.dream_scheduler.tracked_db", - mock_tracked_db, - ), - patch( - "src.deriver.enqueue.enqueue_dream", - side_effect=capture_enqueue_dream, - ), - patch( - "src.crud.get_session", - return_value=mock_db_session, - ), - patch( - "src.crud.get_workspace", - return_value=mock_workspace, - ), - patch( - "src.utils.config_helpers.get_configuration", - return_value=ResolvedConfiguration( - reasoning=ResolvedReasoningConfiguration(enabled=True), - peer_card=ResolvedPeerCardConfiguration(use=True, create=True), - summary=ResolvedSummaryConfiguration( - enabled=True, - messages_per_short_summary=10, - messages_per_long_summary=20, - ), - dream=ResolvedDreamConfiguration(enabled=True), - ), - ), - ): - # Execute the dream - await dream_scheduler.execute_dream( - workspace_name, - DreamType.OMNI, - observer=observer, - observed=observed, - ) - - # Verify that execute_dream queried the document count (2 scalar calls) - assert ( - scalar_call_count == 2 - ), "Should have queried session_name and document count" - - # Verify that enqueue_dream received the CURRENT document count (42), - # proving that execute_dream queries the count at execution time - assert captured_document_count == CURRENT_DOC_COUNT - - class TestThresholdFilter: """Regression tests for Finding 2: threshold must count only explicit-level docs. diff --git a/tests/dreamer/test_dreamer_integration.py b/tests/dreamer/test_dreamer_integration.py index f64728ed..5f60196a 100644 --- a/tests/dreamer/test_dreamer_integration.py +++ b/tests/dreamer/test_dreamer_integration.py @@ -18,7 +18,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.deriver.enqueue import enqueue_dream -from src.dreamer.dream_scheduler import DreamScheduler, set_dream_scheduler +from src.dreamer.dream_scheduler import ( + DreamScheduler, + check_and_schedule_dream, + set_dream_scheduler, +) from src.dreamer.orchestrator import DreamResult, process_dream from src.schemas import ( DreamType, @@ -144,29 +148,36 @@ class TestLastDreamAtCompletionWrite: ) @pytest.mark.asyncio - async def test_completion_preserves_last_dream_document_count( + async def test_completion_writes_guard_pair_atomically( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """Completion write must NOT drop sibling keys in the `dream` sub-object. + """Completion writes last_dream_at AND last_dream_document_count together. - `update_collection_internal_metadata` uses a top-level JSONB `||` merge, - which replaces the whole `"dream"` key. Without a read-modify-write, the - completion write (`last_dream_at`) would wipe `last_dream_document_count` - that enqueue set — causing the next `check_and_schedule_dream` to read 0 - as the baseline and let a fresh dream trigger after the 8h guard expires - even without any new explicit documents. + Both guard fields advance only on successful consolidation, recomputed + inside the row-locked RMW block so the pair stays coherent. Baseline + reflects the actual explicit-doc count at completion, not a stale + enqueue-time snapshot. """ workspace, peer = sample_data - # Pre-seed collection as if enqueue_dream already wrote the baseline. collection = models.Collection( observer=peer.name, observed=peer.name, workspace_name=workspace.name, - internal_metadata={"dream": {"last_dream_document_count": 42}}, + internal_metadata={}, ) db_session.add(collection) + for i in range(7): + db_session.add( + models.Document( + content=f"explicit {i}", + level="explicit", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + ) + ) await db_session.commit() await db_session.refresh(collection) @@ -184,10 +195,9 @@ class TestLastDreamAtCompletionWrite: dream_meta = await _get_dream_metadata(db_session, collection) assert "last_dream_at" in dream_meta, "last_dream_at must be written" - assert dream_meta.get("last_dream_document_count") == 42, ( - "last_dream_document_count from enqueue must be preserved across " - "the completion write — top-level JSONB || would drop it without " - "the read-modify-write in process_dream." + assert dream_meta.get("last_dream_document_count") == 7, ( + "last_dream_document_count must equal the current explicit-doc count " + "at completion time; both guard fields advance together." ) @pytest.mark.asyncio @@ -225,38 +235,33 @@ class TestLastDreamAtCompletionWrite: ) -class TestEnqueueDreamPreservesSiblings: - """Regression test for the symmetric sibling-drop bug in enqueue_dream. +class TestEnqueueDreamLeavesMetadataAlone: + """enqueue_dream must not touch collection.internal_metadata["dream"]. - `update_collection_internal_metadata` uses a top-level JSONB `||` merge, - so writing `{"dream": {"last_dream_document_count": N}}` without a prior - read-modify-write would replace the whole `"dream"` subkey and drop - `last_dream_at` (written by process_dream on the prior completion). - - Symmetric to the orchestrator fix in c8fe40a (and the - test_completion_preserves_last_dream_document_count test above). + After the Loop 4 fix, the guard fields advance only on successful + completion in process_dream. enqueue_dream should preserve whatever + metadata is already on the collection (e.g. a prior completion's + timestamp and baseline) and add nothing of its own. """ @pytest.mark.asyncio - async def test_enqueue_preserves_last_dream_at( + async def test_enqueue_does_not_modify_dream_metadata( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """enqueue_dream must merge into existing dream metadata, not replace it. - - Pre-seeds `last_dream_at` (as if a prior dream completed), calls - `enqueue_dream`, and asserts both: - 1. `last_dream_document_count` is written (the new baseline), and - 2. `last_dream_at` is preserved (the sibling key from the prior completion). - """ workspace, peer = sample_data - prior_timestamp = "2026-04-17T12:00:00+00:00" + prior_metadata = { + "dream": { + "last_dream_at": "2026-04-17T12:00:00+00:00", + "last_dream_document_count": 99, + } + } collection = models.Collection( observer=peer.name, observed=peer.name, workspace_name=workspace.name, - internal_metadata={"dream": {"last_dream_at": prior_timestamp}}, + internal_metadata=prior_metadata, ) db_session.add(collection) await db_session.commit() @@ -267,30 +272,23 @@ class TestEnqueueDreamPreservesSiblings: observer=collection.observer, observed=collection.observed, dream_type=DreamType.OMNI, - document_count=77, session_name=None, ) dream_meta = await _get_dream_metadata(db_session, collection) - assert ( - dream_meta.get("last_dream_document_count") == 77 - ), "enqueue_dream must write last_dream_document_count as the new baseline" - assert dream_meta.get("last_dream_at") == prior_timestamp, ( - "last_dream_at from a prior completion must be preserved across " - "the enqueue write — top-level JSONB || would drop it without the " - "read-modify-write in enqueue_dream." + assert dream_meta == prior_metadata["dream"], ( + "enqueue_dream must leave dream metadata untouched; the guard fields " + "advance only at completion." ) class TestExecuteDreamSessionFilter: """Regression test for the session lookup asymmetry in execute_dream. - The baseline count query in execute_dream filters to `level == "explicit"` - (symmetric with check_and_schedule_dream). The session_name lookup must - filter the same way — otherwise the session is picked from a derived doc - and can disagree with the document set the count is measuring, producing - a dream scoped to a session that wasn't even in the triggering document - cohort. + The session_name lookup filters to `level == "explicit"`, symmetric with + check_and_schedule_dream's count query. Otherwise a derived doc could win + ORDER BY created_at DESC and the dream would be scoped to a session that + wasn't in the triggering document cohort. """ @pytest.mark.asyncio @@ -307,7 +305,8 @@ class TestExecuteDreamSessionFilter: Without the explicit filter on the session lookup, the newer deductive doc's session_name (B) would be returned. With the filter, A is - returned — matching the explicit-only count query immediately below. + returned — matching the explicit-only count query in + check_and_schedule_dream. """ workspace, peer = sample_data @@ -360,8 +359,6 @@ class TestExecuteDreamSessionFilter: db_session.add(deductive_doc) await db_session.commit() - # Capture kwargs passed to enqueue_dream (we don't want to actually - # run it — just verify which session_name execute_dream picks). captured_kwargs: dict[str, Any] = {} async def capture_enqueue_dream( @@ -370,7 +367,6 @@ class TestExecuteDreamSessionFilter: observer: str, observed: str, dream_type: Any, - document_count: int, session_name: str, ) -> None: captured_kwargs.update( @@ -379,7 +375,6 @@ class TestExecuteDreamSessionFilter: "observer": observer, "observed": observed, "dream_type": dream_type, - "document_count": document_count, "session_name": session_name, } ) @@ -433,7 +428,171 @@ class TestExecuteDreamSessionFilter: f"that the count query ignores — the dream would be scoped to a " f"session that wasn't in the triggering cohort." ) - # Sanity: the count must reflect explicit-only too (baseline symmetry). + + +class TestGuardPairCoherence: + """Loop 4 coherence tests for the invariant preserved by the atomic pair + write and the in-flight stampede defense. + + Invariant: From the moment a dream is scheduled until it completes or + fails, no second dream may be enqueued for the same + (workspace, observer, observed) — and the baseline count advances only + when consolidation actually happened. + """ + + @pytest_asyncio.fixture + async def _scheduler(self): + DreamScheduler.reset_singleton() + scheduler = DreamScheduler() + set_dream_scheduler(scheduler) + with ( + patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED", True), + patch("src.dreamer.dream_scheduler.settings.DREAM.DOCUMENT_THRESHOLD", 50), + patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED_TYPES", ["omni"]), + ): + yield scheduler + DreamScheduler.reset_singleton() + + @pytest.mark.asyncio + async def test_pending_queue_item_blocks_second_schedule( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + _scheduler: DreamScheduler, + ): + """In-flight window: pending QueueItem must block a second schedule. + + Walks the stampede timeline: enqueue fires a dream, more explicit + docs arrive past the threshold again, but check_and_schedule_dream + sees the pending queue row and returns False — no second QueueItem. + """ + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + for i in range(50): + db_session.add( + models.Document( + content=f"explicit {i}", + level="explicit", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + ) + ) + await db_session.commit() + await db_session.refresh(collection) + + await enqueue_dream( + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + dream_type=DreamType.OMNI, + session_name=None, + ) + + pending_q = select(models.QueueItem).where( + models.QueueItem.task_type == "dream", + models.QueueItem.processed == False, # noqa: E712 + models.QueueItem.workspace_name == workspace.name, + ) + pending_rows = (await db_session.execute(pending_q)).scalars().all() + assert len(pending_rows) == 1, ( + "enqueue_dream must insert exactly one pending dream QueueItem " + "(baseline for the stampede test)." + ) + + for i in range(50, 100): + db_session.add( + models.Document( + content=f"explicit {i}", + level="explicit", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + ) + ) + await db_session.commit() + await db_session.refresh(collection) + + scheduled = await check_and_schedule_dream(db_session, collection) + + assert scheduled is False, ( + "check_and_schedule_dream must return False while a dream is " + "pending in the queue — the in-flight window must not admit a " + "second schedule regardless of how many explicit docs arrive." + ) + pending_rows_after = (await db_session.execute(pending_q)).scalars().all() + assert len(pending_rows_after) == 1, ( + "No second QueueItem may be inserted while the first is pending. " + f"Found {len(pending_rows_after)} pending rows." + ) + + @pytest.mark.asyncio + async def test_silent_failure_allows_retry_on_same_corpus( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + _scheduler: DreamScheduler, + ): + """Failed dream (run_dream returns None) leaves both guard fields + untouched, so check_and_schedule_dream re-schedules on the same + corpus instead of silently consuming the baseline. + """ + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + for i in range(50): + db_session.add( + models.Document( + content=f"explicit {i}", + level="explicit", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + ) + ) + await db_session.commit() + await db_session.refresh(collection) + + payload = DreamPayload( + dream_type=DreamType.OMNI, + observer=peer.name, + observed=peer.name, + ) + with patch( + "src.dreamer.orchestrator.run_dream", + new=AsyncMock(return_value=None), + ): + await process_dream(payload, workspace.name) + + dream_meta = await _get_dream_metadata(db_session, collection) + assert dream_meta.get("last_dream_document_count", 0) == 0, ( + "Failed dream must not advance last_dream_document_count; " + "pre-Loop-4 the baseline was consumed at enqueue time and a " + "silent failure would lock out retries on the same corpus." + ) assert ( - captured_kwargs["document_count"] == 1 - ), "document_count must reflect explicit-only (1), not total (2)." + "last_dream_at" not in dream_meta + ), "Failed dream must not advance last_dream_at either." + + with patch.object( + _scheduler, "schedule_dream", new_callable=AsyncMock + ) as mock_schedule: + scheduled = await check_and_schedule_dream(db_session, collection) + + assert scheduled is True, ( + "After a silent failure both guards should still allow the " + "same-corpus retry — 50 explicit docs ≥ threshold, no prior " + "last_dream_at, no pending queue item." + ) + assert mock_schedule.called, "schedule_dream must be invoked on the retry path." diff --git a/tests/routes/test_workspaces.py b/tests/routes/test_workspaces.py index f1ce97d6..031e5a90 100644 --- a/tests/routes/test_workspaces.py +++ b/tests/routes/test_workspaces.py @@ -575,32 +575,20 @@ def test_delete_workspace_after_session_deletion(client: TestClient): @pytest.mark.asyncio -async def test_schedule_dream_passes_explicit_only_document_count( +async def test_schedule_dream_invokes_enqueue_dream( client: TestClient, db_session: AsyncSession, sample_data: tuple[Workspace, Peer], ): - """POST /schedule_dream must count only explicit-level docs. + """POST /schedule_dream forwards observer/observed/dream_type to enqueue_dream. - Regression test for the third caller of `enqueue_dream`. The manual - schedule_dream endpoint computes `document_count` and passes it to - `enqueue_dream`, where it becomes `last_dream_document_count` — the - baseline that `check_and_schedule_dream` and `execute_dream` subtract - against to compute the NEW-doc delta. - - Both of those callers filter to `level == "explicit"` (see Loop 2's - fixes). If this route doesn't match, the next auto-scheduled dream - sees a negative/suppressed delta because the baseline was inflated by - deductive/inductive docs. - - Seeds 5 explicit + 10 deductive (total 15) and asserts the count - passed to enqueue_dream is 5, not 15. + After Loop 4, the manual schedule_dream route no longer touches the + baseline count — the orchestrator writes both guard fields atomically on + successful completion. The route's job shrinks to forwarding the dream + request. """ workspace, peer = sample_data - # Seed collection + mixed-level documents. Must commit so the - # endpoint's DB session sees them (separate tracked_db path under - # some configs; safe regardless). collection = models.Collection( observer=peer.name, observed=peer.name, @@ -608,34 +596,8 @@ async def test_schedule_dream_passes_explicit_only_document_count( internal_metadata={}, ) db_session.add(collection) - await db_session.flush() - - for _ in range(5): - db_session.add( - models.Document( - content="explicit doc", - level="explicit", - workspace_name=workspace.name, - observer=peer.name, - observed=peer.name, - ) - ) - for _ in range(10): - db_session.add( - models.Document( - content="deductive doc", - level="deductive", - workspace_name=workspace.name, - observer=peer.name, - observed=peer.name, - ) - ) await db_session.commit() - # Capture kwargs passed to enqueue_dream (patched at the router's - # import site — the route does `from src.deriver.enqueue import - # enqueue_dream` at module load, so patching the symbol on the - # router module is what actually intercepts the call). captured: dict[str, Any] = {} async def fake_enqueue_dream(*args: Any, **kwargs: Any) -> None: @@ -660,8 +622,9 @@ async def test_schedule_dream_passes_explicit_only_document_count( assert response.status_code == 204, response.text assert "kwargs" in captured, "enqueue_dream was not called" - assert captured["kwargs"]["document_count"] == 5, ( - "schedule_dream must pass explicit-only count (5), " - f"got {captured['kwargs']['document_count']} — the unfiltered " - "count would be 15 and would inflate last_dream_document_count." + assert captured["kwargs"]["observer"] == peer.name + assert captured["kwargs"]["observed"] == peer.name + assert "document_count" not in captured["kwargs"], ( + "Loop 4: enqueue_dream no longer accepts document_count; the baseline " + "is written atomically with last_dream_at in process_dream." )