RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Move dream scheduling and stale-claim cleanup into the API

This commit is contained in:
Ulysse Pence 2026-08-27 17:02:44 -02:00
parent 79a3b9d416
commit af46190eff
18 changed files with 561 additions and 1181 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

141
src/backlog.py
View File

@ -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:

6
src/config.py
View File

@ -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"]

2
src/crud/__init__.py
View File

@ -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",

58
src/crud/deriver.py
View File

@ -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(

13
src/crud/representation.py
View File

@ -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(

1
src/deriver/deriver.py
View File

@ -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 += (

21
src/deriver/enqueue.py
View File

@ -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

107
src/deriver/queue_manager.py
View File

@ -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):

8
src/dreamer/__init__.py
View File

@ -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",
]

437
src/dreamer/dream_scheduler.py
View File

@ -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,
)

31
src/reconciler/scheduler.py
View File

@ -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(),

2
tests/conftest.py
View File

@ -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",

9
tests/crud/test_deriver_backlog.py
View File

@ -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,

4
tests/crud/test_representation_manager.py
View File

@ -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(

708
tests/dreamer/test_dream_scheduler.py
View File

@ -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"]

79
tests/dreamer/test_dreamer_integration.py
View File

@ -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."

2
tests/routes/test_scopes.py
View File

@ -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

113
tests/test_db_resilience.py
View File

@ -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