fix(dreamer): atomic guard-pair write + in-flight stampede defense
Loop 4 response to Vineeth's CHANGES_REQUESTED on PR #573. The pre-Loop-4 enqueue-time write of last_dream_document_count was serving double duty: rate limiter AND stampede latch. By arming the 8h guard the moment a dream entered the pipeline, it implicitly blocked a second dream from being scheduled during the in-flight window. Loop 3 relocated the last_dream_at write to completion without moving its sibling baseline, splitting the semantic pair and exposing the latch role that had lived only in Vineeth's head. Invariant (now pinned to check_and_schedule_dream's docstring): 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. Changes: - enqueue_dream: remove the last_dream_document_count write entirely and drop the document_count parameter. enqueue no longer touches dream metadata; the implicit stampede latch is replaced by an explicit queue-backed defense. - process_dream: extend the existing row-locked RMW to write both guard fields atomically. Current explicit-doc count is recomputed inside the locked block (not carried on DreamPayload) so the pair reflects the actual consolidation moment. - check_and_schedule_dream: query QueueItem for pending dreams on this collection's work_unit_keys (mirrors uq_queue_dream_pending_work_unit_key) before arming a timer. Uses queue state as source of truth rather than reflecting it into metadata. - Tests: two new coherence tests under TestGuardPairCoherence — test_pending_queue_item_blocks_second_schedule walks the stampede timeline, test_silent_failure_allows_retry_on_same_corpus verifies failed dreams don't consume the baseline. Existing tests updated to the new contract.
This commit is contained in:
lilyplasticlabs
parent
0172cd1254
commit
d24958dd7c
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue