From 0a7bde73416e48102564276738843d1b47f88f1c Mon Sep 17 00:00:00 2001 From: lilyplasticlabs Date: Fri, 17 Apr 2026 18:43:35 -0400 Subject: [PATCH] fix(dreamer): session lookup symmetry + row lock on dream metadata RMW - dream_scheduler.py: explicit-level filter on execute_dream session lookup (baseline and session pick must agree on the same document set) - crud.collection.get_collection: optional with_for_update flag for callers that need serialized read-modify-write on internal_metadata - enqueue.py + orchestrator.py: pass with_for_update=True on the RMW reads to close the TOCTOU between concurrent enqueue and completion writes Follow-up filed for jsonb_set-based nested updates (docs/factory/backlog/). Co-Authored-By: Claude Opus 4.7 (1M context) --- src/crud/collection.py | 22 +++ src/deriver/enqueue.py | 2 + src/dreamer/dream_scheduler.py | 6 +- src/dreamer/orchestrator.py | 2 + tests/dreamer/test_dreamer_integration.py | 167 +++++++++++++++++++++- 5 files changed, 197 insertions(+), 2 deletions(-) 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 db3432d6..d79350f2 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -521,11 +521,13 @@ async def enqueue_dream( # 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 diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index 8caf5310..3334c742 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -166,7 +166,10 @@ class DreamScheduler: 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 + # 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) @@ -174,6 +177,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) diff --git a/src/dreamer/orchestrator.py b/src/dreamer/orchestrator.py index 3857706a..763eca4a 100644 --- a/src/dreamer/orchestrator.py +++ b/src/dreamer/orchestrator.py @@ -330,6 +330,7 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p # 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. now_iso = datetime.now(timezone.utc).isoformat() async with tracked_db("dream.last_dream_at_write") as db: collection = await crud.get_collection( @@ -337,6 +338,7 @@ DREAM: {payload.dream_type} documents for {workspace_name}/{payload.observer}/{p workspace_name, observer=payload.observer, observed=payload.observed, + with_for_update=True, ) dream_meta = dict(collection.internal_metadata.get("dream", {})) dream_meta["last_dream_at"] = now_iso diff --git a/tests/dreamer/test_dreamer_integration.py b/tests/dreamer/test_dreamer_integration.py index 0d6f613b..f64728ed 100644 --- a/tests/dreamer/test_dreamer_integration.py +++ b/tests/dreamer/test_dreamer_integration.py @@ -18,8 +18,16 @@ 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.orchestrator import DreamResult, process_dream -from src.schemas import DreamType +from src.schemas import ( + DreamType, + ResolvedConfiguration, + ResolvedDreamConfiguration, + ResolvedPeerCardConfiguration, + ResolvedReasoningConfiguration, + ResolvedSummaryConfiguration, +) from src.utils.queue_payload import DreamPayload @@ -272,3 +280,160 @@ class TestEnqueueDreamPreservesSiblings: "the enqueue write — top-level JSONB || would drop it without the " "read-modify-write in enqueue_dream." ) + + +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. + """ + + @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 immediately below. + """ + 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() + + # 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( + workspace_name: str, + *, + observer: str, + observed: str, + dream_type: Any, + document_count: int, + session_name: str, + ) -> None: + captured_kwargs.update( + { + "workspace_name": workspace_name, + "observer": observer, + "observed": observed, + "dream_type": dream_type, + "document_count": document_count, + "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." + ) + # Sanity: the count must reflect explicit-only too (baseline symmetry). + assert ( + captured_kwargs["document_count"] == 1 + ), "document_count must reflect explicit-only (1), not total (2)."