Agent Zero
727d86614e
P0 (7): Auth-bypass removed, migrations fixed, plugin-upload disabled, RLS FORCE+WITH CHECK, plugin double-registration fixed, persistent volume, domain removed P1 (11): User/tenant model, Redis centralized, worker separated, transactional outbox, XSS fixed, DMS chunked streaming, permissions unified, password reset, metrics secured, config/docs fixed, cross-tenant FK P2 (4): Contact model normalized, cross-imports reduced 94%, commands+state machines for contacts/dms/mail/calendar, SPA path-traversal 8 new migrations, 99 unit tests, 13 commands, 8 contracts, 72 files changed
190 lines
6.9 KiB
Python
190 lines
6.9 KiB
Python
"""ARQ worker configuration and entrypoint."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from arq.connections import RedisSettings
|
|
from arq import cron
|
|
|
|
from app.config import get_settings
|
|
from app.core.job_registry import get_all_jobs, get_job, register_job
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ── Distributed lock helpers ─────────────────────────────────────────────────
|
|
# When multiple worker replicas run concurrently, cron jobs must not fire
|
|
# on every replica. We use a short-lived Redis SET NX lock per cron call
|
|
# so only one replica actually executes the job.
|
|
|
|
import redis.asyncio as aioredis # noqa: E402
|
|
import uuid # noqa: E402
|
|
|
|
|
|
async def _acquire_cron_lock(job_name: str, ttl_seconds: int = 120) -> str | None:
|
|
"""Try to acquire a distributed lock for a cron job.
|
|
|
|
Returns a lock token (random UUID) if acquired, or None if another
|
|
replica already holds the lock. The lock auto-expires after
|
|
*ttl_seconds* to avoid deadlocks if a worker crashes mid-job.
|
|
"""
|
|
settings = get_settings()
|
|
client = aioredis.from_url(settings.redis_url)
|
|
token = str(uuid.uuid4())
|
|
lock_key = f"leocrm:cron_lock:{job_name}"
|
|
try:
|
|
acquired = await client.set(lock_key, token, nx=True, ex=ttl_seconds)
|
|
return token if acquired else None
|
|
finally:
|
|
await client.aclose()
|
|
|
|
|
|
async def _release_cron_lock(job_name: str, token: str) -> None:
|
|
"""Release a previously acquired cron lock using a safe compare-and-delete."""
|
|
settings = get_settings()
|
|
client = aioredis.from_url(settings.redis_url)
|
|
lock_key = f"leocrm:cron_lock:{job_name}"
|
|
try:
|
|
# Lua script ensures we only delete if the token matches (avoid
|
|
# releasing a lock that was already expired and re-acquired).
|
|
script = (
|
|
b"if redis.call('get', KEYS[1]) == ARGV[1] "
|
|
b"then return redis.call('del', KEYS[1]) "
|
|
b"else return 0 end"
|
|
)
|
|
await client.eval(script, 1, lock_key, token.encode())
|
|
finally:
|
|
await client.aclose()
|
|
|
|
|
|
def _wrap_cron_with_lock(job_name: str, func: Any, ttl_seconds: int = 120) -> Any:
|
|
"""Wrap a cron callable so it acquires a distributed lock first.
|
|
|
|
If the lock cannot be acquired (another replica is handling it), the
|
|
wrapped function is silently skipped.
|
|
"""
|
|
import functools
|
|
|
|
@functools.wraps(func)
|
|
async def _locked_wrapper(ctx: dict[str, Any], *args: Any, **kwargs: Any) -> Any:
|
|
token = await _acquire_cron_lock(job_name, ttl_seconds=ttl_seconds)
|
|
if token is None:
|
|
logger.debug("Cron job '%s' skipped — lock held by another replica", job_name)
|
|
return None
|
|
try:
|
|
return await func(ctx, *args, **kwargs)
|
|
finally:
|
|
await _release_cron_lock(job_name, token)
|
|
|
|
return _locked_wrapper
|
|
|
|
|
|
def _get_redis_settings() -> RedisSettings:
|
|
"""Get Redis settings from app config."""
|
|
settings = get_settings()
|
|
return RedisSettings.from_dsn(settings.redis_url)
|
|
|
|
|
|
async def on_startup(ctx: dict[str, Any]) -> None:
|
|
"""Called when worker starts."""
|
|
logger.info("ARQ worker starting...")
|
|
# Register search providers (normally done by app startup)
|
|
try:
|
|
from app.core.db import get_session_factory
|
|
from app.plugins.builtins.unified_search.provider_registry import auto_register_providers
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
await auto_register_providers(db)
|
|
logger.info("Search providers registered for worker")
|
|
except Exception:
|
|
logger.warning("Failed to register search providers in worker", exc_info=True)
|
|
|
|
|
|
async def on_shutdown(ctx: dict[str, Any]) -> None:
|
|
"""Called when worker shuts down."""
|
|
logger.info("ARQ worker shutting down...")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Lazy-load plugin jobs via importlib so they register themselves with the
|
|
# job_registry. This avoids circular imports and keeps the worker decoupled
|
|
# from plugin internals.
|
|
# ---------------------------------------------------------------------------
|
|
def _lazy_register_plugin_jobs() -> None:
|
|
"""Import each plugin job module so its register_job() call fires."""
|
|
plugin_job_modules = [
|
|
"app.plugins.builtins.unified_search.jobs",
|
|
"app.plugins.builtins.ai_proactive.jobs",
|
|
"app.plugins.builtins.automation.scheduler",
|
|
"app.plugins.builtins.automation.workflow_timeout",
|
|
"app.plugins.builtins.automation.agent_runner",
|
|
"app.plugins.builtins.automation.execution_engine",
|
|
"app.plugins.builtins.tasks.jobs",
|
|
]
|
|
for mod_name in plugin_job_modules:
|
|
try:
|
|
import importlib
|
|
importlib.import_module(mod_name)
|
|
logger.debug("Lazy-loaded plugin jobs from %s", mod_name)
|
|
except Exception:
|
|
logger.warning("Failed to lazy-load plugin jobs from %s", mod_name, exc_info=True)
|
|
|
|
|
|
# Trigger lazy registration at module level so jobs are available when
|
|
# WorkerSettings.functions is evaluated.
|
|
_lazy_register_plugin_jobs()
|
|
|
|
|
|
# ── Outbox processor job ────────────────────────────────────────────────────
|
|
|
|
async def process_outbox_job(ctx: dict[str, Any]) -> None:
|
|
"""Poll the transactional outbox and publish pending events.
|
|
|
|
Uses a distributed Redis lock so only one worker replica processes the
|
|
outbox at a time. Runs every 5 seconds.
|
|
"""
|
|
from app.core.db import get_session_factory
|
|
from app.core.outbox import process_outbox_batch
|
|
|
|
factory = get_session_factory()
|
|
async with factory() as db:
|
|
try:
|
|
count = await process_outbox_batch(db, batch_size=50)
|
|
if count:
|
|
logger.info("Outbox: published %d events", count)
|
|
except Exception:
|
|
logger.error("Outbox processing failed", exc_info=True)
|
|
await db.rollback()
|
|
|
|
|
|
# Register the outbox job so it appears in get_all_jobs()
|
|
register_job("process_outbox", process_outbox_job)
|
|
|
|
|
|
class WorkerSettings:
|
|
"""ARQ worker settings."""
|
|
functions = get_all_jobs()
|
|
redis_settings = _get_redis_settings()
|
|
on_startup = on_startup
|
|
on_shutdown = on_shutdown
|
|
max_jobs = 10
|
|
job_timeout = 300
|
|
queue_name = "arq:queue"
|
|
cron_jobs = [
|
|
cron(
|
|
_wrap_cron_with_lock("scheduler_tick", get_job("scheduler_tick")),
|
|
minute={0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55},
|
|
),
|
|
cron(
|
|
_wrap_cron_with_lock("tasks_due_reminder", get_job("tasks_due_reminder")),
|
|
hour=8, minute=0,
|
|
),
|
|
# Outbox processor — every 5 seconds, guarded by distributed lock
|
|
cron(
|
|
_wrap_cron_with_lock("process_outbox", process_outbox_job, ttl_seconds=30),
|
|
second="*/5",
|
|
),
|
|
]
|