"""RabbitMQ consumer loop — main entry point for the asf-worker process. Run:: python -m app.worker """ from __future__ import annotations import asyncio import signal import sys from dataclasses import dataclass from typing import Awaitable, Callable import aio_pika import aio_pika.abc import structlog from motor.motor_asyncio import AsyncIOMotorClient from .config import WorkerSettings, get_settings from .handlers import ( WorkerDependencies, noop_handler, echo_handler, ai_handler, indexing_handler, notification_handler, mark_job_done, HandlerResult, ) from asf_events.base import JobRequestEnvelope from asf_domain import JobStatus, JobUpdate log = structlog.get_logger(__name__) JobHandlerFn = Callable[[JobRequestEnvelope, WorkerDependencies], Awaitable[HandlerResult]] """Type alias for a queue handler: (envelope, deps) → coroutine returning HandlerResult.""" # Queue name → handler function HANDLERS_BY_QUEUE: dict[str, Callable] = { "asf.ai.interactive": ai_handler, "asf.ai.background": ai_handler, "asf.indexing": indexing_handler, "asf.code_analysis": ai_handler, # use analyze_code via AI gateway "asf.testing": ai_handler, # use generate_tests via AI gateway "asf.devops": echo_handler, # stub "asf.notifications": notification_handler, } async def _claim_running(deps: WorkerDependencies, envelope: JobRequestEnvelope) -> None: """Atomically mark the job as running + claim worker ownership. Guards against duplicate-delivery races from RabbitMQ (at-least-once). If the job is already in a terminal state we bail out early. """ from datetime import datetime, timezone now = datetime.now(tz=timezone.utc) res = await deps.jobs.update_one( { "_id": envelope.job_id, "tenant_id": envelope.tenant_id, "status": {"$in": [JobStatus.QUEUED.value]}, }, { "$set": { "status": JobStatus.RUNNING.value, "started_at": now, "updated_at": now, "worker": deps.worker_id, }, "$inc": {"retries": 1}, }, ) if res.matched_count == 0: log.warn( "job_already_claimed_or_terminal", job_id=str(envelope.job_id), ) @dataclass class Consumer: settings: WorkerSettings connection: aio_pika.abc.AbstractConnection | None = None channel: aio_pika.abc.AbstractChannel | None = None mongo_client: AsyncIOMotorClient | None = None deps: WorkerDependencies | None = None stop_event: asyncio.Event | None = None async def setup(self) -> None: self.stop_event = asyncio.Event() # ---- Mongo ---- self.mongo_client = AsyncIOMotorClient( self.settings.mongo_uri, uuidRepresentation="standard", tz_aware=True, ) db = self.mongo_client[self.settings.mongo_db] self.deps = WorkerDependencies( db=db, jobs=db["jobs"], ai_executions=db["ai_executions"], artifacts=db["artifacts"], settings=self.settings, worker_id=self.settings.worker_id, ) # ---- Optional AI gateway wiring ---- if self.settings.enable_ai_handlers: try: from asf_ai.app.gateway import AIGateway, AIGatewayConfig # type: ignore[import-not-found] self.deps.ai_gateway = AIGateway(config=AIGatewayConfig( default_provider=self.settings.ai_gateway_provider, )) log.info( "ai_gateway_attached", provider=self.settings.ai_gateway_provider ) except Exception as exc: # noqa: BLE001 log.warn("ai_gateway_unavailable", exc=str(exc)) # ---- RabbitMQ ---- self.connection = await aio_pika.connect_robust(self.settings.rabbitmq_uri) self.channel = await self.connection.channel() await self.channel.set_qos(prefetch_count=self.settings.prefetch_count) jobs_exchange = await self.channel.declare_exchange( "asf.jobs", type=aio_pika.ExchangeType.DIRECT, durable=True ) await self.channel.declare_exchange( "asf.events", type=aio_pika.ExchangeType.TOPIC, durable=True ) for qname in self.settings.queues: q = await self.channel.declare_queue(qname, durable=True) await q.bind(jobs_exchange, routing_key=qname) await q.consume(self._make_consumer_callback(qname)) log.info("subscribed", queue=qname) def _make_consumer_callback(self, qname: str): handler = HANDLERS_BY_QUEUE.get(qname, noop_handler) async def _cb(message: aio_pika.abc.AbstractIncomingMessage) -> None: async with message.process(requeue=True, reject_on_redelivered=False): assert self.deps is not None try: envelope = JobRequestEnvelope.model_validate_json(message.body) except Exception: # noqa: BLE001 log.exception("bad_message_body", queue=qname) return # ack & drop poison # 1. claim ownership in Mongo await _claim_running(self.deps, envelope) # 2. run the handler try: result: HandlerResult = await handler(envelope, self.deps) except Exception as exc: # noqa: BLE001 log.exception("handler_unhandled_exception", queue=qname) result = HandlerResult( status=JobStatus.FAILED, error=f"{type(exc).__name__}: {exc}", ) # 3. persist final state try: await mark_job_done(self.deps, envelope.job_id, result) except Exception: # noqa: BLE001 log.exception("job_persistence_failed", job_id=str(envelope.job_id)) return _cb async def run_forever(self) -> None: assert self.stop_event is not None log.info( "asf_worker_started", queues=self.settings.queues, worker_id=self.settings.worker_id, ) await self.stop_event.wait() async def close(self) -> None: log.info("asf_worker_shutting_down") if self.channel is not None: try: await self.channel.close() except Exception: # noqa: BLE001 pass if self.connection is not None: try: await self.connection.close() except Exception: # noqa: BLE001 pass if self.mongo_client is not None: self.mongo_client.close() def build_consumer(settings: WorkerSettings | None = None) -> Consumer: return Consumer(settings=settings or get_settings()) async def run_worker_async() -> None: settings = get_settings() consumer = build_consumer(settings) await consumer.setup() loop = asyncio.get_running_loop() def _graceful() -> None: if consumer.stop_event is not None: consumer.stop_event.set() for sig in (signal.SIGINT, signal.SIGTERM): try: loop.add_signal_handler(sig, _graceful) except NotImplementedError: # Windows pass try: await consumer.run_forever() finally: await consumer.close() def run_worker() -> None: """Synchronous entry point (e.g. python -m app.worker).""" try: asyncio.run(run_worker_async()) except KeyboardInterrupt: print("Interrupted. Goodbye.", file=sys.stderr) if __name__ == "__main__": run_worker()