diff --git a/src/crud/collection.py b/src/crud/collection.py index f0496d50..2790cec1 100644 --- a/src/crud/collection.py +++ b/src/crud/collection.py @@ -81,6 +81,7 @@ async def get_collection( *, observer: str, observed: str, + with_for_update: bool = False, ) -> models.Collection: """ Get a collection by observer/observed for a workspace. @@ -90,6 +91,11 @@ async def get_collection( workspace_name: Name of the workspace observer: Name of the observing peer (owns the collection) observed: Name of the observed peer + with_for_update: If True, acquire a row-level lock (SELECT ... FOR UPDATE) + on the collection. Bypasses the cache so the lock is actually held + by the current transaction. Callers using this flag must wrap the + read and subsequent write in the same transaction (the lock is + released on commit/rollback). Returns: The collection if found @@ -97,6 +103,22 @@ async def get_collection( Raises: ResourceNotFoundException: If the collection does not exist """ + if with_for_update: + # Row-lock path: go direct to DB (skip cache) so the FOR UPDATE lock + # is actually acquired on the row in the current transaction. The + # cached dict path would return without issuing SELECT ... FOR UPDATE. + stmt = ( + select(models.Collection) + .where(models.Collection.workspace_name == workspace_name) + .where(models.Collection.observer == observer) + .where(models.Collection.observed == observed) + .with_for_update() + ) + collection = await db.scalar(stmt) + if collection is None: + raise ResourceNotFoundException("Collection not found") + return collection + data = await _fetch_collection(db, workspace_name, observer, observed) if data is None: raise ResourceNotFoundException("Collection not found") diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index 71b00a2b..cbd032c2 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -1,5 +1,4 @@ import logging -from datetime import datetime, timezone from typing import Any, Literal from sqlalchemy import exists, insert, select @@ -436,27 +435,26 @@ 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 + are written atomically in process_dream on successful completion. + 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: Current document count for metadata update 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, @@ -467,11 +465,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( @@ -491,7 +484,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( @@ -512,25 +504,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) - now_iso = datetime.now(timezone.utc).isoformat() - await crud.update_collection_internal_metadata( - db_session, - workspace_name, - observer, - observed, - update_data={ - "dream": { - "last_dream_document_count": document_count, - "last_dream_at": now_iso, - } - }, - ) - # 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 1a339d21..eb2e2177 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,13 +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 session and get current document count async with tracked_db("dream_session_lookup") as db: stmt = ( select(models.Document.session_name) @@ -174,6 +172,7 @@ class DreamScheduler: models.Document.workspace_name == workspace_name, models.Document.observer == observer, models.Document.observed == observed, + models.Document.level == "explicit", ) .order_by(models.Document.created_at.desc()) .limit(1) @@ -186,14 +185,6 @@ class DreamScheduler: ) return - # Get current document count at execution time (not stale from scheduling) - count_stmt = select(func.count(models.Document.id)).where( - models.Document.workspace_name == workspace_name, - models.Document.observer == observer, - models.Document.observed == observed, - ) - current_document_count = int(await db.scalar(count_stmt) or 0) - session = await crud.get_session( db, workspace_name=workspace_name, session_name=session_name ) @@ -212,7 +203,6 @@ class DreamScheduler: observer=observer, observed=observed, dream_type=dream_type, - document_count=current_document_count, session_name=session_name, ) @@ -231,13 +221,18 @@ async def check_and_schedule_dream( collection: models.Collection, ) -> bool: """ - Check if a collection has reached the document threshold and schedule a timer-based dream. + 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. Document threshold is reached + 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 @@ -249,21 +244,21 @@ 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") - # Count current documents in the collection + # Count explicit-level docs only: dreamer output (deductive/inductive/ + # contradiction) would inflate the threshold and create a feedback loop. count_stmt = select(func.count(models.Document.id)).where( models.Document.workspace_name == collection.workspace_name, models.Document.observer == collection.observer, models.Document.observed == collection.observed, + models.Document.level == "explicit", ) - current_document_count = int(await db.scalar(count_stmt) or 0) + current_explicit_count = int(await db.scalar(count_stmt) or 0) - # Calculate documents added since last dream - documents_since_last_dream = current_document_count - last_dream_document_count + documents_since_last_dream = current_explicit_count - last_dream_document_count logger.debug( "Dream check", @@ -271,16 +266,14 @@ async def check_and_schedule_dream( "workspace_name": collection.workspace_name, "observer": collection.observer, "observed": collection.observed, - "current_document_count": current_document_count, + "current_explicit_count": current_explicit_count, "last_dream_document_count": last_dream_document_count, "documents_since_last_dream": documents_since_last_dream, "document_threshold": settings.DREAM.DOCUMENT_THRESHOLD, }, ) - # 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) @@ -299,11 +292,43 @@ async def check_and_schedule_dream( f"Invalid last_dream_at timestamp: {last_dream_at}, error: {e}" ) + # Queue is source of truth for in-flight dreams; mirrors + # uq_queue_dream_pending_work_unit_key. + 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 000d45a6..d09e5b67 100644 --- a/src/dreamer/orchestrator.py +++ b/src/dreamer/orchestrator.py @@ -17,11 +17,13 @@ import logging import time import uuid from dataclasses import dataclass +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 @@ -323,6 +325,34 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p + f"duration={result.total_duration_ms:.0f}ms" ) + # Both guard fields advance together only on successful consolidation. + now_iso = datetime.now(timezone.utc).isoformat() + async with tracked_db("dream.guard_pair_write") as db: + collection = await crud.get_collection( + db, + workspace_name, + observer=payload.observer, + 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, + payload.observer, + payload.observed, + update_data={"dream": dream_meta}, + ) + except Exception as e: logger.error( f"Error processing dream task {payload.dream_type} for {payload.observer}/{payload.observed}: {str(e)}", diff --git a/src/routers/workspaces.py b/src/routers/workspaces.py index 90530e92..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,21 +222,11 @@ async def schedule_dream( observed = request.observed if request.observed is not None else request.observer dream_type = request.dream_type - # Count documents in the collection - count_stmt = select(func.count(models.Document.id)).where( - models.Document.workspace_name == workspace_id, - models.Document.observer == observer, - models.Document.observed == observed, - ) - 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 new file mode 100644 index 00000000..bea5c256 --- /dev/null +++ b/tests/deriver/test_enqueue_dream.py @@ -0,0 +1,59 @@ +"""Regression tests for `enqueue_dream` metadata write shape. + +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, patch + +import pytest + +from src import schemas +from src.deriver.enqueue import enqueue_dream + + +class TestEnqueueDreamMetadataShape: + @pytest.mark.asyncio + 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", + new_callable=AsyncMock, + ) as mock_update, + patch( + "src.deriver.enqueue.crud.get_collection", + new_callable=AsyncMock, + ) as mock_get_collection, + patch( + "src.deriver.enqueue.tracked_db", + ) as mock_db_ctx, + ): + 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( + workspace_name="test_workspace", + observer="alice", + observed="bob", + dream_type=schemas.DreamType.OMNI, + session_name=None, + ) + + 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 06f904c1..4db81d74 100644 --- a/tests/dreamer/test_dream_scheduler.py +++ b/tests/dreamer/test_dream_scheduler.py @@ -4,8 +4,14 @@ from typing import Any from unittest.mock import AsyncMock, patch import pytest +from sqlalchemy.ext.asyncio import AsyncSession -from src.dreamer.dream_scheduler import DreamScheduler, set_dream_scheduler +from src import models +from src.dreamer.dream_scheduler import ( + DreamScheduler, + check_and_schedule_dream, + set_dream_scheduler, +) from src.schemas import DreamType from src.utils.work_unit import construct_work_unit_key @@ -279,136 +285,140 @@ class TestCancelDreamsForObserved: assert key_ws2 in dream_scheduler.pending_dreams -class TestDocumentCountAtExecutionTime: - """Regression tests for Bug #2: Stale document count used in metadata update. +class TestThresholdFilter: + """Regression tests for Finding 2: threshold must count only explicit-level docs. - 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. + Previously the threshold counted all documents in a collection, including + dreamer output (deductive/inductive/contradiction). This created a feedback + loop where each dream's output inflated the trigger for the next dream. + The fix filters the count to `level == "explicit"` only. """ - @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. + @pytest.fixture(autouse=True) + def _pin_dream_config(self): + """Pin DOCUMENT_THRESHOLD=50 and ENABLED_TYPES=['omni'] for this class. - 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. + These tests assume the default thresholds; a developer's local env + (e.g. DREAM_DOCUMENT_THRESHOLD=5 for faster manual testing) would + otherwise invalidate the 30/60/10 fixtures below. Scoped to this + class only — do NOT widen; other tests may have different assumptions. """ - from contextlib import asynccontextmanager - from unittest.mock import MagicMock + with ( + patch("src.dreamer.dream_scheduler.settings.DREAM.DOCUMENT_THRESHOLD", 50), + patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED_TYPES", ["omni"]), + ): + yield - from src import models - from src.schemas import ( - ResolvedConfiguration, - ResolvedDreamConfiguration, - ResolvedPeerCardConfiguration, - ResolvedReasoningConfiguration, - ResolvedSummaryConfiguration, + async def _make_collection( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ) -> models.Collection: + """Helper: create a Collection in the test workspace with no dream metadata.""" + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + await db_session.commit() + return collection + + async def _insert_doc( + self, + db_session: AsyncSession, + collection: models.Collection, + level: str, + ) -> None: + """Helper: insert one Document at the given level.""" + db_session.add( + models.Document( + content="test", + level=level, + workspace_name=collection.workspace_name, + observer=collection.observer, + observed=collection.observed, + ) ) - workspace_name = "test_workspace" - observer = "bob" - observed = "bob" - session_name = "test_session" + @pytest.mark.asyncio + async def test_mixed_levels_below_explicit_threshold( + self, + dream_scheduler: DreamScheduler, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """30 explicit + 40 deductive + 10 inductive → should NOT trigger. - # The document count that the database will return - CURRENT_DOC_COUNT = 42 + Total doc count = 80 (would trigger under the buggy unfiltered count), + but explicit count = 30 < threshold 50, so the correct behavior is to + NOT schedule a dream. This is the core regression: the fix must reject + this scenario. + """ + collection = await self._make_collection(db_session, sample_data) + for _ in range(30): + await self._insert_doc(db_session, collection, "explicit") + for _ in range(40): + await self._insert_doc(db_session, collection, "deductive") + for _ in range(10): + await self._insert_doc(db_session, collection, "inductive") + await db_session.commit() - # Track what document_count is passed to enqueue_dream - captured_document_count: int | None = None + with patch.object(dream_scheduler, "schedule_dream", new_callable=AsyncMock): + scheduled = await check_and_schedule_dream(db_session, collection) - 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 + assert scheduled is False, ( + "Threshold should filter on explicit level only — dreamer output " + "(deductive/inductive) must not count toward the trigger." + ) - # 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) + @pytest.mark.asyncio + async def test_explicit_only_at_threshold( + self, + dream_scheduler: DreamScheduler, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """60 explicit + 0 derived → should trigger (60 ≥ threshold 50).""" + collection = await self._make_collection(db_session, sample_data) + for _ in range(60): + await self._insert_doc(db_session, collection, "explicit") + await db_session.commit() - # Mock scalar to return session_name for first call, document count for second - scalar_call_count = 0 + with patch.object( + dream_scheduler, "schedule_dream", new_callable=AsyncMock + ) as mock_schedule: + scheduled = await check_and_schedule_dream(db_session, collection) - 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 + assert scheduled is True + assert mock_schedule.called, "schedule_dream should fire when threshold met" - mock_session.scalar = mock_scalar + @pytest.mark.asyncio + async def test_contradiction_excluded_from_count( + self, + dream_scheduler: DreamScheduler, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Contradiction-level docs are dreamer output — must not count. - @asynccontextmanager - async def mock_tracked_db(_: str | None = None): - yield mock_session + 100 contradictions + 10 explicit → explicit=10 < threshold=50, no trigger. + Confirms the positive `== "explicit"` filter excludes contradiction by + construction (same as deductive/inductive). + """ + collection = await self._make_collection(db_session, sample_data) + for _ in range(100): + await self._insert_doc(db_session, collection, "contradiction") + for _ in range(10): + await self._insert_doc(db_session, collection, "explicit") + await db_session.commit() - 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, - ) + with patch.object(dream_scheduler, "schedule_dream", new_callable=AsyncMock): + scheduled = await check_and_schedule_dream(db_session, collection) - # 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 + assert scheduled is False class TestEnqueueCancelsDreamsCorrectly: diff --git a/tests/dreamer/test_dreamer_integration.py b/tests/dreamer/test_dreamer_integration.py new file mode 100644 index 00000000..5f60196a --- /dev/null +++ b/tests/dreamer/test_dreamer_integration.py @@ -0,0 +1,598 @@ +"""Integration tests for the dream completion write. + +Finding 3 (code-level) relocates `last_dream_at` from enqueue time to +dream-completion time (in `process_dream`). These tests exercise the real +Postgres JSONB merge via `tracked_db` to verify the write lands in the +collection's internal_metadata on successful dreams — and critically, +does NOT land on failures or exceptions. +""" + +from datetime import datetime, timedelta +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest +import pytest_asyncio +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from src import models +from src.deriver.enqueue import enqueue_dream +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, + ResolvedConfiguration, + ResolvedDreamConfiguration, + ResolvedPeerCardConfiguration, + ResolvedReasoningConfiguration, + ResolvedSummaryConfiguration, +) +from src.utils.queue_payload import DreamPayload + + +@pytest_asyncio.fixture +async def seeded_collection( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], +) -> models.Collection: + """Create a Collection with an empty dream metadata dict.""" + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + await db_session.commit() + await db_session.refresh(collection) + return collection + + +def _make_dream_result() -> DreamResult: + """Build a minimal non-null DreamResult for happy-path tests.""" + return DreamResult( + run_id="test_run_01", + specialists_run=["deduction", "induction"], + deduction_success=True, + induction_success=True, + surprisal_enabled=False, + surprisal_conclusion_count=0, + total_iterations=3, + total_duration_ms=1234.5, + input_tokens=100, + output_tokens=50, + ) + + +async def _get_dream_metadata( + db_session: AsyncSession, collection: models.Collection +) -> dict[str, Any]: + """Re-fetch collection and return its internal_metadata['dream'] dict (or {}).""" + await db_session.refresh(collection) + stmt = select(models.Collection).where(models.Collection.id == collection.id) + refreshed = (await db_session.execute(stmt)).scalar_one() + dream_meta: dict[str, Any] = refreshed.internal_metadata.get("dream", {}) + return dream_meta + + +class TestLastDreamAtCompletionWrite: + """Regression tests for Finding 3: `last_dream_at` written at completion.""" + + @pytest.mark.asyncio + async def test_happy_path_writes_last_dream_at( + self, + db_session: AsyncSession, + seeded_collection: models.Collection, + ): + """Non-null DreamResult → `last_dream_at` is set in internal_metadata.""" + payload = DreamPayload( + dream_type=DreamType.OMNI, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + ) + + with patch( + "src.dreamer.orchestrator.run_dream", + new=AsyncMock(return_value=_make_dream_result()), + ): + await process_dream(payload, seeded_collection.workspace_name) + + dream_meta = await _get_dream_metadata(db_session, seeded_collection) + assert ( + "last_dream_at" in dream_meta + ), "process_dream must write last_dream_at when run_dream returns a result" + # Must be a tz-aware UTC ISO timestamp. A naive datetime.now().isoformat() + # would pass a loose "T in string" check but corrupt the 8h guard math + # against tz-aware now() comparisons downstream. + parsed = datetime.fromisoformat(dream_meta["last_dream_at"]) + assert ( + parsed.tzinfo is not None + ), f"last_dream_at must be timezone-aware, got {dream_meta['last_dream_at']!r}" + assert parsed.utcoffset() == timedelta( + 0 + ), f"last_dream_at must be UTC, got offset {parsed.utcoffset()}" + + @pytest.mark.asyncio + async def test_failure_path_leaves_last_dream_at_null( + self, + db_session: AsyncSession, + seeded_collection: models.Collection, + ): + """run_dream returns None → `last_dream_at` stays absent. + + Lenient success criteria: the guard only advances on completion of a + non-null DreamResult. Failed runs (None return) must not count. + """ + payload = DreamPayload( + dream_type=DreamType.OMNI, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + ) + + with patch( + "src.dreamer.orchestrator.run_dream", + new=AsyncMock(return_value=None), + ): + await process_dream(payload, seeded_collection.workspace_name) + + dream_meta = await _get_dream_metadata(db_session, seeded_collection) + assert "last_dream_at" not in dream_meta, ( + "last_dream_at must NOT be written when run_dream returns None " + "(failed dream). The guard should not falsely advance." + ) + + @pytest.mark.asyncio + async def test_completion_writes_guard_pair_atomically( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Completion writes last_dream_at AND last_dream_document_count together. + + 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 + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + 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) + + payload = DreamPayload( + dream_type=DreamType.OMNI, + observer=collection.observer, + observed=collection.observed, + ) + + with patch( + "src.dreamer.orchestrator.run_dream", + new=AsyncMock(return_value=_make_dream_result()), + ): + await process_dream(payload, collection.workspace_name) + + 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") == 7, ( + "last_dream_document_count must equal the current explicit-doc count " + "at completion time; both guard fields advance together." + ) + + @pytest.mark.asyncio + async def test_exception_path_leaves_last_dream_at_null( + self, + db_session: AsyncSession, + seeded_collection: models.Collection, + ): + """run_dream raises → `last_dream_at` stays absent. + + `process_dream` catches exceptions (logs + marks task processed without + re-raising) so the queue worker doesn't get stuck retrying. The guard + write must not happen in the exception path — it's inside the + `if result is not None` block, which never executes if an exception + bypassed the assignment. + """ + payload = DreamPayload( + dream_type=DreamType.OMNI, + observer=seeded_collection.observer, + observed=seeded_collection.observed, + ) + + with patch( + "src.dreamer.orchestrator.run_dream", + new=AsyncMock(side_effect=RuntimeError("simulated specialist crash")), + ): + # process_dream swallows exceptions internally; no re-raise expected + await process_dream(payload, seeded_collection.workspace_name) + + dream_meta = await _get_dream_metadata(db_session, seeded_collection) + assert "last_dream_at" not in dream_meta, ( + "last_dream_at must NOT be written when run_dream raises. " + "process_dream swallows the exception but the guard write must " + "not occur." + ) + + +class TestEnqueueDreamLeavesMetadataAlone: + """enqueue_dream must not touch collection.internal_metadata["dream"]. + + 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_does_not_modify_dream_metadata( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + 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=prior_metadata, + ) + db_session.add(collection) + await db_session.commit() + await db_session.refresh(collection) + + await enqueue_dream( + workspace_name=workspace.name, + observer=collection.observer, + observed=collection.observed, + dream_type=DreamType.OMNI, + session_name=None, + ) + + dream_meta = await _get_dream_metadata(db_session, collection) + 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 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 + async def test_session_name_picked_from_latest_explicit_doc( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Latest explicit session wins even when a newer deductive doc exists. + + Seeds: + - Session A (older): one explicit-level Document + - Session B (newer): one deductive-level Document (dreamer output) + + 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 in + check_and_schedule_dream. + """ + workspace, peer = sample_data + + # Pre-create the collection so crud.get_collection inside enqueue_dream + # finds something (process_dream's baseline write is not exercised here; + # we're only asserting the kwargs passed to enqueue_dream). + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + + # Two sessions: A (older), B (newer). Insert A first so its created_at + # is strictly earlier than B's. + session_a = models.Session(name="session_a", workspace_name=workspace.name) + db_session.add(session_a) + await db_session.commit() + await db_session.refresh(session_a) + + session_b = models.Session(name="session_b", workspace_name=workspace.name) + db_session.add(session_b) + await db_session.commit() + await db_session.refresh(session_b) + + # Older explicit doc in session A. + explicit_doc = models.Document( + content="explicit observation", + level="explicit", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + session_name=session_a.name, + ) + db_session.add(explicit_doc) + await db_session.commit() + + # Newer deductive doc in session B. Without the explicit filter on the + # session lookup, this doc's session_name (B) would win on ORDER BY + # created_at DESC — even though the count query ignores it. + deductive_doc = models.Document( + content="deductive observation", + level="deductive", + workspace_name=workspace.name, + observer=peer.name, + observed=peer.name, + session_name=session_b.name, + ) + db_session.add(deductive_doc) + await db_session.commit() + + captured_kwargs: dict[str, Any] = {} + + async def capture_enqueue_dream( + workspace_name: str, + *, + observer: str, + observed: str, + dream_type: Any, + session_name: str, + ) -> None: + captured_kwargs.update( + { + "workspace_name": workspace_name, + "observer": observer, + "observed": observed, + "dream_type": dream_type, + "session_name": session_name, + } + ) + + # Fresh scheduler instance; ENABLED patched so execute_dream runs. + DreamScheduler.reset_singleton() + scheduler = DreamScheduler() + set_dream_scheduler(scheduler) + + try: + with ( + patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED", True), + patch( + "src.deriver.enqueue.enqueue_dream", + side_effect=capture_enqueue_dream, + ), + 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), + ), + ), + ): + await scheduler.execute_dream( + workspace.name, + DreamType.OMNI, + observer=peer.name, + observed=peer.name, + ) + finally: + DreamScheduler.reset_singleton() + + assert captured_kwargs, ( + "enqueue_dream must be called — execute_dream returned early, " + "likely because the session lookup returned no rows (check that " + "the explicit filter matches at least one doc in the fixture)." + ) + assert captured_kwargs["session_name"] == session_a.name, ( + f"Session lookup must filter to level=='explicit' to match the " + f"baseline count query. Got session_name=" + f"{captured_kwargs['session_name']!r}, expected {session_a.name!r} " + f"(the older session with the only explicit doc). Picking " + f"{session_b.name!r} means the session came from a derived doc " + f"that the count query ignores — the dream would be scoped to a " + f"session that wasn't in the triggering cohort." + ) + + +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 ( + "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 36bc8312..031e5a90 100644 --- a/tests/routes/test_workspaces.py +++ b/tests/routes/test_workspaces.py @@ -1,9 +1,12 @@ from typing import Any +from unittest.mock import AsyncMock, patch import pytest from fastapi.testclient import TestClient from nanoid import generate as generate_nanoid +from sqlalchemy.ext.asyncio import AsyncSession +from src import models from src.models import Peer, Workspace @@ -569,3 +572,59 @@ def test_delete_workspace_after_session_deletion(client: TestClient): # Now workspace deletion should succeed response = client.delete(f"/v3/workspaces/{workspace_name}") assert response.status_code == 202 + + +@pytest.mark.asyncio +async def test_schedule_dream_invokes_enqueue_dream( + client: TestClient, + db_session: AsyncSession, + sample_data: tuple[Workspace, Peer], +): + """POST /schedule_dream forwards observer/observed/dream_type to enqueue_dream. + + 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 + + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata={}, + ) + db_session.add(collection) + await db_session.commit() + + captured: dict[str, Any] = {} + + async def fake_enqueue_dream(*args: Any, **kwargs: Any) -> None: + captured["args"] = args + captured["kwargs"] = kwargs + + with ( + patch("src.routers.workspaces.settings.DREAM.ENABLED", True), + patch( + "src.routers.workspaces.enqueue_dream", + new=AsyncMock(side_effect=fake_enqueue_dream), + ), + ): + response = client.post( + f"/v3/workspaces/{workspace.name}/schedule_dream", + json={ + "observer": peer.name, + "observed": peer.name, + "dream_type": "omni", + }, + ) + + assert response.status_code == 204, response.text + assert "kwargs" in captured, "enqueue_dream was not called" + 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." + )