company-ai-platform/app/modules/subscriptions/services/delivery.py

from dataclasses import dataclass
from datetime import UTC, datetime, time, timedelta
from typing import Any, Protocol
from uuid import uuid4
from zoneinfo import ZoneInfo

import httpx
from fastapi import HTTPException, status
from sqlalchemy import and_, func, or_, select, update
from sqlalchemy.orm import Session

from app.core.http.pagination import bounded_limit
from app.core.utils.time import utc_now
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.feishu.errors import FeishuAPIError
from app.modules.feishu_users.constants import FeishuUserRole, FeishuUserStatus
from app.modules.feishu_users.models import FeishuUser
from app.modules.subscriptions.constants import (
    DAILY_DELIVERY_LIMIT_REACHED,
    DELIVERY_LEASE_SECONDS,
    DELIVERY_NOT_FOUND,
    DELIVERY_RETRY_DELAYS_SECONDS,
    MAX_DAILY_DELIVERIES,
    PushDeliveryStatus,
    PushSubscriptionStatus,
    SubscriptionAuditAction,
    SubscriptionTargetType,
)
from app.modules.subscriptions.models import PushDelivery, PushSubscription
from app.modules.subscriptions.services.schedule import (
    is_in_quiet_hours,
    next_quiet_end,
)


@dataclass(frozen=True, slots=True)
class DeliveryGenerationRequest:
    prompt: str
    owner_id: int | None
    use_personal_context: bool
    use_company_rules: bool
    allow_tools: bool = False
    record_history: bool = False
    infer_preferences: bool = False
    write_memory: bool = False


@dataclass(frozen=True, slots=True)
class DeliverySendRequest:
    receive_id: str
    receive_id_type: str
    tenant_key: str
    text: str
    uuid: str


class DeliveryGenerator(Protocol):
    def generate(self, request: DeliveryGenerationRequest) -> str:
        """Generate delivery text without side effects or business-data tools."""


class DeliverySender(Protocol):
    def send(self, request: DeliverySendRequest) -> dict[str, Any]:
        """Send a message and return the provider response."""


class RetryableDeliveryError(RuntimeError):
    """A temporary generation or provider failure."""


class PermanentDeliveryError(RuntimeError):
    """A delivery failure that must not be retried."""


class DeliveryService:
    """Process durable deliveries with fencing and database-driven retry timing."""

    def __init__(
        self,
        db: Session,
        *,
        generator: DeliveryGenerator,
        sender: DeliverySender,
        lease_seconds: int = DELIVERY_LEASE_SECONDS,
    ):
        self.db = db
        self.generator = generator
        self.sender = sender
        self.lease_seconds = lease_seconds

    def process(
        self,
        delivery_code: str,
        *,
        now: datetime | None = None,
        worker_id: str = "subscription-delivery",
    ) -> PushDelivery:
        current = _naive_utc(now or utc_now())
        existing = self._get(delivery_code)
        if existing.status in {
            PushDeliveryStatus.SENT,
            PushDeliveryStatus.FAILED,
            PushDeliveryStatus.SKIPPED,
        }:
            return existing

        lock_owner = f"{worker_id}:{uuid4().hex}"
        if not self._claim(existing.id, lock_owner, current):
            self.db.rollback()
            return self._get(delivery_code)
        delivery = self._get(delivery_code)
        subscription, owner = self._load_context(delivery.subscription_id)

        skip_reason = self._skip_reason(subscription, owner)
        if skip_reason is not None:
            return self._finish_skipped(delivery, lock_owner, owner, skip_reason)

        if (
            owner is not None
            and is_in_quiet_hours(
                current,
                owner.timezone,
                owner.quiet_hours_start,
                owner.quiet_hours_end,
            )
        ):
            quiet_end = next_quiet_end(
                current,
                owner.timezone,
                owner.quiet_hours_start,
                owner.quiet_hours_end,
            )
            return self._defer_for_quiet_hours(delivery, lock_owner, quiet_end)

        if not self._start_attempt(delivery.id, lock_owner, current):
            self.db.rollback()
            return self._get(delivery_code)
        delivery = self._get(delivery_code)
        try:
            content = delivery.rendered_content
            if content is None:
                content = str(
                    self.generator.generate(
                        self._generation_request(subscription)
                    )
                ).strip()
                if not content:
                    raise PermanentDeliveryError("Delivery generator returned empty content")
                self._save_content(delivery.id, lock_owner, content, current)
            subscription, owner = self._lock_send_context(delivery.subscription_id)
            skip_reason = self._skip_reason(subscription, owner)
            if skip_reason is None and owner is not None:
                if self._sent_today(owner, current) >= MAX_DAILY_DELIVERIES:
                    skip_reason = DAILY_DELIVERY_LIMIT_REACHED
            if skip_reason is not None:
                return self._finish_skipped(
                    delivery,
                    lock_owner,
                    owner,
                    skip_reason,
                )
            response = self.sender.send(
                DeliverySendRequest(
                    receive_id=(
                        owner.open_id
                        if subscription.target_type == SubscriptionTargetType.USER
                        else subscription.target_id
                    ),
                    receive_id_type=(
                        "open_id"
                        if subscription.target_type == SubscriptionTargetType.USER
                        else "chat_id"
                    ),
                    tenant_key=owner.tenant_key,
                    text=content,
                    uuid=delivery.message_uuid,
                )
            )
            self._validate_provider_response(response)
        except Exception as exc:
            self.db.rollback()
            return self._finish_failure(delivery_code, lock_owner, current, exc)
        return self._finish_sent(
            delivery_code,
            lock_owner,
            current,
            owner,
            response,
        )

    def process_due(
        self,
        *,
        now: datetime | None = None,
        limit: int = 100,
        worker_id: str = "subscription-delivery",
    ) -> list[PushDelivery]:
        current = _naive_utc(now or utc_now())
        stmt = (
            select(PushDelivery.code)
            .where(
                or_(
                    and_(
                        PushDelivery.status.in_(
                            [
                                PushDeliveryStatus.PENDING,
                                PushDeliveryStatus.RETRY,
                            ]
                        ),
                        PushDelivery.next_attempt_at.is_not(None),
                        PushDelivery.next_attempt_at <= current,
                    ),
                    and_(
                        PushDelivery.status == PushDeliveryStatus.PROCESSING,
                        PushDelivery.locked_until.is_not(None),
                        PushDelivery.locked_until <= current,
                    ),
                )
            )
            .order_by(PushDelivery.next_attempt_at.asc(), PushDelivery.id.asc())
            .limit(bounded_limit(limit))
        )
        codes = list(self.db.execute(stmt).scalars())
        return [
            self.process(
                code,
                now=current,
                worker_id=worker_id,
            )
            for code in codes
        ]

    def _claim(self, delivery_id: int, lock_owner: str, current: datetime) -> bool:
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.id == delivery_id,
                or_(
                    and_(
                        PushDelivery.status.in_(
                            [
                                PushDeliveryStatus.PENDING,
                                PushDeliveryStatus.RETRY,
                            ]
                        ),
                        PushDelivery.next_attempt_at.is_not(None),
                        PushDelivery.next_attempt_at <= current,
                    ),
                    and_(
                        PushDelivery.status == PushDeliveryStatus.PROCESSING,
                        PushDelivery.locked_until.is_not(None),
                        PushDelivery.locked_until <= current,
                    ),
                ),
                or_(
                    PushDelivery.locked_until.is_(None),
                    PushDelivery.locked_until <= current,
                ),
            )
            .values(
                status=PushDeliveryStatus.PROCESSING,
                locked_by=lock_owner,
                locked_until=current + timedelta(seconds=self.lease_seconds),
            )
            .execution_options(synchronize_session=False)
        )
        self.db.commit()
        return result.rowcount == 1

    def _start_attempt(
        self,
        delivery_id: int,
        lock_owner: str,
        current: datetime,
    ) -> bool:
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.id == delivery_id,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                attempt_count=PushDelivery.attempt_count + 1,
                locked_until=current + timedelta(seconds=self.lease_seconds),
            )
            .execution_options(synchronize_session=False)
        )
        self.db.commit()
        return result.rowcount == 1

    def _save_content(
        self,
        delivery_id: int,
        lock_owner: str,
        content: str,
        current: datetime,
    ) -> None:
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.id == delivery_id,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                rendered_content=content,
                locked_until=current + timedelta(seconds=self.lease_seconds),
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            self.db.rollback()
            raise RetryableDeliveryError("Delivery lease was lost")
        self.db.commit()

    def _finish_sent(
        self,
        delivery_code: str,
        lock_owner: str,
        current: datetime,
        owner: FeishuUser | None,
        response: dict[str, Any],
    ) -> PushDelivery:
        provider_message_id = _provider_message_id(response)
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.code == delivery_code,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                status=PushDeliveryStatus.SENT,
                next_attempt_at=None,
                provider_message_id=provider_message_id,
                last_error=None,
                sent_at=current,
                locked_by=None,
                locked_until=None,
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            self.db.rollback()
            return self._get(delivery_code)
        record = self._get(delivery_code)
        self._audit(
            owner,
            record,
            SubscriptionAuditAction.DELIVERY_SENT,
            {"status": PushDeliveryStatus.SENT},
        )
        self.db.commit()
        self.db.refresh(record)
        return record

    def _finish_failure(
        self,
        delivery_code: str,
        lock_owner: str,
        current: datetime,
        exc: Exception,
    ) -> PushDelivery:
        record = self._get(delivery_code)
        retryable = _is_retryable(exc)
        retry_index = record.attempt_count - 1
        will_retry = retryable and 0 <= retry_index < len(
            DELIVERY_RETRY_DELAYS_SECONDS
        )
        next_attempt_at = (
            current + timedelta(seconds=DELIVERY_RETRY_DELAYS_SECONDS[retry_index])
            if will_retry
            else None
        )
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.code == delivery_code,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                status=(
                    PushDeliveryStatus.RETRY
                    if will_retry
                    else PushDeliveryStatus.FAILED
                ),
                next_attempt_at=next_attempt_at,
                last_error=f"{type(exc).__name__}: {exc}"[:2000],
                locked_by=None,
                locked_until=None,
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            self.db.rollback()
            return self._get(delivery_code)
        record = self._get(delivery_code)
        if not will_retry:
            _, owner = self._load_context(record.subscription_id)
            self._audit(
                owner,
                record,
                SubscriptionAuditAction.DELIVERY_FAILED,
                {
                    "status": PushDeliveryStatus.FAILED,
                    "attempt_count": record.attempt_count,
                },
            )
        self.db.commit()
        self.db.refresh(record)
        return record

    def _finish_skipped(
        self,
        delivery: PushDelivery,
        lock_owner: str,
        owner: FeishuUser | None,
        reason: str,
    ) -> PushDelivery:
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.id == delivery.id,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                status=PushDeliveryStatus.SKIPPED,
                next_attempt_at=None,
                last_error=reason,
                locked_by=None,
                locked_until=None,
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            self.db.rollback()
            return self._get(delivery.code)
        record = self._get(delivery.code)
        self._audit(
            owner,
            record,
            SubscriptionAuditAction.DELIVERY_SKIPPED,
            {"status": PushDeliveryStatus.SKIPPED, "reason": reason},
        )
        self.db.commit()
        self.db.refresh(record)
        return record

    def _defer_for_quiet_hours(
        self,
        delivery: PushDelivery,
        lock_owner: str,
        quiet_end: datetime,
    ) -> PushDelivery:
        status_value = (
            PushDeliveryStatus.PENDING
            if delivery.attempt_count == 0
            else PushDeliveryStatus.RETRY
        )
        result = self.db.execute(
            update(PushDelivery)
            .where(
                PushDelivery.id == delivery.id,
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_by == lock_owner,
            )
            .values(
                status=status_value,
                next_attempt_at=quiet_end,
                locked_by=None,
                locked_until=None,
            )
            .execution_options(synchronize_session=False)
        )
        if result.rowcount != 1:
            self.db.rollback()
            return self._get(delivery.code)
        self.db.commit()
        return self._get(delivery.code)

    def _generation_request(
        self,
        subscription: PushSubscription,
    ) -> DeliveryGenerationRequest:
        personal = subscription.target_type == SubscriptionTargetType.USER
        return DeliveryGenerationRequest(
            prompt=subscription.prompt,
            owner_id=subscription.owner_id if personal else None,
            use_personal_context=personal,
            use_company_rules=not personal,
        )

    def _skip_reason(
        self,
        subscription: PushSubscription | None,
        owner: FeishuUser | None,
    ) -> str | None:
        if subscription is None:
            return "Subscription was removed"
        if owner is None or owner.status != FeishuUserStatus.ACTIVE:
            return "Feishu user is disabled"
        if subscription.status not in {
            PushSubscriptionStatus.ACTIVE,
            PushSubscriptionStatus.COMPLETED,
        }:
            return "Subscription is not active"
        if subscription.target_type == SubscriptionTargetType.USER:
            if subscription.target_id != owner.open_id:
                return "Private subscription target no longer matches its owner"
            return None
        if subscription.target_type == SubscriptionTargetType.CHAT:
            if owner.role != FeishuUserRole.ADMIN or not subscription.target_id:
                return "Group subscription is no longer authorized"
            return None
        return "Unsupported subscription target"

    def _load_context(
        self,
        subscription_id: int,
    ) -> tuple[PushSubscription | None, FeishuUser | None]:
        subscription = self.db.get(PushSubscription, subscription_id)
        owner = (
            self.db.get(FeishuUser, subscription.owner_id)
            if subscription is not None
            else None
        )
        return subscription, owner

    def _lock_send_context(
        self,
        subscription_id: int,
    ) -> tuple[PushSubscription | None, FeishuUser | None]:
        """Recheck authorization and serialize the final per-owner send decision."""

        subscription = self.db.execute(
            select(PushSubscription)
            .where(PushSubscription.id == subscription_id)
            .with_for_update()
            .execution_options(populate_existing=True)
        ).scalar_one_or_none()
        owner = (
            self.db.execute(
                select(FeishuUser)
                .where(FeishuUser.id == subscription.owner_id)
                .with_for_update()
                .execution_options(populate_existing=True)
            ).scalar_one_or_none()
            if subscription is not None
            else None
        )
        return subscription, owner

    def _sent_today(self, owner: FeishuUser, current: datetime) -> int:
        zone = ZoneInfo(owner.timezone)
        local_now = current.replace(tzinfo=UTC).astimezone(zone)
        local_start = datetime.combine(local_now.date(), time.min, tzinfo=zone)
        local_end = local_start + timedelta(days=1)
        start_utc = local_start.astimezone(UTC).replace(tzinfo=None)
        end_utc = local_end.astimezone(UTC).replace(tzinfo=None)
        return int(
            self.db.scalar(
                select(func.count())
                .select_from(PushDelivery)
                .join(PushSubscription)
                .where(
                    PushSubscription.owner_id == owner.id,
                    PushDelivery.status == PushDeliveryStatus.SENT,
                    PushDelivery.sent_at >= start_utc,
                    PushDelivery.sent_at < end_utc,
                )
            )
            or 0
        )

    def _get(self, delivery_code: str) -> PushDelivery:
        record = self.db.execute(
            select(PushDelivery).where(PushDelivery.code == delivery_code)
        ).scalar_one_or_none()
        if record is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=DELIVERY_NOT_FOUND,
            )
        return record

    def _audit(
        self,
        owner: FeishuUser | None,
        delivery: PushDelivery,
        action: str,
        response: dict[str, Any],
    ) -> None:
        AuditService(self.db).record(
            AuditLogCreate(
                actor=owner.code if owner is not None else "subscription-system",
                source="subscriptions",
                action=action,
                target_type="push-delivery",
                target_id=delivery.code,
                response_payload=response,
            )
        )

    @staticmethod
    def _validate_provider_response(response: dict[str, Any]) -> None:
        if not isinstance(response, dict):
            raise RetryableDeliveryError("Message provider returned an invalid response")
        if "code" in response and response.get("code") != 0:
            raise RetryableDeliveryError(
                f"Message provider returned business code {response.get('code')}"
            )


def _provider_message_id(response: dict[str, Any]) -> str | None:
    direct = response.get("message_id")
    nested = response.get("data")
    value = direct or (nested.get("message_id") if isinstance(nested, dict) else None)
    return str(value) if value is not None else None


def _is_retryable(exc: Exception) -> bool:
    if isinstance(exc, PermanentDeliveryError):
        return False
    if isinstance(exc, RetryableDeliveryError):
        return True
    if isinstance(exc, FeishuAPIError):
        return exc.retryable
    if isinstance(exc, httpx.HTTPStatusError):
        code = exc.response.status_code
        return code == 429 or code >= 500
    if isinstance(exc, (httpx.TimeoutException, httpx.NetworkError)):
        return True
    if isinstance(exc, HTTPException):
        return exc.status_code == 429 or exc.status_code >= 500
    return True


def _naive_utc(value: datetime) -> datetime:
    if value.tzinfo is None:
        return value
    return value.astimezone(UTC).replace(tzinfo=None)