agentic.email/backend/app/services/handoff_delivery.py

import asyncio
import email.utils
import logging
from datetime import datetime, timezone

from bson import ObjectId

from app.database import mongo_db

logger = logging.getLogger(__name__)

JOB_MIRROR = "MIRROR_TO_HUMAN"
JOB_RELAY = "RELAY_TO_EXTERNAL"

STATUS_PENDING = "PENDING"
STATUS_SENDING = "SENDING"
STATUS_SENT = "SENT"
STATUS_FAILED = "FAILED"
STATUS_AMBIGUOUS = "AMBIGUOUS"


class HandoffPreflightError(RuntimeError):
    """A deterministic error that happened before SMTP submission."""


def _now() -> datetime:
    return datetime.now(timezone.utc)


async def enqueue_handoff_delivery(
    database,
    *,
    kind: str,
    account_id: str,
    handoff_id: ObjectId,
    source_message_id: ObjectId,
    agent_email: str,
) -> dict:
    idempotency_key = f"{kind}:{account_id}:{handoff_id}:{source_message_id}"
    now = _now()
    document = {
        "idempotency_key": idempotency_key,
        "kind": kind,
        "account_id": account_id,
        "handoff_id": handoff_id,
        "source_message_id": source_message_id,
        "agent_email": agent_email.lower(),
        "generated_message_id": email.utils.make_msgid(
            domain=agent_email.rsplit("@", 1)[-1].lower()
        ),
        "status": STATUS_PENDING,
        "attempt_count": 0,
        "next_attempt_at": now,
        "created_at": now,
        "updated_at": now,
    }
    await database.handoff_delivery_jobs.update_one(
        {"idempotency_key": idempotency_key},
        {"$setOnInsert": document},
        upsert=True,
    )
    return await database.handoff_delivery_jobs.find_one(
        {"idempotency_key": idempotency_key}
    )


async def process_handoff_delivery_once(database=mongo_db) -> bool:
    now = _now()
    job = await database.handoff_delivery_jobs.find_one_and_update(
        {
            "status": STATUS_PENDING,
            "next_attempt_at": {"$lte": now},
        },
        {
            "$set": {
                "status": STATUS_SENDING,
                "started_at": now,
                "updated_at": now,
            },
            "$inc": {"attempt_count": 1},
        },
        sort=[("created_at", 1), ("_id", 1)],
        return_document=True,
    )
    if not job:
        return False

    try:
        from app.services.human_handoff_service import deliver_handoff_job

        target_message_id = await deliver_handoff_job(database, job)
    except HandoffPreflightError as error:
        logger.error(
            "handoff_delivery_preflight_failed job=%s kind=%s error=%s",
            job["_id"],
            job["kind"],
            error,
        )
        await database.handoff_delivery_jobs.update_one(
            {"_id": job["_id"], "status": STATUS_SENDING},
            {
                "$set": {
                    "status": STATUS_FAILED,
                    "error": str(error),
                    "failed_at": _now(),
                    "updated_at": _now(),
                }
            },
        )
        if job["kind"] == JOB_RELAY:
            await database.messages.update_one(
                {"_id": job["source_message_id"]},
                {
                    "$set": {
                        "relay_status": STATUS_FAILED,
                        "relay_error": str(error),
                        "relay_failed_at": _now(),
                    }
                },
            )
    except Exception as error:
        # SMTP exceptions can be ambiguous: the remote MTA may have accepted the
        # message before the acknowledgement was lost. Never auto-retry these.
        logger.exception(
            "handoff_delivery_ambiguous job=%s kind=%s",
            job["_id"],
            job["kind"],
        )
        await database.handoff_delivery_jobs.update_one(
            {"_id": job["_id"], "status": STATUS_SENDING},
            {
                "$set": {
                    "status": STATUS_AMBIGUOUS,
                    "error": str(error),
                    "ambiguous_at": _now(),
                    "updated_at": _now(),
                }
            },
        )
        if job["kind"] == JOB_RELAY:
            await database.messages.update_one(
                {"_id": job["source_message_id"]},
                {
                    "$set": {
                        "relay_status": STATUS_AMBIGUOUS,
                        "relay_error": str(error),
                        "relay_failed_at": _now(),
                    }
                },
            )
    else:
        await database.handoff_delivery_jobs.update_one(
            {"_id": job["_id"], "status": STATUS_SENDING},
            {
                "$set": {
                    "status": STATUS_SENT,
                    "target_message_id": target_message_id,
                    "sent_at": _now(),
                    "updated_at": _now(),
                }
            },
        )
    return True


class HandoffDeliveryWorker:
    def __init__(self, poll_seconds: float = 1.0):
        self.poll_seconds = poll_seconds
        self._task: asyncio.Task | None = None
        self._stop = asyncio.Event()

    async def start(self) -> None:
        if self._task and not self._task.done():
            return
        # A process that died during SMTP submission cannot know whether the
        # remote MTA accepted it. Quarantine instead of risking a duplicate.
        await mongo_db.handoff_delivery_jobs.update_many(
            {"status": STATUS_SENDING},
            {
                "$set": {
                    "status": STATUS_AMBIGUOUS,
                    "error": "Worker restarted during SMTP submission.",
                    "ambiguous_at": _now(),
                    "updated_at": _now(),
                }
            },
        )
        self._stop.clear()
        self._task = asyncio.create_task(self._run())

    async def stop(self) -> None:
        self._stop.set()
        if self._task:
            await self._task
            self._task = None

    async def _run(self) -> None:
        while not self._stop.is_set():
            processed = False
            try:
                processed = await process_handoff_delivery_once()
            except Exception:
                logger.exception("handoff_delivery_worker_iteration_failed")
            if processed:
                continue
            try:
                await asyncio.wait_for(
                    self._stop.wait(),
                    timeout=self.poll_seconds,
                )
            except TimeoutError:
                pass


handoff_delivery_worker = HandoffDeliveryWorker()