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

feat: refactor to use ReconcilerScheduler

This commit is contained in:
Rajat Ahuja 2026-01-13 15:15:03 -05:00
parent 9ced00e2d5
commit cae29b7100
15 changed files with 465 additions and 128 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.env.template
View File

@ -253,3 +253,4 @@ VECTOR_STORE_MIGRATED=false
# LanceDB-specific settings (local embedded mode)
# VECTOR_STORE_LANCEDB_PATH=./lancedb_data
#VECTOR_STORE_RECONCILIATION_INTERVAL_SECONDS

1
config.toml.example
View File

@ -184,3 +184,4 @@ DIMENSIONS = 1536
# TURBOPUFFER_API_KEY = "your-turbopuffer-api-key"
# TURBOPUFFER_REGION = "us-east-1"
LANCEDB_PATH = "./lancedb_data"
RECONCILIATION_INTERVAL_SECONDS=

22
migrations/versions/119a52b73c60_support_external_embeddings.py
View File

@ -7,6 +7,8 @@ This migration:
hybrid sync/soft delete pattern for vector store consistency.
3. Adds sync_state, last_sync_at, and sync_attempts columns to documents and
message_embeddings tables for tracking vector store synchronization status.
4. Adds partial unique index on queue table for reconciler task deduplication,
ensuring only one pending reconciler task exists per work_unit_key.
Revision ID: 119a52b73c60
Revises: 7c0d9a4e3b1f
@ -181,11 +183,31 @@ def upgrade() -> None:
schema=schema,
)
# Add partial unique index on queue table for reconciler task deduplication
# This ensures only one pending reconciler task exists per work_unit_key
if not index_exists("queue", "uq_queue_work_unit_key", inspector):
op.create_index(
"uq_queue_work_unit_key",
"queue",
["work_unit_key"],
unique=True,
schema=schema,
postgresql_where=sa.text("task_type = 'reconciler' AND processed = false"),
)
def downgrade() -> None:
"""Remove deleted_at columns and revert embedding columns."""
inspector = sa.inspect(op.get_bind())
# Drop reconciler queue index if it exists
if index_exists("queue", "uq_queue_work_unit_key", inspector):
op.drop_index(
"uq_queue_work_unit_key",
table_name="queue",
schema=schema,
)
# Drop message_embeddings indexes if they exist
if index_exists(
"message_embeddings", "ix_message_embeddings_sync_state_last_sync_at", inspector

4
src/config.py
View File

@ -563,6 +563,10 @@ class VectorStoreSettings(HonchoSettings):
# LanceDB-specific settings (local embedded mode)
LANCEDB_PATH: str = "./lancedb_data"
RECONCILIATION_INTERVAL_SECONDS: Annotated[int, Field(default=300, gt=0)] = (
300 # 5 minutes
)
@model_validator(mode="after")
def _require_api_key_for_turbopuffer(self) -> "VectorStoreSettings":
if self.TYPE == "turbopuffer" and not self.TURBOPUFFER_API_KEY:

58
src/deriver/consumer.py
View File

@ -10,12 +10,15 @@ from src.deriver.deriver import process_representation_tasks_batch
from src.dreamer.dreamer import process_dream
from src.exceptions import ResourceNotFoundException
from src.models import Message
from src.schemas import ResolvedConfiguration
from src.reconciler.queue_cleanup import cleanup_queue_items
from src.reconciler.sync_vectors import run_vector_reconciliation_cycle
from src.schemas import ReconcilerType, ResolvedConfiguration
from src.utils import summarizer
from src.utils.logging import log_performance_metrics
from src.utils.queue_payload import (
DeletionPayload,
DreamPayload,
ReconcilerPayload,
SummaryPayload,
WebhookPayload,
)
@ -119,6 +122,19 @@ async def process_item(queue_item: models.QueueItem) -> None:
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_deletion(validated, workspace_name)
elif task_type == "reconciler":
with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"):
try:
validated = ReconcilerPayload(**queue_payload)
except ValidationError as e:
logger.error(
"Invalid reconciler payload received: %s. Payload: %s",
str(e),
queue_payload,
)
raise ValueError(f"Invalid payload structure: {str(e)}") from e
await process_reconciler(validated)
else:
raise ValueError(f"Invalid task type: {task_type}")
@ -220,3 +236,43 @@ async def process_deletion(
else:
raise ValueError(f"Unsupported deletion type: {deletion_type}")
async def process_reconciler(payload: ReconcilerPayload) -> None:
"""
Process a reconciler task from the queue.
Currently supports:
- sync_vectors: Syncs pending documents/message embeddings to vector store
and cleans up soft-deleted documents.
- cleanup_queue: Removes old processed queue items.
Args:
payload: The reconciler payload containing the reconciler type
"""
reconciler_type = payload.reconciler_type
if reconciler_type == ReconcilerType.SYNC_VECTORS:
logger.debug("Processing sync_vectors task")
metrics = await run_vector_reconciliation_cycle()
if (
metrics.total_synced > 0
or metrics.total_failed > 0
or metrics.total_cleaned > 0
):
logger.info(
"Reconciliation complete: synced %s docs, %s message embeddings; failed %s docs, %s message embeddings; cleaned %s docs",
metrics.documents_synced,
metrics.message_embeddings_synced,
metrics.documents_failed,
metrics.message_embeddings_failed,
metrics.documents_cleaned,
)
elif reconciler_type == ReconcilerType.CLEANUP_QUEUE:
logger.debug("Processing cleanup_queue task")
await cleanup_queue_items()
else:
raise ValueError(f"Unsupported reconciler type: {reconciler_type}")

143
src/deriver/queue_manager.py
View File

@ -23,13 +23,17 @@ from src.deriver.consumer import (
process_item,
process_representation_batch,
)
from src.deriver.vector_reconciliation import run_vector_reconciliation_cycle
from src.dreamer.dream_scheduler import (
DreamScheduler,
get_dream_scheduler,
set_dream_scheduler,
)
from src.models import QueueItem
from src.reconciler import (
ReconcilerScheduler,
get_reconciler_scheduler,
set_reconciler_scheduler,
)
from src.schemas import ResolvedConfiguration
from src.sentry import initialize_sentry
from src.utils.work_unit import parse_work_unit_key
@ -50,17 +54,12 @@ class WorkerOwnership(NamedTuple):
aqs_id: str # The ID of the ActiveQueueSession that the worker is processing
QUEUE_CLEANUP_INTERVAL_SECONDS = 43200 # 12 hours
RECONCILIATION_INTERVAL_SECONDS = 300 # 5 minutes
class QueueManager:
def __init__(self):
self.shutdown_event: asyncio.Event = asyncio.Event()
self.active_tasks: set[asyncio.Task[None]] = set()
self.worker_ownership: dict[str, WorkerOwnership] = {}
self.queue_empty_flag: asyncio.Event = asyncio.Event()
self._maintenance_task: asyncio.Task[None] | None = None
# Initialize from settings
self.workers: int = settings.DERIVER.WORKERS
@ -74,6 +73,14 @@ class QueueManager:
else:
self.dream_scheduler = existing_scheduler
# Get or create the singleton reconciler scheduler
existing_reconciler = get_reconciler_scheduler()
if existing_reconciler is None:
self.reconciler_scheduler: ReconcilerScheduler = ReconcilerScheduler()
set_reconciler_scheduler(self.reconciler_scheduler)
else:
self.reconciler_scheduler = existing_reconciler
# Initialize Sentry if enabled, using settings
if settings.SENTRY.ENABLED:
initialize_sentry(integrations=[AsyncioIntegration()])
@ -116,11 +123,11 @@ class QueueManager:
)
logger.debug("Signal handlers registered")
# Start background maintenance loop (handles both queue cleanup and vector cleanup)
# Start the reconciler scheduler
try:
self._maintenance_task = asyncio.create_task(self._maintenance_loop())
await self.reconciler_scheduler.start()
except Exception:
logger.exception("Failed to start maintenance loop")
logger.exception("Failed to start reconciler scheduler")
# Run the polling loop directly in this task
logger.debug("Starting polling loop directly")
@ -137,6 +144,9 @@ class QueueManager:
# Cancel all pending dreams
await self.dream_scheduler.shutdown()
# Stop the reconciler scheduler
await self.reconciler_scheduler.shutdown()
if self.active_tasks:
logger.info(
f"Waiting for {len(self.active_tasks)} active tasks to complete..."
@ -168,14 +178,6 @@ class QueueManager:
finally:
self.worker_ownership.clear()
# Cancel maintenance loop if running
if self._maintenance_task is not None:
from contextlib import suppress
self._maintenance_task.cancel()
with suppress(asyncio.CancelledError):
await self._maintenance_task
##########################
# Polling and Scheduling #
##########################
@ -323,113 +325,6 @@ class QueueManager:
# Queue Worker Logic #
######################
async def cleanup_queue_items(self) -> None:
"""Delete processed queue items.
Successfully processed queue items are deleted immediately,
while errored queue items are deleted after retention window."""
async with tracked_db("cleanup_queue_items") as db:
now = datetime.now(timezone.utc)
error_cutoff = now - timedelta(
seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS
)
await db.execute(
delete(models.QueueItem).where(
models.QueueItem.processed
& (
models.QueueItem.error.is_(None)
| (
models.QueueItem.error.is_not(None)
& (models.QueueItem.created_at < error_cutoff)
)
)
)
)
await db.commit()
async def _maintenance_loop(self) -> None:
"""
Run periodic maintenance tasks.
- Queue cleanup: every 12 hours (remove old processed/errored queue items)
- Reconciliation: every 5 minutes (sync vectors + clean up soft deletes)
Always runs to keep vector stores consistent with the DB
"""
# Track when each task should next run
next_queue_cleanup = datetime.now(timezone.utc)
next_vector_reconciliation = datetime.now(timezone.utc)
try:
while not self.shutdown_event.is_set():
now = datetime.now(timezone.utc)
# Run queue cleanup if due
if now >= next_queue_cleanup:
try:
await self.cleanup_queue_items()
except Exception:
logger.exception("Error during queue cleanup")
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception()
next_queue_cleanup = now + timedelta(
seconds=QUEUE_CLEANUP_INTERVAL_SECONDS
)
# Run vector store reconciliation if due
if now >= next_vector_reconciliation:
try:
logger.info("Running vector reconciliation cycle")
await self._run_reconciliation()
except Exception:
logger.exception("Error during vector reconciliation")
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception()
next_vector_reconciliation = now + timedelta(
seconds=RECONCILIATION_INTERVAL_SECONDS
)
# Sleep until next task is due or shutdown
# Filter out None values when computing next task time
task_times = [next_queue_cleanup]
task_times.append(next_vector_reconciliation)
next_task_time = min(task_times)
sleep_seconds = max(
0, (next_task_time - datetime.now(timezone.utc)).total_seconds()
)
try:
await asyncio.wait_for(
self.shutdown_event.wait(),
timeout=sleep_seconds
or 1, # At least 1 second to avoid busy loop
)
break # Shutdown event set
except asyncio.TimeoutError:
# Timeout means it's time for next task
pass
except asyncio.CancelledError:
logger.debug("Maintenance loop cancelled")
raise
async def _run_reconciliation(self) -> None:
"""Run vector store reconciliation for sync + cleanup."""
metrics = await run_vector_reconciliation_cycle()
if (
metrics.total_synced > 0
or metrics.total_failed > 0
or metrics.total_cleaned > 0
):
logger.info(
"Reconciliation: synced %s docs, %s message embeddings; failed %s docs, %s message embeddings; cleaned %s docs",
metrics.documents_synced,
metrics.message_embeddings_synced,
metrics.documents_failed,
metrics.message_embeddings_failed,
metrics.documents_cleaned,
)
async def _handle_processing_error(
self,
error: Exception,

13
src/reconciler/__init__.py Normal file
View File

@ -0,0 +1,13 @@
"""Reconciler for self-healing background tasks like vector sync and cleanup."""
from .scheduler import (
ReconcilerScheduler,
get_reconciler_scheduler,
set_reconciler_scheduler,
)
__all__ = [
"ReconcilerScheduler",
"get_reconciler_scheduler",
"set_reconciler_scheduler",
]

45
src/reconciler/queue_cleanup.py Normal file
View File

@ -0,0 +1,45 @@
"""
Queue cleanup job.
This module provides a periodic cleanup job that removes old processed queue items.
"""
import logging
from datetime import datetime, timedelta, timezone
from sqlalchemy import delete
from src import models
from src.config import settings
from src.dependencies import tracked_db
logger = logging.getLogger(__name__)
async def cleanup_queue_items() -> None:
"""
Delete processed queue items.
Successfully processed queue items are deleted immediately,
while errored queue items are deleted after retention window.
"""
async with tracked_db("cleanup_queue_items") as db:
now = datetime.now(timezone.utc)
error_cutoff = now - timedelta(
seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS
)
await db.execute(
delete(models.QueueItem).where(
models.QueueItem.processed
& (
models.QueueItem.error.is_(None)
| (
models.QueueItem.error.is_not(None)
& (models.QueueItem.created_at < error_cutoff)
)
)
)
)
await db.commit()
logger.info("Queue cleanup completed")

283
src/reconciler/scheduler.py Normal file
View File

@ -0,0 +1,283 @@
"""
Reconciler scheduler for self-healing background tasks.
This module provides a scheduler for running reconciliation and cleanup tasks
like vector store sync and soft-delete cleanup. It ensures only one task of each
type runs at a time across multiple deriver instances by using the queue table
for coordination.
"""
import asyncio
import contextlib
import logging
from datetime import datetime, timedelta, timezone
import sentry_sdk
from pydantic import BaseModel
from sqlalchemy import exists, select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
from src.dependencies import tracked_db
from src.models import QueueItem
logger = logging.getLogger(__name__)
# System workspace used for global reconciler tasks
SYSTEM_WORKSPACE_NAME = "__system__"
class ReconcilerTask(BaseModel):
"""Definition of a reconciler task."""
name: str
work_unit_key: str
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(
name="sync_vectors",
work_unit_key="global:sync_vectors",
interval_seconds=settings.VECTOR_STORE.RECONCILIATION_INTERVAL_SECONDS,
),
"cleanup_queue": ReconcilerTask(
name="cleanup_queue",
work_unit_key="global:cleanup_queue",
interval_seconds=QUEUE_CLEANUP_INTERVAL_SECONDS,
),
}
_reconciler_scheduler: "ReconcilerScheduler | None" = None
def set_reconciler_scheduler(scheduler: "ReconcilerScheduler") -> None:
"""Set the global reconciler scheduler reference."""
global _reconciler_scheduler
_reconciler_scheduler = scheduler
def get_reconciler_scheduler() -> "ReconcilerScheduler | None":
"""Get the global reconciler scheduler reference."""
return _reconciler_scheduler
class ReconcilerScheduler:
"""
Scheduler for self-healing reconciliation and cleanup tasks.
Ensures only one task of each type runs at a time across multiple deriver
instances by using the queue table for coordination. This provides:
- Vector store synchronization (syncing pending documents/embeddings)
- Soft-delete cleanup (removing soft-deleted documents from vector stores)
- Self-healing behavior (retrying failed syncs)
- Extensible task registry for adding new maintenance tasks
Each task type has its own interval and work_unit_key, allowing multiple
different tasks to be queued simultaneously while preventing duplicates
of the same task type.
"""
_instance: "ReconcilerScheduler | 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 ReconcilerScheduler._initialized:
self._scheduler_task: asyncio.Task[None] | None = None
self._shutdown_event: asyncio.Event = asyncio.Event()
# Track next run time for each task
self._next_run: dict[str, datetime] = {}
ReconcilerScheduler._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 start(self) -> None:
"""Start the reconciler scheduler loop."""
if self._scheduler_task is not None:
logger.warning("ReconcilerScheduler already running")
return
self._shutdown_event.clear()
# Initialize next run times to first interval
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._scheduler_task = asyncio.create_task(self._scheduler_loop())
logger.info(
"ReconcilerScheduler started with %d tasks: %s",
len(RECONCILER_TASKS),
list(RECONCILER_TASKS.keys()),
)
async def shutdown(self) -> None:
"""Stop the reconciler scheduler."""
if self._scheduler_task is None:
return
logger.info("Shutting down ReconcilerScheduler...")
self._shutdown_event.set()
try:
await asyncio.wait_for(self._scheduler_task, timeout=5.0)
except asyncio.TimeoutError:
logger.warning("ReconcilerScheduler shutdown timed out, cancelling task")
self._scheduler_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._scheduler_task
self._scheduler_task = None
self._next_run.clear()
logger.info("ReconcilerScheduler stopped")
async def _scheduler_loop(self) -> None:
"""
Main scheduler loop that enqueues tasks based on their intervals.
Each task has its own interval and the loop checks all tasks on each
iteration, enqueueing any that are due.
"""
try:
while not self._shutdown_event.is_set():
now = datetime.now(timezone.utc)
# Check each task and enqueue if due
for task_name, task in RECONCILER_TASKS.items():
next_run = self._next_run.get(task_name, now)
if now >= next_run:
try:
enqueued = await self._try_enqueue_task(task)
if enqueued:
logger.debug("Enqueued task: %s", task_name)
except Exception as e:
logger.exception("Error enqueueing task %s", task_name)
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
# Schedule next run regardless of whether we enqueued
# (if already pending/in-progress, we'll skip next time too)
self._next_run[task_name] = now + timedelta(
seconds=task.interval_seconds
)
# Calculate sleep time until next task is due
if self._next_run:
next_task_time = min(self._next_run.values())
sleep_seconds = max(
1.0, # At least 1 second to avoid busy loop
(next_task_time - datetime.now(timezone.utc)).total_seconds(),
)
else:
sleep_seconds = 60.0 # Default if no tasks
try:
await asyncio.wait_for(
self._shutdown_event.wait(),
timeout=sleep_seconds,
)
break # Shutdown event was set
except asyncio.TimeoutError:
# Timeout means interval elapsed, continue loop
pass
except asyncio.CancelledError:
logger.debug("ReconcilerScheduler loop cancelled")
raise
async def _try_enqueue_task(self, task: ReconcilerTask) -> bool:
"""
Attempt to enqueue a reconciler task.
This is idempotent - if the task is already in-progress (has an
ActiveQueueSession) or pending in the queue, the enqueue is skipped.
Args:
task: The task definition to enqueue
Returns:
True if a task was enqueued, False if skipped
"""
async with tracked_db("reconciler_enqueue") as db:
# Ensure the system workspace exists (for FK constraint)
await self._ensure_system_workspace(db)
# Check if task is already in progress
in_progress_check = select(
exists(
select(models.ActiveQueueSession.id).where(
models.ActiveQueueSession.work_unit_key == task.work_unit_key
)
)
)
is_in_progress = await db.scalar(in_progress_check)
if is_in_progress:
logger.debug("Task %s already in progress, skipping enqueue", task.name)
return False
# Check if there's already a pending task
pending_check = select(
exists(
select(QueueItem.id).where(
QueueItem.work_unit_key == task.work_unit_key,
QueueItem.processed == False, # noqa: E712
)
)
)
is_pending = await db.scalar(pending_check)
if is_pending:
logger.debug(
"Task %s already pending in queue, skipping enqueue", task.name
)
return False
# Enqueue the task using ORM
queue_item = QueueItem(
work_unit_key=task.work_unit_key,
payload={
"reconciler_type": task.name,
},
session_id=None,
task_type="reconciler",
workspace_name=SYSTEM_WORKSPACE_NAME,
message_id=None,
)
db.add(queue_item)
try:
await db.commit()
except IntegrityError:
# Another instance already enqueued this task (unique constraint)
await db.rollback()
logger.debug(
"Task %s already enqueued by another instance, skipping", task.name
)
return False
logger.info("Enqueued reconciler task: %s", task.name)
return True
async def _ensure_system_workspace(self, db: AsyncSession) -> None:
"""Ensure the system workspace exists for reconciler tasks."""
# Use upsert to create workspace if it doesn't exist
stmt = pg_insert(models.Workspace).values(name=SYSTEM_WORKSPACE_NAME)
stmt = stmt.on_conflict_do_nothing(index_elements=["name"])
await db.execute(stmt)

0
src/deriver/vector_reconciliation.py → src/reconciler/sync_vectors.py
View File

7
src/schemas.py
View File

@ -27,6 +27,13 @@ class DreamType(str, Enum):
OMNI = "omni"
class ReconcilerType(str, Enum):
"""Types of reconciler tasks that can be performed."""
SYNC_VECTORS = "sync_vectors"
CLEANUP_QUEUE = "cleanup_queue"
class ReasoningConfiguration(BaseModel):
enabled: bool | None = Field(
default=None,

9
src/utils/queue_payload.py
View File

@ -3,7 +3,7 @@ from typing import Any, Literal
from pydantic import BaseModel, ConfigDict
from src.schemas import DreamType, ResolvedConfiguration
from src.schemas import DreamType, ReconcilerType, ResolvedConfiguration
class BasePayload(BaseModel):
@ -67,6 +67,13 @@ class DeletionPayload(BasePayload):
resource_id: str
class ReconcilerPayload(BasePayload):
"""Payload for reconciler tasks (vector sync, queue cleanup, self-healing)."""
task_type: Literal["reconciler"] = "reconciler"
reconciler_type: ReconcilerType
def create_webhook_payload(
event_type: str,
data: dict[str, Any],

4
src/utils/types.py
View File

@ -11,6 +11,8 @@ class GetOrCreateResult(NamedTuple, Generic[T]):
SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"]
TaskType = Literal["webhook", "summary", "representation", "dream", "deletion"]
TaskType = Literal[
"webhook", "summary", "representation", "dream", "deletion", "reconciler"
]
VectorSyncState = Literal["synced", "pending", "failed"]
DocumentLevel = Literal["explicit", "deductive", "inductive", "contradiction"]

1
tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py
View File

@ -12,6 +12,7 @@ INDEXES = (
("documents", "ix_documents_sync_state_last_sync_at"),
("message_embeddings", "ix_message_embeddings_sync_state"),
("message_embeddings", "ix_message_embeddings_sync_state_last_sync_at"),
("queue", "uq_queue_work_unit_key"),
)

2
tests/deriver/test_vector_reconciliation.py
View File

@ -15,7 +15,7 @@ from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.deriver.vector_reconciliation import (
from src.reconciler.sync_vectors import (
MAX_SYNC_ATTEMPTS,
ReconciliationMetrics,
_get_documents_needing_sync, # pyright: ignore[reportPrivateUsage]