import asyncio import logging import uuid from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta import httpx import redis.asyncio as redis import structlog from opentelemetry import trace from sqlalchemy import delete, select from app.db import ( Database, DeliveryOutbox, Dialog, Message, MessageAttachment, SafetyTask, UserIdentity, ) from app.integrations import ( DependencyFailure, OpenLinesClient, S3Client, SafetyClient, fresh_openlines_payload, ) from app.logging_security import redact_event from app.notification_models import ClientUploadDraft from app.notification_service import expire_notifications from app.realtime import RealtimeFanout from app.services import ( ensure_delivery_outbox, load_active_client_profile, load_settings, openlines_delivery_payload, publish_dialog_status, publish_message_status, ) from app.settings import Settings, get_settings from app.telemetry import TelemetryRuntime, add_trace_context, init_telemetry, origin_links log = structlog.get_logger() def configure_logging(level: str, telemetry: TelemetryRuntime | None = None) -> None: logging.basicConfig(level=level, format="%(message)s") if telemetry and telemetry.logging_handler not in logging.getLogger().handlers: telemetry.logging_handler.addFilter( lambda record: not record.name.startswith("opentelemetry") ) logging.getLogger().addHandler(telemetry.logging_handler) structlog.configure( processors=[ structlog.contextvars.merge_contextvars, add_trace_context, redact_event, structlog.processors.TimeStamper(fmt="iso", utc=True, key="timestamp"), structlog.stdlib.add_log_level, structlog.processors.JSONRenderer(), ], logger_factory=structlog.stdlib.LoggerFactory(), wrapper_class=structlog.stdlib.BoundLogger, cache_logger_on_first_use=True, ) def run_worker(service_name: str, target: Callable[[], Awaitable[None]]) -> None: settings = get_settings() telemetry = init_telemetry(service_name) configure_logging(settings.log_level, telemetry) structlog.contextvars.bind_contextvars(**{"service.name": service_name}) try: asyncio.run(target()) finally: if telemetry: telemetry.shutdown() @asynccontextmanager async def worker_http_clients( settings: Settings, ) -> AsyncIterator[tuple[httpx.AsyncClient, httpx.AsyncClient]]: async with ( httpx.AsyncClient() as openlines_http, httpx.AsyncClient(verify=settings.message_safety_ca_file) as safety_http, ): yield openlines_http, safety_http async def delivery_once( db: Database, client: OpenLinesClient, s3: S3Client, fanout: RealtimeFanout, settings: Settings, worker_id: str, batch_size: int = 20, ) -> int: tracer = trace.get_tracer("han.api.delivery-worker") async with db.sessions() as session: rows = ( ( await session.execute( select(DeliveryOutbox) .where( DeliveryOutbox.status.in_(["pending", "retry"]), DeliveryOutbox.next_attempt_at <= datetime.now(UTC), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) ids = [row.id for row in rows] for row in rows: row.status = "processing" row.locked_at = datetime.now(UTC) row.locked_by = worker_id await session.commit() for row_id in ids: async with db.sessions() as session: row = await session.get(DeliveryOutbox, row_id, with_for_update=True) if row is None: continue with ( tracer.start_as_current_span( "delivery.process", links=origin_links(row.traceparent), ), structlog.contextvars.bound_contextvars( delivery_outbox_id=str(row.id), message_id=str(row.message_id), ), ): message = None dialog = None try: payload = await fresh_openlines_payload(row.payload_json, s3) await client.send(row.message_id, payload, f"worker-{worker_id}") row.status = "delivered" message = await session.get(Message, row.message_id) if message: message.delivery_status = "delivered" dialog = await session.get(Dialog, message.dialog_id) if dialog: dialog.status = "waiting_for_company" dialog.last_message_at = datetime.now(UTC) except DependencyFailure: row.attempt_count += 1 row.status = "dead_letter" if row.attempt_count >= 12 else "retry" row.next_attempt_at = datetime.now(UTC) + timedelta( seconds=min(3600, 2**row.attempt_count) ) row.last_error_code = "dependency_unavailable" row.locked_at = None row.locked_by = None await session.commit() if message: await publish_message_status(fanout, message, settings) if dialog: await publish_dialog_status(fanout, dialog) return len(ids) async def safety_once( db: Database, safety: SafetyClient, s3: S3Client, fanout: RealtimeFanout, settings: Settings, worker_id: str, batch_size: int = 20, ) -> int: tracer = trace.get_tracer("han.api.safety-recovery-worker") async with db.sessions() as session: rows = ( ( await session.execute( select(SafetyTask) .where( SafetyTask.status.in_(["polling", "failed"]), SafetyTask.next_poll_at <= datetime.now(UTC), SafetyTask.expires_at > datetime.now(UTC), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) ids = [row.id for row in rows] for row in rows: row.locked_at, row.locked_by = datetime.now(UTC), worker_id await session.commit() for task_id in ids: async with db.sessions() as session: task = await session.get(SafetyTask, task_id, with_for_update=True) if task is None: continue message = None try: with tracer.start_as_current_span( "safety.recovery.poll", links=origin_links(task.traceparent), ): verdict = await safety.poll( task.poll_location, f"worker-{worker_id}", ) message = await session.get(Message, task.message_id) attachment = ( await session.get(MessageAttachment, task.attachment_id) if task.attachment_id else None ) if verdict["_status"] == 200 and message: message.safety_processing_mode = verdict["processing_mode"] message.safety_config_version = verdict["config_version"] message.safety_rules_version = verdict["rules_version"] if attachment and attachment.quarantine_object_key: destination = f"attachments/dialogs/{message.dialog_id}/{attachment.id}" await s3.promote( attachment.quarantine_object_key, destination, version_id=attachment.quarantine_version_id or "", etag=attachment.quarantine_etag or "", ) attachment.storage_bucket = s3.settings.selectel_s3_bucket_attachments attachment.object_key = destination attachment.quarantine_object_key = None attachment.scan_status = ( "bypassed" if verdict["processing_mode"] == "mock" else "clean" ) message.safety_status = "allowed" task.status = "completed" dialog = await session.get(Dialog, message.dialog_id) user = ( await session.get(UserIdentity, dialog.user_id) if dialog else None ) if dialog and user: profile = await load_active_client_profile(session, user.id) await ensure_delivery_outbox( session, message_id=message.id, external_chat_id=dialog.id, payload_json=openlines_delivery_payload( message=message, dialog_id=dialog.id, user=user, profile=profile, attachment=attachment, ), next_attempt_at=datetime.now(UTC), traceparent=task.traceparent, ) elif verdict["_status"] == 403 and message: message.safety_processing_mode = verdict["processing_mode"] message.safety_config_version = verdict["config_version"] message.safety_rules_version = verdict["rules_version"] message.text = "" message.safety_status = "blocked" message.delivery_status = "rejected" if attachment and attachment.quarantine_object_key: await s3.delete_quarantine(attachment.quarantine_object_key) attachment.scan_status = "infected" task.status = "completed" else: if verdict["_status"] == 202: task.poll_location = verdict["_location"] task.next_poll_at = datetime.now(UTC) + timedelta(seconds=2) except DependencyFailure as exc: task.attempt_count += 1 terminal = exc.terminal or ( exc.code == "task_not_found" and task.attempt_count >= 2 ) task.status = "terminal_failed" if terminal else "failed" task.last_error_code = exc.code if terminal and message: message.delivery_status = "failed" task.next_poll_at = datetime.now(UTC) + timedelta( seconds=min(300, 2**task.attempt_count) ) task.locked_at = None task.locked_by = None await session.commit() if message: await publish_message_status(fanout, message, settings) return len(ids) async def cleanup_once(db: Database, s3: S3Client, batch_size: int = 100) -> int: async with db.sessions() as session: rows = ( ( await session.execute( select(MessageAttachment) .where( MessageAttachment.record_status == "A", MessageAttachment.quarantine_object_key.is_not(None), MessageAttachment.upload_expires_at < datetime.now(UTC), MessageAttachment.message_id.is_(None), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) for row in rows: if row.quarantine_object_key: await s3.delete_quarantine(row.quarantine_object_key) row.record_status = "D" row.status_changed_at = datetime.now(UTC) row.status_change_reason = "expired_quarantine_cleanup" await session.commit() return len(rows) async def loop(kind: str) -> None: settings = get_settings() db = Database(settings.database_url) s3 = S3Client(settings) redis_rt = redis.from_url(settings.redis_realtime_url, decode_responses=True) fanout = RealtimeFanout(redis_rt) worker_id = f"{kind}-{uuid.uuid4()}" try: async with worker_http_clients(settings) as (openlines_http, safety_http): safety = SafetyClient(settings, safety_http) openlines = OpenLinesClient(settings, openlines_http) while True: count = 0 if kind == "delivery": count = await delivery_once(db, openlines, s3, fanout, settings, worker_id) elif kind == "safety": count = await safety_once(db, safety, s3, fanout, settings, worker_id) else: count = await cleanup_once(db, s3) if not count: await asyncio.sleep(settings.worker_poll_interval_sec) finally: await redis_rt.aclose() await db.close() async def notification_expire_loop() -> None: settings = get_settings() db = Database(settings.database_url) try: while True: async with db.sessions() as session: snapshot = await load_settings(session) run_at = snapshot.values["notification.expire_job.run_at"] hour, minute = (int(value) for value in run_at.split(":", 1)) now = datetime.now(UTC) target = now.replace(hour=hour, minute=minute, second=0, microsecond=0) if target <= now: target += timedelta(days=1) await asyncio.sleep((target - now).total_seconds()) async with db.sessions() as session: personal, guest = await expire_notifications(session) log.info( "notification.expired_batch", personal_count=personal, guest_count=guest, ) finally: await db.close() async def notification_draft_cleanup_once(db: Database, s3: S3Client) -> int: async with db.sessions() as session: snapshot = await load_settings(session) cutoff = datetime.now(UTC) - timedelta( days=snapshot.integer("notification.upload_draft.ttl_days") ) rows = ( ( await session.execute( select(ClientUploadDraft) .where(ClientUploadDraft.created_at < cutoff) .with_for_update(skip_locked=True) .limit(100) ) ) .scalars() .all() ) for row in rows: if row.state != "submitted": if row.quarantine_object_key: await s3.delete_quarantine(row.quarantine_object_key) elif row.object_key: await s3.delete(row.storage_bucket, row.object_key) await session.execute( delete(ClientUploadDraft).where(ClientUploadDraft.id == row.id) ) await session.commit() return len(rows) async def notification_draft_cleanup_loop() -> None: settings = get_settings() db = Database(settings.database_url) s3 = S3Client(settings) try: while True: count = await notification_draft_cleanup_once(db, s3) if count < 100: await asyncio.sleep(86400) finally: await db.close() def delivery_main() -> None: run_worker("delivery-worker", lambda: loop("delivery")) def safety_main() -> None: run_worker("safety-recovery-worker", lambda: loop("safety")) def cleanup_main() -> None: run_worker("cleanup-worker", lambda: loop("cleanup")) def notification_expire_main() -> None: run_worker("notification-expire-worker", notification_expire_loop) def notification_draft_cleanup_main() -> None: run_worker("notification-draft-cleanup-worker", notification_draft_cleanup_loop)