From af46190efff541733fd0a92a3ab1c2045277e5c1 Mon Sep 17 00:00:00 2001 From: Ulysse Pence Date: Thu, 27 Aug 2026 17:02:44 -0200 Subject: [PATCH] Move dream scheduling and stale-claim cleanup into the API --- src/backlog.py | 141 ++++- src/config.py | 6 + src/crud/__init__.py | 2 + src/crud/deriver.py | 58 +- src/crud/representation.py | 13 +- src/deriver/deriver.py | 1 - src/deriver/enqueue.py | 21 - src/deriver/queue_manager.py | 107 ---- src/dreamer/__init__.py | 8 +- src/dreamer/dream_scheduler.py | 437 ++----------- src/reconciler/scheduler.py | 31 +- tests/conftest.py | 2 + tests/crud/test_deriver_backlog.py | 9 +- tests/crud/test_representation_manager.py | 4 - tests/dreamer/test_dream_scheduler.py | 708 +++++++++------------- tests/dreamer/test_dreamer_integration.py | 79 +-- tests/routes/test_scopes.py | 2 +- tests/test_db_resilience.py | 113 ---- 18 files changed, 561 insertions(+), 1181 deletions(-) diff --git a/src/backlog.py b/src/backlog.py index 45a731d2..4335bea3 100644 --- a/src/backlog.py +++ b/src/backlog.py @@ -3,6 +3,7 @@ import asyncio import contextlib import logging +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, cast @@ -13,12 +14,21 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models from src.config import settings from src.dependencies import tracked_db +from src.dreamer import execute_dream +from src.schemas import DreamType from src.telemetry.prometheus import metrics as prometheus_metrics from src.utils.work_unit import construct_work_unit_key logger = logging.getLogger(__name__) -DREAM_POLL_INTERVAL_SECONDS = 300 + +@dataclass +class DueDream: + workspace_name: str + observer: str + observed: str + dream_type: DreamType + documents_since_last_dream: int class BacklogMetricsPoller: @@ -75,6 +85,8 @@ class BacklogMetricsPoller: continue async def _refresh(self) -> None: + due_dreams: list[DueDream] | None = None + async with tracked_db("backlog_metrics", read_only=True) as db: backlog = await crud.get_deriver_backlog(db) prometheus_metrics.prometheus_metrics.set_deriver_backlog( @@ -88,15 +100,43 @@ class BacklogMetricsPoller: ) if self._dream_poll_due(): - prometheus_metrics.prometheus_metrics.set_dreams_pending( - count=await count_pending_dreams(db) + due_dreams = await find_due_dreams(db) + + if due_dreams is None: + return + + prometheus_metrics.prometheus_metrics.set_dreams_pending(count=len(due_dreams)) + + for dream in due_dreams: + try: + await execute_dream( + dream.workspace_name, + dream.dream_type, + observer=dream.observer, + observed=dream.observed, + trigger_reason="document_threshold", + delay_reason="idle_timeout", + documents_since_last_dream_at_schedule=dream.documents_since_last_dream, + document_threshold=settings.DREAM.DOCUMENT_THRESHOLD, ) + except Exception as e: + logger.error( + "Failed to enqueue dream for %s/%s/%s: %s", + dream.workspace_name, + dream.observer, + dream.observed, + e, + ) + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(e) def _dream_poll_due(self) -> bool: now = datetime.now(timezone.utc) if self._next_dream_poll is not None and now < self._next_dream_poll: return False - self._next_dream_poll = now + timedelta(seconds=DREAM_POLL_INTERVAL_SECONDS) + self._next_dream_poll = now + timedelta( + seconds=settings.DREAM.POLL_INTERVAL_SECONDS + ) return True @@ -108,9 +148,21 @@ async def _seconds_since_last_vector_sync(db: AsyncSession) -> float: return (datetime.now(timezone.utc) - newest).total_seconds() -async def count_pending_dreams(db: AsyncSession) -> int: - if not settings.DREAM.ENABLED: - return 0 +async def find_due_dreams(db: AsyncSession) -> list[DueDream]: + """Collections whose next dream is due, read-only. + + A dream is due when the collection has gained DOCUMENT_THRESHOLD explicit + documents since its last dream, has been quiet for IDLE_TIMEOUT_MINUTES, is + past the MIN_HOURS_BETWEEN_DREAMS gate, and has no dream queued or attempted + since its newest explicit document. + """ + dream_types = [ + DreamType(dream_type) + for dream_type in settings.DREAM.ENABLED_TYPES + if dream_type == DreamType.OMNI.value + ] + if not settings.DREAM.ENABLED or not dream_types: + return [] explicit_counts = ( select( @@ -118,6 +170,7 @@ async def count_pending_dreams(db: AsyncSession) -> int: models.Document.observer, models.Document.observed, func.count(models.Document.id).label("explicit_count"), + func.max(models.Document.created_at).label("newest_created_at"), ) .where(models.Document.level == "explicit") .group_by( @@ -136,6 +189,7 @@ async def count_pending_dreams(db: AsyncSession) -> int: models.Collection.observed, models.Collection.internal_metadata, func.coalesce(explicit_counts.c.explicit_count, 0), + explicit_counts.c.newest_created_at, ).outerjoin( explicit_counts, (models.Collection.workspace_name == explicit_counts.c.workspace_name) @@ -146,7 +200,8 @@ async def count_pending_dreams(db: AsyncSession) -> int: ).all() now = datetime.now(timezone.utc) - due: list[str] = [] + idle_cutoff = now - timedelta(minutes=settings.DREAM.IDLE_TIMEOUT_MINUTES) + candidates: dict[str, tuple[DueDream, datetime]] = {} for row in rows: workspace_name = cast(str, row[0]) @@ -154,6 +209,7 @@ async def count_pending_dreams(db: AsyncSession) -> int: observed = cast(str, row[2]) internal_metadata = cast("dict[str, Any] | None", row[3]) explicit_count = cast(int, row[4]) + newest_created_at = cast("datetime | None", row[5]) dream_metadata: dict[str, Any] = (internal_metadata or {}).get("dream", {}) since_last_dream = explicit_count - int( @@ -162,41 +218,60 @@ async def count_pending_dreams(db: AsyncSession) -> int: if since_last_dream < settings.DREAM.DOCUMENT_THRESHOLD: continue + if newest_created_at is None or newest_created_at > idle_cutoff: + continue + last_dream_at = cast("str | None", dream_metadata.get("last_dream_at")) if last_dream_at and _within_min_hours_gate(last_dream_at, now): continue - for dream_type in settings.DREAM.ENABLED_TYPES: - due.append( - construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": observer, - "observed": observed, - "dream_type": dream_type, - }, - ) + for dream_type in dream_types: + work_unit_key = construct_work_unit_key( + workspace_name, + { + "task_type": "dream", + "observer": observer, + "observed": observed, + "dream_type": dream_type.value, + }, + ) + candidates[work_unit_key] = ( + DueDream( + workspace_name=workspace_name, + observer=observer, + observed=observed, + dream_type=dream_type, + documents_since_last_dream=since_last_dream, + ), + newest_created_at, ) - if not due: - return 0 + if not candidates: + return [] - already_queued = set( - ( - await db.execute( - select(models.QueueItem.work_unit_key).where( - models.QueueItem.task_type == "dream", - ~models.QueueItem.processed, - models.QueueItem.work_unit_key.in_(due), - ) + attempt_rows = ( + await db.execute( + select( + models.QueueItem.work_unit_key, + func.max(models.QueueItem.created_at), ) + .where( + models.QueueItem.task_type == "dream", + models.QueueItem.work_unit_key.in_(candidates.keys()), + ) + .group_by(models.QueueItem.work_unit_key) ) - .scalars() - .all() - ) + ).all() + newest_attempts: dict[str, datetime] = { + cast(str, row[0]): cast(datetime, row[1]) for row in attempt_rows + } - return len([key for key in due if key not in already_queued]) + return [ + due + for work_unit_key, (due, newest_created_at) in candidates.items() + if work_unit_key not in newest_attempts + or newest_attempts[work_unit_key] < newest_created_at + ] def _within_min_hours_gate(last_dream_at: str, now: datetime) -> bool: diff --git a/src/config.py b/src/config.py index 0115b598..04adbc6b 100644 --- a/src/config.py +++ b/src/config.py @@ -918,6 +918,11 @@ class DeriverSettings(HonchoSettings): int, Field(default=30 * 24 * 3600, gt=0) ] = 30 * 24 * 3600 # 30 days default + # Spacing between runs of the processed-queue-item cleanup + QUEUE_CLEANUP_INTERVAL_SECONDS: Annotated[int, Field(default=12 * 3600, ge=1)] = ( + 12 * 3600 + ) + @staticmethod def _MODEL_CONFIG_DEFAULT() -> ConfiguredModelSettings: # Minimal default: transport + model only. Any other knobs would merge @@ -1352,6 +1357,7 @@ class DreamSettings(HonchoSettings): ENABLED: bool = True DOCUMENT_THRESHOLD: Annotated[int, Field(default=50, gt=0, le=1000)] = 50 IDLE_TIMEOUT_MINUTES: Annotated[int, Field(default=60, gt=0, le=1440)] = 60 + POLL_INTERVAL_SECONDS: Annotated[int, Field(default=300, ge=1)] = 300 MIN_HOURS_BETWEEN_DREAMS: Annotated[int, Field(default=8, gt=0, le=72)] = 8 ENABLED_TYPES: list[str] = ["omni"] diff --git a/src/crud/__init__.py b/src/crud/__init__.py index 1129648c..8d196ba0 100644 --- a/src/crud/__init__.py +++ b/src/crud/__init__.py @@ -4,6 +4,7 @@ from .collection import ( update_collection_internal_metadata, ) from .deriver import ( + cleanup_stale_work_units, get_deriver_backlog, get_deriver_status, get_queue_status, @@ -109,6 +110,7 @@ __all__ = [ "get_or_create_collection", "update_collection_internal_metadata", # Deriver + "cleanup_stale_work_units", "get_deriver_backlog", "get_deriver_status", "get_queue_status", diff --git a/src/crud/deriver.py b/src/crud/deriver.py index ca6d4d19..8c65e89e 100644 --- a/src/crud/deriver.py +++ b/src/crud/deriver.py @@ -1,14 +1,15 @@ from collections.abc import Sequence -from datetime import timedelta +from datetime import datetime, timedelta, timezone from logging import getLogger from typing import Any -from sqlalchemy import ColumnElement, Select, case, func, or_, select +from sqlalchemy import ColumnElement, Select, case, delete, func, or_, select from sqlalchemy.engine import Row from sqlalchemy.ext.asyncio import AsyncSession from src import models, schemas from src.config import settings +from src.dependencies import tracked_db logger = getLogger(__name__) @@ -109,29 +110,46 @@ def representation_batch_threshold_clause( def unclaimed_work_unit_clause( work_unit_key: ColumnElement[str], - *, - tolerate_stale_claims: bool = False, ) -> ColumnElement[bool]: - claim = select(models.ActiveQueueSession.id).where( - models.ActiveQueueSession.work_unit_key == work_unit_key + return ( + ~select(models.ActiveQueueSession.id) + .where(models.ActiveQueueSession.work_unit_key == work_unit_key) + .exists() ) - if tolerate_stale_claims: - claim = claim.where( - models.ActiveQueueSession.last_updated - >= func.now() - - timedelta(minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES) + +async def cleanup_stale_work_units() -> None: + """Clean up stale work units""" + async with tracked_db("cleanup_stale_work_units") as db: + cutoff = datetime.now(timezone.utc) - timedelta( + minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES ) - return ~claim.exists() + stale_ids = ( + ( + await db.execute( + select(models.ActiveQueueSession.id) + .where(models.ActiveQueueSession.last_updated < cutoff) + .order_by(models.ActiveQueueSession.last_updated) + .with_for_update(skip_locked=True) + ) + ) + .scalars() + .all() + ) + + # Delete only the records we successfully got locks for + if stale_ids: + await db.execute( + delete(models.ActiveQueueSession).where( + models.ActiveQueueSession.id.in_(stale_ids) + ) + ) + await db.commit() async def get_deriver_backlog(db: AsyncSession) -> schemas.DeriverBacklog: - """Count outstanding deriver work across the whole database. - - Read-only, and deliberately unfiltered by workspace: the caller wants the - service-wide total, not one tenant's share of it. - """ + """Count outstanding deriver work across the whole database""" token_stats = ( select( models.QueueItem.work_unit_key, @@ -161,11 +179,7 @@ async def get_deriver_backlog(db: AsyncSession) -> schemas.DeriverBacklog: token_stats, work_units.c.work_unit_key == token_stats.c.work_unit_key, ) - .where( - unclaimed_work_unit_clause( - work_units.c.work_unit_key, tolerate_stale_claims=True - ) - ) + .where(unclaimed_work_unit_clause(work_units.c.work_unit_key)) ) threshold_clause = representation_batch_threshold_clause( diff --git a/src/crud/representation.py b/src/crud/representation.py index 6fafb842..d8cdb3e6 100644 --- a/src/crud/representation.py +++ b/src/crud/representation.py @@ -12,9 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, exceptions, models, schemas from src.config import settings from src.dependencies import tracked_db -from src.dreamer.dream_scheduler import check_and_schedule_dream from src.embedding_client import EmbeddingTokenLimitError, embedding_client -from src.schemas import ResolvedConfiguration from src.telemetry.events import EmbeddingCallPurpose from src.telemetry.logging import accumulate_metric from src.utils.formatting import format_datetime_utc @@ -65,7 +63,6 @@ class RepresentationManager: message_ids: list[int], session_name: str, message_created_at: datetime.datetime, - message_level_configuration: ResolvedConfiguration, ) -> crud.CreateDocumentsResult: """ Save Representation objects to the collection as a set of documents. @@ -133,7 +130,6 @@ class RepresentationManager: message_ids, session_name, message_created_at, - message_level_configuration, ) create_document_duration = (time.perf_counter() - create_document_start) * 1000 @@ -154,10 +150,9 @@ class RepresentationManager: message_ids: list[int], session_name: str, message_created_at: datetime.datetime, - message_level_configuration: ResolvedConfiguration, ) -> crud.CreateDocumentsResult: # get_or_create_collection already handles IntegrityError with rollback and a retry - collection = await crud.get_or_create_collection( + await crud.get_or_create_collection( db, self.workspace_name, observer=self.observer, @@ -203,12 +198,6 @@ class RepresentationManager: deduplicate=settings.DERIVER.DEDUPLICATE, ) - if message_level_configuration.dream.enabled: - try: - await check_and_schedule_dream(db, collection) - except Exception as e: - logger.warning(f"Failed to check dream scheduling: {e}") - return accepted_documents_result async def get_working_representation( diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py index f76c4d52..39e2b86e 100644 --- a/src/deriver/deriver.py +++ b/src/deriver/deriver.py @@ -222,7 +222,6 @@ async def process_representation_tasks_batch( message_ids, latest_message.session_name, latest_message.created_at, - message_level_configuration, ) ) agg_representation_result.exact_dup_existing_count += ( diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py index 2f70e787..a51e46e9 100644 --- a/src/deriver/enqueue.py +++ b/src/deriver/enqueue.py @@ -7,7 +7,6 @@ from sqlalchemy.ext.asyncio import AsyncSession from src import crud, models, schemas from src.config import settings from src.dependencies import tracked_db -from src.dreamer.dream_scheduler import get_dream_scheduler from src.exceptions import ValidationException from src.models import QueueItem from src.schemas import MessageConfiguration, ResolvedConfiguration @@ -31,26 +30,6 @@ async def enqueue(payload: list[dict[str, Any]]) -> None: payload: List of message payload dictionaries """ - # Cancel any pending dreams for affected collections since user is active again. - # This cancels dreams for all collections where observed=peer_name, which covers - # both self-observation and peer-to-peer observation cases. - dream_scheduler = get_dream_scheduler() - if dream_scheduler and payload: - cancelled_dreams: set[str] = set() - for message in payload: - workspace_name = message.get("workspace_name") - peer_name = message.get("peer_name") - if workspace_name and peer_name: - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, peer_name - ) - cancelled_dreams.update(cancelled) - - if cancelled_dreams: - logger.info( - f"Cancelled {len(cancelled_dreams)} pending dreams due to new activity" - ) - async with tracked_db("message_enqueue") as db_session: try: # Determine if batch or single processing diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index d40bed3d..03e1a30b 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -2,11 +2,9 @@ import asyncio import contextlib import random import signal -import time from asyncio import Task from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from datetime import datetime, timedelta, timezone from logging import getLogger from typing import Any, NamedTuple, cast @@ -29,12 +27,6 @@ from src.deriver.consumer import ( process_item, process_representation_batch, ) -from src.dreamer.dream_scheduler import ( - DreamScheduler, - check_and_schedule_dream, - get_dream_scheduler, - set_dream_scheduler, -) from src.models import QueueItem from src.schemas import ResolvedConfiguration from src.telemetry import prometheus_metrics @@ -131,28 +123,10 @@ class QueueManager: settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS ) - # Monotonic timestamp of the last stale-work-unit cleanup ATTEMPT. - # None -> the first poll always runs cleanup (recovers rows left stale - # by a crashed predecessor immediately). - self._last_stale_cleanup_attempt: float | None = None - # Jittered gate width (seconds) sampled ONCE per attempt, so the deadline - # for the next run is fixed when the timestamp is set rather than - # re-rolled on every poll (which would make the effective spacing a - # random walk and untestable at non-zero jitter ratios). - self._stale_cleanup_gate_seconds: float = 0.0 - # Initialize from settings self.workers: int = settings.DERIVER.WORKERS self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers) - # Get or create the singleton dream scheduler - existing_scheduler = get_dream_scheduler() - if existing_scheduler is None: - self.dream_scheduler: DreamScheduler = DreamScheduler() - set_dream_scheduler(self.dream_scheduler) - else: - self.dream_scheduler = existing_scheduler - # Initialize Sentry if enabled, using settings if settings.SENTRY.ENABLED: initialize_sentry( @@ -197,11 +171,6 @@ class QueueManager: ) logger.debug("Signal handlers registered") - try: - await self._reschedule_pending_dreams() - except Exception: - logger.exception("Failed to reschedule pending dreams at startup") - # Run the polling loop directly in this task logger.debug("Starting polling loop directly") try: @@ -210,28 +179,11 @@ class QueueManager: finally: await self.cleanup() - async def _reschedule_pending_dreams(self) -> None: - if not settings.DREAM.ENABLED: - return - - async with tracked_db("reschedule_pending_dreams") as db: - collections = (await db.execute(select(models.Collection))).scalars().all() - rescheduled = 0 - for collection in collections: - if await check_and_schedule_dream(db, collection): - rescheduled += 1 - - if rescheduled: - logger.info("Rescheduled %d dream(s) at startup", rescheduled) - async def shutdown(self, sig: signal.Signals) -> None: """Handle graceful shutdown""" logger.info(f"Received exit signal {sig.name}...") self.shutdown_event.set() - # Cancel all pending dreams - await self.dream_scheduler.shutdown() - if self.active_tasks: logger.info( f"Waiting for {len(self.active_tasks)} active tasks to complete..." @@ -267,64 +219,6 @@ class QueueManager: # Polling and Scheduling # ########################## - async def _maybe_cleanup_stale_work_units(self) -> None: - """Run stale-work-unit cleanup at most once per (jittered) interval. - - Staleness is a minutes-timescale condition (STALE_SESSION_TIMEOUT_MINUTES), - but the polling loop fires on a seconds timescale on every deriver - instance — running cleanup unconditionally per poll multiplies into - unnecessary write transactions. Gate it locally: - concurrent cleaners on other instances remain safe via FOR UPDATE SKIP - LOCKED, so no cross-instance coordination is required, and the jittered - gate (sampled once per attempt) keeps instances from re-synchronizing - their cleanup runs. The gate tracks the last ATTEMPT (set before - running), so a failing cleanup waits a full interval instead of retrying - every poll against a DB that is already struggling. An interval of 0 - preserves run-every-poll behavior. - """ - interval = settings.DERIVER.STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS - if ( - interval > 0.0 - and self._last_stale_cleanup_attempt is not None - and time.monotonic() - self._last_stale_cleanup_attempt - < self._stale_cleanup_gate_seconds - ): - return - # Record the attempt and fix the next deadline before running, so the - # gate width is stable for this cycle and a failing cleanup still waits. - self._last_stale_cleanup_attempt = time.monotonic() - self._stale_cleanup_gate_seconds = self._jitter(interval) - await self.cleanup_stale_work_units() - - async def cleanup_stale_work_units(self) -> None: - """Clean up stale work units""" - async with tracked_db("cleanup_stale_work_units") as db: - cutoff = datetime.now(timezone.utc) - timedelta( - minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES - ) - - stale_ids = ( - ( - await db.execute( - select(models.ActiveQueueSession.id) - .where(models.ActiveQueueSession.last_updated < cutoff) - .order_by(models.ActiveQueueSession.last_updated) - .with_for_update(skip_locked=True) - ) - ) - .scalars() - .all() - ) - - # Delete only the records we successfully got locks for - if stale_ids: - await db.execute( - delete(models.ActiveQueueSession).where( - models.ActiveQueueSession.id.in_(stale_ids) - ) - ) - await db.commit() - async def get_and_claim_work_units(self) -> dict[str, str]: """ Get available work units that aren't being processed. @@ -531,7 +425,6 @@ class QueueManager: continue try: - await self._maybe_cleanup_stale_work_units() claimed_work_units = await self.get_and_claim_work_units() if claimed_work_units: if self._is_tenant_work(claimed_work_units): diff --git a/src/dreamer/__init__.py b/src/dreamer/__init__.py index bce5c95b..41a57f3e 100644 --- a/src/dreamer/__init__.py +++ b/src/dreamer/__init__.py @@ -1,11 +1,7 @@ -from .dream_scheduler import ( - check_and_schedule_dream, - get_dream_scheduler, -) +from .dream_scheduler import execute_dream from .orchestrator import process_dream __all__ = [ - "get_dream_scheduler", - "check_and_schedule_dream", + "execute_dream", "process_dream", ] diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py index 43bba428..bf686350 100644 --- a/src/dreamer/dream_scheduler.py +++ b/src/dreamer/dream_scheduler.py @@ -1,406 +1,71 @@ -import asyncio -import contextlib -from datetime import datetime, timezone from logging import getLogger -import sentry_sdk -from sqlalchemy import exists, func, select -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import select from src import models -from src.config import settings from src.dependencies import tracked_db from src.schemas import DreamType -from src.utils.work_unit import construct_work_unit_key, parse_work_unit_key logger = getLogger(__name__) -_dream_scheduler: "DreamScheduler | None" = None +async def execute_dream( + workspace_name: str, + dream_type: DreamType, + *, + observer: str, + observed: str, + trigger_reason: str | None = None, + delay_reason: str | None = None, + documents_since_last_dream_at_schedule: int | None = None, + document_threshold: int | None = None, +) -> None: + """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 + async with tracked_db("dream_session_lookup") as db: + stmt = ( + select(models.Document.session_name) + .where( + 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) + ) + session_name = await db.scalar(stmt) -def set_dream_scheduler(dream_scheduler: "DreamScheduler") -> None: - """Set the global dream scheduler reference.""" - global _dream_scheduler - _dream_scheduler = dream_scheduler - - -def get_dream_scheduler() -> "DreamScheduler | None": - """Get the global dream scheduler reference.""" - return _dream_scheduler - - -class DreamScheduler: - _instance: "DreamScheduler | None" = None - _initialized: bool = False - - def __new__(cls): - if cls._instance is None: - cls._instance = super().__new__(cls) - return cls._instance - - def __init__(self): - # Only initialize once - if not DreamScheduler._initialized: - self.pending_dreams: dict[str, asyncio.Task[None]] = {} - DreamScheduler._initialized = True - - @classmethod - def reset_singleton(cls) -> None: - """Reset the singleton instance. Only use this in tests.""" - cls._instance = None - cls._initialized = False - - async def schedule_dream( - self, - work_unit_key: str, - workspace_name: str, - delay_minutes: int, - dream_type: DreamType, - *, - observer: str, - observed: str, - trigger_reason: str | None = None, - delay_reason: str | None = None, - documents_since_last_dream_at_schedule: int | None = None, - document_threshold: int | None = None, - ) -> None: - """Schedule a dream for a collection after a delay. - - telemetry kwargs are captured at schedule time and threaded - through the queue payload so DreamRunEvent can attribute the dream - back to its scheduling context. - """ - if not settings.DREAM.ENABLED: + if not session_name: + logger.warning( + f"No documents found for {workspace_name}/{observer}/{observed}, skipping dream" + ) return - # Cancel any existing dream for this collection - await self.cancel_dream(work_unit_key) - - task = asyncio.create_task( - self._delayed_dream( - work_unit_key, - workspace_name, - delay_minutes, - dream_type, - observer=observer, - observed=observed, - trigger_reason=trigger_reason, - delay_reason=delay_reason, - documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, - document_threshold=document_threshold, - ) + session = await crud.get_session( + db, workspace_name=workspace_name, session_name=session_name ) - self.pending_dreams[work_unit_key] = task - task.add_done_callback(lambda t: self.pending_dreams.pop(work_unit_key, None)) + workspace = await crud.get_workspace(db, workspace_name=workspace_name) - async def cancel_dream(self, work_unit_key: str) -> bool: - """Cancel a pending dream. Returns True if a dream was cancelled.""" - if work_unit_key in self.pending_dreams: - task = self.pending_dreams.pop(work_unit_key) - task.cancel() - # Wait for the task to actually finish (including its done callback) - with contextlib.suppress(asyncio.CancelledError): - await task - return True - return False + configuration = get_configuration(None, session, workspace) - async def cancel_dreams_for_observed( - self, workspace_name: str, observed: str - ) -> set[str]: - """ - Cancel all pending dreams where the observed peer matches. - - This handles both self-observation (observer=observed) and peer-to-peer - observation (observer!=observed) dreams. - - Args: - workspace_name: The workspace to match - observed: The observed peer name to match - - Returns: - Set of work_unit_keys that were cancelled - """ - cancelled: set[str] = set() - - # Collect keys to cancel (can't modify dict while iterating) - keys_to_cancel: list[str] = [] - for work_unit_key in self.pending_dreams: - parsed = parse_work_unit_key(work_unit_key) - if parsed.workspace_name == workspace_name and parsed.observed == observed: - keys_to_cancel.append(work_unit_key) - - # Cancel each matching dream - for key in keys_to_cancel: - if await self.cancel_dream(key): - cancelled.add(key) - - return cancelled - - async def _delayed_dream( - self, - work_unit_key: str, - workspace_name: str, - delay_minutes: int, - dream_type: DreamType, - *, - observer: str, - observed: str, - trigger_reason: str | None = None, - delay_reason: str | None = None, - documents_since_last_dream_at_schedule: int | None = None, - document_threshold: int | None = None, - ) -> None: - try: - await asyncio.sleep(delay_minutes * 60) - - await self.execute_dream( - workspace_name, - dream_type, - observer=observer, - observed=observed, - trigger_reason=trigger_reason, - delay_reason=delay_reason, - documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, - document_threshold=document_threshold, - ) - logger.info("Executed dream for %s", work_unit_key) - - except asyncio.CancelledError: - logger.debug("Dream task cancelled for %s", work_unit_key) - except Exception as e: - logger.error("Error in delayed dream for %s: %s", work_unit_key, e) - if settings.SENTRY.ENABLED: - sentry_sdk.capture_exception(e) - - async def execute_dream( - self, - workspace_name: str, - dream_type: DreamType, - *, - observer: str, - observed: str, - trigger_reason: str | None = None, - delay_reason: str | None = None, - documents_since_last_dream_at_schedule: int | None = None, - document_threshold: int | None = None, - ) -> None: - """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 - - async with tracked_db("dream_session_lookup") as db: - stmt = ( - select(models.Document.session_name) - .where( - 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) - ) - session_name = await db.scalar(stmt) - - if not session_name: - logger.warning( - f"No documents found for {workspace_name}/{observer}/{observed}, skipping dream" - ) - return - - session = await crud.get_session( - db, workspace_name=workspace_name, session_name=session_name - ) - workspace = await crud.get_workspace(db, workspace_name=workspace_name) - - configuration = get_configuration(None, session, workspace) - - if not configuration.dream.enabled: - logger.debug( - f"Dreams disabled for {workspace_name}/{session_name}, skipping dream" - ) - return - - await enqueue_dream( - workspace_name, - observer=observer, - observed=observed, - dream_type=dream_type, - session_name=session_name, - trigger_reason=trigger_reason, - delay_reason=delay_reason, - documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, - document_threshold=document_threshold, - ) - - async def shutdown(self) -> None: - """Cancel all pending dreams during shutdown.""" - if self.pending_dreams: - logger.info(f"Cancelling {len(self.pending_dreams)} pending dreams...") - for task in self.pending_dreams.values(): - task.cancel() - await asyncio.gather(*self.pending_dreams.values(), return_exceptions=True) - self.pending_dreams.clear() - - -async def check_and_schedule_dream( - db: AsyncSession, - 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 pending in the queue for this collection (in-flight check) - 5. No dream is already scheduled for this collection - - Args: - db: Database session - collection: Collection model to check - - Returns: - True if a dream timer was scheduled, False otherwise - """ - if not settings.DREAM.ENABLED: - return False - - 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 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_explicit_count = int(await db.scalar(count_stmt) or 0) - - documents_since_last_dream = current_explicit_count - last_dream_document_count - - logger.debug( - "Dream check", - extra={ - "workspace_name": collection.workspace_name, - "observer": collection.observer, - "observed": collection.observed, - "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, - }, - ) - - if documents_since_last_dream >= settings.DREAM.DOCUMENT_THRESHOLD: - # capture *why* this schedule fired (threshold) and - # *how* it will fire (idle vs immediate). The two gates were - # intentionally split — collapsing them into a single trigger_reason - # would lose the scheduling semantics. - trigger_reason = "document_threshold" - delay_reason = ( - "idle_timeout" if settings.DREAM.IDLE_TIMEOUT_MINUTES > 0 else "immediate" - ) - if last_dream_at: - try: - last_dream_time = datetime.fromisoformat(last_dream_at) - hours_since_last_dream = ( - datetime.now(timezone.utc) - last_dream_time - ).total_seconds() / 3600 - - if hours_since_last_dream < settings.DREAM.MIN_HOURS_BETWEEN_DREAMS: - logger.debug( - f"Skipping dream for {collection.observer}/{collection.observed}: only {hours_since_last_dream:.1f} hours " - + f"since last dream (minimum: {settings.DREAM.MIN_HOURS_BETWEEN_DREAMS})" - ) - # delay_reason = "min_hours_gate" if we DID schedule, but - # we don't — return early. Telemetry only records dreams - # that actually fire. - return False - except (ValueError, TypeError) as e: - logger.warning( - 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: + if not configuration.dream.enabled: logger.debug( - "Skipping dream schedule for %s/%s: pending dream already in queue", - collection.observer, - collection.observed, + f"Dreams disabled for {workspace_name}/{session_name}, skipping dream" ) - return False + return - dream_scheduler = get_dream_scheduler() - if dream_scheduler: - for dream_type in enabled_dream_types: - dream_work_unit_key = construct_work_unit_key( - collection.workspace_name, - { - "task_type": "dream", - "observer": collection.observer, - "observed": collection.observed, - "dream_type": dream_type, - }, - ) - await dream_scheduler.schedule_dream( - dream_work_unit_key, - collection.workspace_name, - settings.DREAM.IDLE_TIMEOUT_MINUTES, - dream_type=DreamType(dream_type), - observer=collection.observer, - observed=collection.observed, - trigger_reason=trigger_reason, - delay_reason=delay_reason, - documents_since_last_dream_at_schedule=documents_since_last_dream, - document_threshold=settings.DREAM.DOCUMENT_THRESHOLD, - ) - logger.debug( - "Scheduled dream", - extra={ - "workspace_name": collection.workspace_name, - "observer": collection.observer, - "observed": collection.observed, - "documents_since_last_dream": documents_since_last_dream, - "document_threshold": settings.DREAM.DOCUMENT_THRESHOLD, - "dream_type": dream_type, - }, - ) - return True - - return False + await enqueue_dream( + workspace_name, + observer=observer, + observed=observed, + dream_type=dream_type, + session_name=session_name, + trigger_reason=trigger_reason, + delay_reason=delay_reason, + documents_since_last_dream_at_schedule=documents_since_last_dream_at_schedule, + document_threshold=document_threshold, + ) diff --git a/src/reconciler/scheduler.py b/src/reconciler/scheduler.py index 21f37b51..adf5340c 100644 --- a/src/reconciler/scheduler.py +++ b/src/reconciler/scheduler.py @@ -17,7 +17,7 @@ from pydantic import BaseModel from sqlalchemy import exists, select from sqlalchemy.exc import IntegrityError -from src import models +from src import crud, models from src.config import settings from src.dependencies import tracked_db from src.models import QueueItem @@ -34,9 +34,6 @@ class ReconcilerTask(BaseModel): interval_seconds: int -# Task intervals -QUEUE_CLEANUP_INTERVAL_SECONDS = 12 * 3600 # 12 hours - # Task registry - add new tasks here RECONCILER_TASKS: dict[str, ReconcilerTask] = { "sync_vectors": ReconcilerTask( @@ -47,7 +44,7 @@ RECONCILER_TASKS: dict[str, ReconcilerTask] = { "cleanup_queue": ReconcilerTask( name="cleanup_queue", work_unit_key="reconciler:cleanup_queue", - interval_seconds=QUEUE_CLEANUP_INTERVAL_SECONDS, + interval_seconds=settings.DERIVER.QUEUE_CLEANUP_INTERVAL_SECONDS, ), } @@ -97,6 +94,7 @@ class ReconcilerScheduler: self._shutdown_event: asyncio.Event = asyncio.Event() # Track next run time for each task self._next_run: dict[str, datetime] = {} + self._next_stale_cleanup: datetime | None = None ReconcilerScheduler._initialized = True @classmethod @@ -116,6 +114,7 @@ class ReconcilerScheduler: now = datetime.now(timezone.utc) for task_name, task in RECONCILER_TASKS.items(): self._next_run[task_name] = now + timedelta(seconds=task.interval_seconds) + self._next_stale_cleanup = now self._scheduler_task = asyncio.create_task(self._scheduler_loop()) logger.info( @@ -142,6 +141,7 @@ class ReconcilerScheduler: self._scheduler_task = None self._next_run.clear() + self._next_stale_cleanup = None logger.info("ReconcilerScheduler stopped") async def _scheduler_loop(self) -> None: @@ -166,6 +166,20 @@ class ReconcilerScheduler: # endregion await record_pending_embeddings_backlog() + if ( + self._next_stale_cleanup is not None + and now >= self._next_stale_cleanup + ): + try: + await crud.cleanup_stale_work_units() + except Exception as e: + logger.exception("Error cleaning up stale work units") + if settings.SENTRY.ENABLED: + sentry_sdk.capture_exception(e) + self._next_stale_cleanup = now + timedelta( + seconds=settings.DERIVER.STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS + ) + # Check each task and enqueue if due for task_name, task in RECONCILER_TASKS.items(): next_run = self._next_run.get(task_name, now) @@ -186,8 +200,11 @@ class ReconcilerScheduler: ) # Calculate sleep time until next task is due - if self._next_run: - next_task_time = min(self._next_run.values()) + next_run_times = list(self._next_run.values()) + if self._next_stale_cleanup is not None: + next_run_times.append(self._next_stale_cleanup) + if next_run_times: + next_task_time = min(next_run_times) sleep_seconds = max( 1.0, # At least 1 second to avoid busy loop (next_task_time - datetime.now(timezone.utc)).total_seconds(), diff --git a/tests/conftest.py b/tests/conftest.py index 090d5395..3da45152 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -954,6 +954,7 @@ def mock_tracked_db(request: pytest.FixtureRequest): # Use ExitStack (not a parenthesized `with`) to stay under CPython's # 20-statically-nested-block limit as this list grows. tracked_db_targets = [ + "src.backlog.tracked_db", "src.dependencies.tracked_db", "src.deriver.queue_manager.tracked_db", "src.deriver.consumer.tracked_db", @@ -969,6 +970,7 @@ def mock_tracked_db(request: pytest.FixtureRequest): "src.webhooks.webhook_delivery.tracked_db", "src.utils.agent_tools.tracked_db", "src.utils.search.tracked_db", + "src.crud.deriver.tracked_db", "src.crud.document.tracked_db", "src.crud.message.tracked_db", "src.reconciler.sync_vectors.tracked_db", diff --git a/tests/crud/test_deriver_backlog.py b/tests/crud/test_deriver_backlog.py index cc9b6de9..e61b4670 100644 --- a/tests/crud/test_deriver_backlog.py +++ b/tests/crud/test_deriver_backlog.py @@ -192,7 +192,7 @@ class TestDeriverBacklog: assert backlog.eligible_work_units == 0 - async def test_stale_claim_does_not_hide_work( + async def test_stale_claim_is_reaped_and_stops_hiding_work( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], @@ -221,9 +221,12 @@ class TestDeriverBacklog: ) await db_session.commit() - backlog = await crud.get_deriver_backlog(db_session) + assert (await crud.get_deriver_backlog(db_session)).eligible_work_units == 0 - assert backlog.eligible_work_units == 1 + await crud.cleanup_stale_work_units() + await db_session.commit() + + assert (await crud.get_deriver_backlog(db_session)).eligible_work_units == 1 async def test_processed_items_are_not_counted( self, diff --git a/tests/crud/test_representation_manager.py b/tests/crud/test_representation_manager.py index 3c2f57e4..a50993a6 100644 --- a/tests/crud/test_representation_manager.py +++ b/tests/crud/test_representation_manager.py @@ -516,7 +516,6 @@ class TestRepresentationManagerSave: message_ids=[1], session_name="session", message_created_at=datetime.now(timezone.utc), - message_level_configuration=_resolved_config(), ) assert len(saved.created_documents) == 1 @@ -574,7 +573,6 @@ class TestRepresentationManagerSave: message_ids=[1], session_name="session", message_created_at=datetime.now(timezone.utc), - message_level_configuration=_resolved_config(), ) assert len(saved.created_documents) == 1 @@ -627,7 +625,6 @@ class TestRepresentationManagerSave: message_ids=[1], session_name="session", message_created_at=datetime.now(timezone.utc), - message_level_configuration=_resolved_config(), ) assert len(saved.created_documents) == 0 @@ -682,7 +679,6 @@ class TestRepresentationManagerSave: message_ids=[1], session_name="session", message_created_at=datetime.now(timezone.utc), - message_level_configuration=_resolved_config(), ) mock_embed.assert_awaited_once_with( diff --git a/tests/dreamer/test_dream_scheduler.py b/tests/dreamer/test_dream_scheduler.py index 4db81d74..2d830c55 100644 --- a/tests/dreamer/test_dream_scheduler.py +++ b/tests/dreamer/test_dream_scheduler.py @@ -1,338 +1,50 @@ -"""Regression tests for dream scheduler bug fixes.""" +"""Tests for the API-side dream due-check and for enqueueing a due dream.""" -from typing import Any +import datetime from unittest.mock import AsyncMock, patch import pytest from sqlalchemy.ext.asyncio import AsyncSession from src import models -from src.dreamer.dream_scheduler import ( - DreamScheduler, - check_and_schedule_dream, - set_dream_scheduler, -) +from src.backlog import find_due_dreams +from src.dreamer.dream_scheduler import execute_dream from src.schemas import DreamType from src.utils.work_unit import construct_work_unit_key -@pytest.fixture -def dream_scheduler(): - """Create a fresh DreamScheduler instance for each test.""" - # Reset the singleton before each test - DreamScheduler.reset_singleton() - scheduler = DreamScheduler() - set_dream_scheduler(scheduler) - # Patch DREAM.ENABLED to True so tests work regardless of local config - with patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED", True): - yield scheduler - # Cleanup - DreamScheduler.reset_singleton() +def _now() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc) -class TestCancelDreamsForObserved: - """Regression tests for Bug #1: Peer-to-peer observation dreams not cancelled on activity. - - Previously, when a message arrived from peer Bob, only the self-observation dream - (observer=Bob, observed=Bob) was cancelled. Peer-to-peer observation dreams - (observer=Alice, observed=Bob) were NOT cancelled, allowing dreams to fire - during active conversation. - """ - - @pytest.mark.asyncio - async def test_cancel_self_observation_dream(self, dream_scheduler: DreamScheduler): - """Cancelling dreams for observed peer should cancel self-observation dreams.""" - workspace_name = "test_workspace" - peer_name = "bob" - - # Schedule a self-observation dream (observer=bob, observed=bob) - work_unit_key = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": peer_name, - "observed": peer_name, - "dream_type": "omni", - }, - ) - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - await dream_scheduler.schedule_dream( - work_unit_key, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer=peer_name, - observed=peer_name, - ) - - # Verify dream is pending - assert work_unit_key in dream_scheduler.pending_dreams - - # Cancel dreams for observed peer - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, peer_name - ) - - # Verify the dream was cancelled - assert work_unit_key in cancelled - assert work_unit_key not in dream_scheduler.pending_dreams - - @pytest.mark.asyncio - async def test_cancel_peer_to_peer_observation_dream( - self, dream_scheduler: DreamScheduler - ): - """Cancelling dreams for observed peer should also cancel peer-to-peer dreams. - - This is the core regression test for Bug #1: previously, if Alice was - observing Bob and Bob sent a message, the dream for (observer=Alice, - observed=Bob) would NOT be cancelled. - """ - workspace_name = "test_workspace" - observer = "alice" # Alice is watching Bob - observed = "bob" # Bob sends a message - - # Schedule a peer-to-peer observation dream (observer=alice, observed=bob) - work_unit_key = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": observer, - "observed": observed, - "dream_type": "omni", - }, - ) - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - await dream_scheduler.schedule_dream( - work_unit_key, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer=observer, - observed=observed, - ) - - # Verify dream is pending - assert work_unit_key in dream_scheduler.pending_dreams - - # When Bob sends a message, cancel all dreams where observed=bob - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, observed - ) - - # Verify the peer-to-peer dream was cancelled - assert work_unit_key in cancelled - assert work_unit_key not in dream_scheduler.pending_dreams - - @pytest.mark.asyncio - async def test_cancel_multiple_observers_same_observed( - self, dream_scheduler: DreamScheduler - ): - """When observed peer sends a message, ALL dreams observing them should cancel.""" - workspace_name = "test_workspace" - observed = "bob" - observers = ["alice", "charlie", "bob"] # Multiple observers including self - - work_unit_keys: list[str] = [] - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - for observer in observers: - work_unit_key = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": observer, - "observed": observed, - "dream_type": "omni", - }, - ) - work_unit_keys.append(work_unit_key) - - await dream_scheduler.schedule_dream( - work_unit_key, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer=observer, - observed=observed, - ) - - # Verify all dreams are pending - assert len(dream_scheduler.pending_dreams) == 3 - - # Cancel all dreams where observed=bob - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, observed - ) - - # All three should be cancelled - assert len(cancelled) == 3 - for key in work_unit_keys: - assert key in cancelled - assert len(dream_scheduler.pending_dreams) == 0 - - @pytest.mark.asyncio - async def test_does_not_cancel_dreams_for_different_observed( - self, dream_scheduler: DreamScheduler - ): - """Cancelling dreams for one observed peer should not affect others.""" - workspace_name = "test_workspace" - - # Dream for Alice observing Bob - key_alice_bob = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": "alice", - "observed": "bob", - "dream_type": "omni", - }, - ) - - # Dream for Alice observing Charlie (should NOT be cancelled) - key_alice_charlie = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": "alice", - "observed": "charlie", - "dream_type": "omni", - }, - ) - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - await dream_scheduler.schedule_dream( - key_alice_bob, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer="alice", - observed="bob", - ) - await dream_scheduler.schedule_dream( - key_alice_charlie, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer="alice", - observed="charlie", - ) - - assert len(dream_scheduler.pending_dreams) == 2 - - # Cancel only dreams where observed=bob - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, "bob" - ) - - # Only the bob dream should be cancelled - assert key_alice_bob in cancelled - assert key_alice_charlie not in cancelled - assert key_alice_charlie in dream_scheduler.pending_dreams - - @pytest.mark.asyncio - async def test_does_not_cancel_dreams_for_different_workspace( - self, dream_scheduler: DreamScheduler - ): - """Cancelling dreams should be scoped to the correct workspace.""" - observed = "bob" - - key_ws1 = construct_work_unit_key( - "workspace1", - { - "task_type": "dream", - "observer": "alice", - "observed": observed, - "dream_type": "omni", - }, - ) - key_ws2 = construct_work_unit_key( - "workspace2", - { - "task_type": "dream", - "observer": "alice", - "observed": observed, - "dream_type": "omni", - }, - ) - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - await dream_scheduler.schedule_dream( - key_ws1, - "workspace1", - delay_minutes=60, - dream_type=DreamType.OMNI, - observer="alice", - observed=observed, - ) - await dream_scheduler.schedule_dream( - key_ws2, - "workspace2", - delay_minutes=60, - dream_type=DreamType.OMNI, - observer="alice", - observed=observed, - ) - - # Cancel only in workspace1 - cancelled = await dream_scheduler.cancel_dreams_for_observed( - "workspace1", observed - ) - - assert key_ws1 in cancelled - assert key_ws2 not in cancelled - assert key_ws2 in dream_scheduler.pending_dreams +async def _make_collection( + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + internal_metadata: dict[str, object] | None = None, +) -> models.Collection: + workspace, peer = sample_data + collection = models.Collection( + observer=peer.name, + observed=peer.name, + workspace_name=workspace.name, + internal_metadata=internal_metadata or {}, + ) + db_session.add(collection) + await db_session.commit() + return collection -class TestThresholdFilter: - """Regression tests for Finding 2: threshold must count only explicit-level docs. - - 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.fixture(autouse=True) - def _pin_dream_config(self): - """Pin DOCUMENT_THRESHOLD=50 and ENABLED_TYPES=['omni'] for this class. - - 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. - """ - with ( - patch("src.dreamer.dream_scheduler.settings.DREAM.DOCUMENT_THRESHOLD", 50), - patch("src.dreamer.dream_scheduler.settings.DREAM.ENABLED_TYPES", ["omni"]), - ): - yield - - 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.""" +async def _insert_docs( + db_session: AsyncSession, + collection: models.Collection, + level: str, + count: int, + *, + age_minutes: int = 0, + session_name: str | None = None, +) -> None: + created_at = _now() - datetime.timedelta(minutes=age_minutes) + for _ in range(count): db_session.add( models.Document( content="test", @@ -340,133 +52,293 @@ class TestThresholdFilter: workspace_name=collection.workspace_name, observer=collection.observer, observed=collection.observed, + session_name=session_name, + created_at=created_at, ) ) + await db_session.commit() - @pytest.mark.asyncio - async def test_mixed_levels_below_explicit_threshold( + +async def _insert_dream_item( + db_session: AsyncSession, + collection: models.Collection, + *, + age_minutes: int, + processed: bool, + error: str | None = None, +) -> None: + work_unit_key = construct_work_unit_key( + collection.workspace_name, + { + "task_type": "dream", + "observer": collection.observer, + "observed": collection.observed, + "dream_type": DreamType.OMNI.value, + }, + ) + db_session.add( + models.QueueItem( + work_unit_key=work_unit_key, + payload={"task_type": "dream"}, + task_type="dream", + workspace_name=collection.workspace_name, + processed=processed, + error=error, + created_at=_now() - datetime.timedelta(minutes=age_minutes), + ) + ) + await db_session.commit() + + +@pytest.fixture(autouse=True) +def _pin_dream_config(): + with ( + patch("src.backlog.settings.DREAM.ENABLED", True), + patch("src.backlog.settings.DREAM.DOCUMENT_THRESHOLD", 50), + patch("src.backlog.settings.DREAM.ENABLED_TYPES", ["omni"]), + patch("src.backlog.settings.DREAM.IDLE_TIMEOUT_MINUTES", 60), + patch("src.backlog.settings.DREAM.MIN_HOURS_BETWEEN_DREAMS", 8), + ): + yield + + +@pytest.mark.asyncio +class TestFindDueDreams: + async def test_below_threshold_is_not_due( self, - dream_scheduler: DreamScheduler, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """30 explicit + 40 deductive + 10 inductive → should NOT trigger. + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 30, age_minutes=90) - 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() + assert await find_due_dreams(db_session) == [] - with patch.object(dream_scheduler, "schedule_dream", new_callable=AsyncMock): - scheduled = await check_and_schedule_dream(db_session, collection) + async def test_derived_levels_do_not_count( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 30, age_minutes=90) + await _insert_docs(db_session, collection, "deductive", 40, age_minutes=90) + await _insert_docs(db_session, collection, "contradiction", 40, age_minutes=90) - assert scheduled is False, ( - "Threshold should filter on explicit level only — dreamer output " - "(deductive/inductive) must not count toward the trigger." + assert await find_due_dreams(db_session) == [] + + async def test_threshold_met_but_not_idle_is_not_due( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=1) + + assert await find_due_dreams(db_session) == [] + + async def test_threshold_met_and_idle_is_due( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + + due = await find_due_dreams(db_session) + + assert len(due) == 1 + assert due[0].observer == collection.observer + assert due[0].observed == collection.observed + assert due[0].dream_type is DreamType.OMNI + assert due[0].documents_since_last_dream == 60 + + async def test_documents_since_last_dream_uses_stored_count( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection( + db_session, sample_data, {"dream": {"last_dream_document_count": 40}} + ) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + + assert await find_due_dreams(db_session) == [] + + async def test_min_hours_gate_blocks_a_recent_dream( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + last_dream_at = (_now() - datetime.timedelta(hours=2)).isoformat() + collection = await _make_collection( + db_session, sample_data, {"dream": {"last_dream_at": last_dream_at}} + ) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + + assert await find_due_dreams(db_session) == [] + + async def test_pending_dream_item_blocks( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + await _insert_dream_item( + db_session, collection, age_minutes=10, processed=False ) - @pytest.mark.asyncio - async def test_explicit_only_at_threshold( + assert await find_due_dreams(db_session) == [] + + async def test_failed_dream_waits_for_new_documents( 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() + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + await _insert_dream_item( + db_session, collection, age_minutes=80, processed=True, error="boom" + ) - with patch.object( - dream_scheduler, "schedule_dream", new_callable=AsyncMock - ) as mock_schedule: - scheduled = await check_and_schedule_dream(db_session, collection) + assert await find_due_dreams(db_session) == [] - assert scheduled is True - assert mock_schedule.called, "schedule_dream should fire when threshold met" - - @pytest.mark.asyncio - async def test_contradiction_excluded_from_count( + async def test_failed_dream_retries_after_new_documents( self, - dream_scheduler: DreamScheduler, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """Contradiction-level docs are dreamer output — must not count. + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + await _insert_dream_item( + db_session, collection, age_minutes=80, processed=True, error="boom" + ) + await _insert_docs(db_session, collection, "explicit", 1, age_minutes=70) - 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") + due = await find_due_dreams(db_session) + + assert len(due) == 1 + assert due[0].documents_since_last_dream == 61 + + async def test_dreams_disabled_returns_nothing( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + + with patch("src.backlog.settings.DREAM.ENABLED", False): + assert await find_due_dreams(db_session) == [] + + async def test_card_refresh_is_never_enqueued( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + collection = await _make_collection(db_session, sample_data) + await _insert_docs(db_session, collection, "explicit", 60, age_minutes=90) + + with patch("src.backlog.settings.DREAM.ENABLED_TYPES", ["card_refresh"]): + assert await find_due_dreams(db_session) == [] + + +@pytest.mark.asyncio +class TestExecuteDream: + async def test_enqueues_the_dream( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + workspace, peer = sample_data + session = models.Session( + name=f"dream-session-{peer.name}", workspace_name=workspace.name + ) + db_session.add(session) await db_session.commit() - with patch.object(dream_scheduler, "schedule_dream", new_callable=AsyncMock): - scheduled = await check_and_schedule_dream(db_session, collection) + collection = await _make_collection(db_session, sample_data) + await _insert_docs( + db_session, + collection, + "explicit", + 1, + age_minutes=90, + session_name=session.name, + ) - assert scheduled is False + with patch( + "src.deriver.enqueue.enqueue_dream", new_callable=AsyncMock + ) as mock_enqueue: + await execute_dream( + workspace.name, + DreamType.OMNI, + observer=peer.name, + observed=peer.name, + trigger_reason="document_threshold", + delay_reason="idle_timeout", + ) + assert mock_enqueue.called + assert mock_enqueue.call_args.kwargs["session_name"] == session.name -class TestEnqueueCancelsDreamsCorrectly: - """Integration test verifying the full flow of message enqueue cancelling dreams.""" - - @pytest.mark.asyncio - async def test_enqueue_cancels_peer_to_peer_dreams( - self, dream_scheduler: DreamScheduler + async def test_skips_when_no_documents( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], ): - """When a message is enqueued, it should cancel all dreams for that observed peer.""" + workspace, peer = sample_data + await _make_collection(db_session, sample_data) - workspace_name = "test_workspace" - observed = "bob" - - # Schedule dreams for multiple observers watching bob - keys: list[Any] = [] - for observer in ["alice", "charlie", "bob"]: - key = construct_work_unit_key( - workspace_name, - { - "task_type": "dream", - "observer": observer, - "observed": observed, - "dream_type": "omni", - }, - ) - keys.append(key) - - with patch.object(dream_scheduler, "execute_dream", new_callable=AsyncMock): - await dream_scheduler.schedule_dream( - key, - workspace_name, - delay_minutes=60, - dream_type=DreamType.OMNI, - observer=observer, - observed=observed, - ) - - assert len(dream_scheduler.pending_dreams) == 3 - - # Mock the database operations in enqueue - with patch("src.deriver.enqueue.tracked_db"): - # The enqueue function should cancel dreams via cancel_dreams_for_observed - # We just test that the scheduler method was called correctly - cancelled = await dream_scheduler.cancel_dreams_for_observed( - workspace_name, observed + with patch( + "src.deriver.enqueue.enqueue_dream", new_callable=AsyncMock + ) as mock_enqueue: + await execute_dream( + workspace.name, + DreamType.OMNI, + observer=peer.name, + observed=peer.name, ) - # All three dreams should be cancelled - assert len(cancelled) == 3 - assert len(dream_scheduler.pending_dreams) == 0 + assert not mock_enqueue.called + + +@pytest.mark.asyncio +class TestPollerEnqueueFailures: + async def test_one_failing_dream_does_not_block_the_others( + self, + db_session: AsyncSession, # pyright: ignore[reportUnusedParameter] + ): + from src.backlog import BacklogMetricsPoller, DueDream + + due = [ + DueDream( + workspace_name="ws", + observer=name, + observed=name, + dream_type=DreamType.OMNI, + documents_since_last_dream=60, + ) + for name in ("first", "second", "third") + ] + attempted: list[str] = [] + + async def flaky( + _workspace_name: str, + _dream_type: DreamType, + *, + observer: str, + observed: str, # pyright: ignore[reportUnusedParameter] + **_kwargs: object, + ) -> None: + attempted.append(observer) + if observer == "first": + raise ValueError("boom") + + with ( + patch("src.backlog.find_due_dreams", new=AsyncMock(return_value=due)), + patch("src.backlog.execute_dream", new=AsyncMock(side_effect=flaky)), + ): + await BacklogMetricsPoller()._refresh() # pyright: ignore[reportPrivateUsage] + + assert attempted == ["first", "second", "third"] diff --git a/tests/dreamer/test_dreamer_integration.py b/tests/dreamer/test_dreamer_integration.py index fb9b1894..5b6dc037 100644 --- a/tests/dreamer/test_dreamer_integration.py +++ b/tests/dreamer/test_dreamer_integration.py @@ -7,7 +7,7 @@ collection's internal_metadata on successful dreams — and critically, does NOT land on failures or exceptions. """ -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import AsyncMock, patch @@ -17,12 +17,9 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from src import models +from src.backlog import find_due_dreams from src.deriver.enqueue import enqueue_dream -from src.dreamer.dream_scheduler import ( - DreamScheduler, - check_and_schedule_dream, - set_dream_scheduler, -) +from src.dreamer.dream_scheduler import execute_dream from src.dreamer.orchestrator import DreamResult, process_dream from src.schemas import ( DreamType, @@ -286,7 +283,7 @@ 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 + the due-check'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. """ @@ -305,8 +302,7 @@ 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 in - check_and_schedule_dream. + returned — matching the explicit-only count query in the due-check. """ workspace, peer = sample_data @@ -380,14 +376,7 @@ class TestExecuteDreamSessionFilter: } ) - # 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), + with ( patch( "src.deriver.enqueue.enqueue_dream", side_effect=capture_enqueue_dream, @@ -406,14 +395,12 @@ class TestExecuteDreamSessionFilter: ), ), ): - await scheduler.execute_dream( + await 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, " @@ -442,30 +429,27 @@ class TestGuardPairCoherence: """ @pytest_asyncio.fixture - async def _scheduler(self): - DreamScheduler.reset_singleton() - scheduler = DreamScheduler() - set_dream_scheduler(scheduler) + async def _dream_config(self): 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"]), + patch("src.backlog.settings.DREAM.ENABLED", True), + patch("src.backlog.settings.DREAM.DOCUMENT_THRESHOLD", 50), + patch("src.backlog.settings.DREAM.ENABLED_TYPES", ["omni"]), + patch("src.backlog.settings.DREAM.IDLE_TIMEOUT_MINUTES", 60), ): - yield scheduler - DreamScheduler.reset_singleton() + yield @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, + _dream_config: None, ): """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. + docs arrive past the threshold again, but the due-check sees the + pending queue row and reports nothing due — no second QueueItem. """ workspace, peer = sample_data collection = models.Collection( @@ -475,6 +459,7 @@ class TestGuardPairCoherence: internal_metadata={}, ) db_session.add(collection) + idle_created_at = datetime.now(timezone.utc) - timedelta(minutes=90) for i in range(50): db_session.add( models.Document( @@ -483,6 +468,7 @@ class TestGuardPairCoherence: workspace_name=workspace.name, observer=peer.name, observed=peer.name, + created_at=idle_created_at, ) ) await db_session.commit() @@ -515,17 +501,18 @@ class TestGuardPairCoherence: workspace_name=workspace.name, observer=peer.name, observed=peer.name, + created_at=idle_created_at, ) ) await db_session.commit() await db_session.refresh(collection) - scheduled = await check_and_schedule_dream(db_session, collection) + due = await find_due_dreams(db_session) - 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." + assert due == [], ( + "The due-check must report nothing 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, ( @@ -538,11 +525,11 @@ class TestGuardPairCoherence: self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], - _scheduler: DreamScheduler, + _dream_config: None, ): """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. + untouched, so the due-check re-schedules on the same corpus instead + of silently consuming the baseline. """ workspace, peer = sample_data collection = models.Collection( @@ -552,6 +539,7 @@ class TestGuardPairCoherence: internal_metadata={}, ) db_session.add(collection) + idle_created_at = datetime.now(timezone.utc) - timedelta(minutes=90) for i in range(50): db_session.add( models.Document( @@ -560,6 +548,7 @@ class TestGuardPairCoherence: workspace_name=workspace.name, observer=peer.name, observed=peer.name, + created_at=idle_created_at, ) ) await db_session.commit() @@ -586,14 +575,10 @@ class TestGuardPairCoherence: "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) + due = await find_due_dreams(db_session) - assert scheduled is True, ( + assert len(due) == 1, ( "After a silent failure both guards should still allow the " - "same-corpus retry — 50 explicit docs ≥ threshold, no prior " + "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_scopes.py b/tests/routes/test_scopes.py index a68a3725..c47ffd8f 100644 --- a/tests/routes/test_scopes.py +++ b/tests/routes/test_scopes.py @@ -738,7 +738,7 @@ async def test_crud_get_peer_resolves_scope_and_dotted_names( ): """crud.get_peer must accept names outside RESOURCE_NAME_PATTERN. - This is the Dreamer's preflight path: DreamScheduler passes + This is the Dreamer's preflight path: the dream check passes ``collection.observer`` straight through, and scope peers have ``observe_others=true``, so ``(scope.x, peer)`` collections exist and get dreamt. While get_peer took a PeerCreate, every such dream died at preflight diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py index 8b60154a..ec883494 100644 --- a/tests/test_db_resilience.py +++ b/tests/test_db_resilience.py @@ -152,9 +152,6 @@ async def test_polling_loop_idle_sleeps_once_per_cycle( sleeps: list[float] = [] polls = {"n": 0} - async def fake_cleanup() -> None: - return None - async def fake_claim() -> dict[str, str]: polls["n"] += 1 if polls["n"] >= 5: @@ -164,7 +161,6 @@ async def test_polling_loop_idle_sleeps_once_per_cycle( async def fake_sleep(seconds: float) -> None: sleeps.append(seconds) - monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) # queue_manager calls asyncio.sleep on the stdlib module; patch it there. monkeypatch.setattr(asyncio, "sleep", fake_sleep) @@ -212,9 +208,6 @@ async def test_reconciler_work_does_not_reset_backoff( sleeps: list[float] = [] polls = {"n": 0} - async def fake_cleanup() -> None: - return None - async def fake_claim() -> dict[str, str]: polls["n"] += 1 if polls["n"] >= 5: @@ -228,7 +221,6 @@ async def test_reconciler_work_does_not_reset_backoff( async def fake_sleep(seconds: float) -> None: sleeps.append(seconds) - monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) monkeypatch.setattr(qm, "process_work_unit", fake_process) monkeypatch.setattr(asyncio, "sleep", fake_sleep) @@ -259,9 +251,6 @@ async def test_tenant_work_still_resets_backoff( sleeps: list[float] = [] polls = {"n": 0} - async def fake_cleanup() -> None: - return None - async def fake_claim() -> dict[str, str]: polls["n"] += 1 if polls["n"] >= 5: @@ -274,7 +263,6 @@ async def test_tenant_work_still_resets_backoff( async def fake_sleep(seconds: float) -> None: sleeps.append(seconds) - monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) monkeypatch.setattr(qm, "process_work_unit", fake_process) monkeypatch.setattr(asyncio, "sleep", fake_sleep) @@ -318,104 +306,3 @@ def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None: assert value() == start -@pytest.mark.asyncio -async def test_stale_cleanup_time_gate(monkeypatch: pytest.MonkeyPatch) -> None: - """cleanup_stale_work_units runs at most once per gate interval per - instance (staleness is a minutes-timescale condition; per-poll cleanup - multiplies into needless fleet-wide write transactions). First poll always - runs it so a crashed predecessor's stale rows are recovered immediately.""" - monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) - monkeypatch.setattr( - settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 - ) - - from src.deriver import queue_manager as qm_mod - - qm = qm_mod.QueueManager() - runs = {"n": 0} - - async def fake_cleanup() -> None: - runs["n"] += 1 - - monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) - - clock = {"now": 1_000.0} - monkeypatch.setattr( - "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] - ) - - # First call runs (no prior attempt recorded). - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert runs["n"] == 1 - - # Inside the gate window: skipped. - clock["now"] += 10.0 - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert runs["n"] == 1 - - # Past the gate window: runs again. - clock["now"] += 60.0 - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert runs["n"] == 2 - - -@pytest.mark.asyncio -async def test_stale_cleanup_gate_failed_attempt_waits_full_interval( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """The gate records the ATTEMPT before running, so a failing cleanup is not - retried on every poll against a DB that is already struggling.""" - monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) - monkeypatch.setattr( - settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0 - ) - - from src.deriver import queue_manager as qm_mod - - qm = qm_mod.QueueManager() - attempts = {"n": 0} - - async def failing_cleanup() -> None: - attempts["n"] += 1 - raise RuntimeError("db unavailable") - - monkeypatch.setattr(qm, "cleanup_stale_work_units", failing_cleanup) - - clock = {"now": 1_000.0} - monkeypatch.setattr( - "src.deriver.queue_manager.time.monotonic", lambda: clock["now"] - ) - - with pytest.raises(RuntimeError): - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert attempts["n"] == 1 - - # Immediately after the failure: still gated, no hammering. - clock["now"] += 1.0 - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert attempts["n"] == 1 - - -@pytest.mark.asyncio -async def test_stale_cleanup_gate_zero_interval_runs_every_poll( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Interval 0.0 preserves legacy run-on-every-poll behavior.""" - monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) - monkeypatch.setattr( - settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 0.0 - ) - - from src.deriver import queue_manager as qm_mod - - qm = qm_mod.QueueManager() - runs = {"n": 0} - - async def fake_cleanup() -> None: - runs["n"] += 1 - - monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) - - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage] - assert runs["n"] == 2