company-ai-platform/app/modules/observability/service.py

from collections.abc import Callable
from datetime import timedelta
from typing import Any

from sqlalchemy import and_, func, or_, select, text
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.core.constants import ActorValue
from app.core.utils.time import utc_now
from app.modules.ai_memory.service import AIMemoryService
from app.modules.audit.constants import (
    AuditAction,
    AuditRiskLevel,
    AuditSource,
    AuditTargetType,
)
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.events.constants import EventStatus
from app.modules.events.services import EventService
from app.modules.feishu.app_tickets import FeishuAppTicketService
from app.modules.feishu.constants import FeishuAppType
from app.modules.feishu_users.constants import FeishuUserStatus
from app.modules.feishu_users.models import FeishuUser
from app.modules.observability.constants import (
    HeartbeatStatus,
    ObservabilityKey,
    ObservabilityMetricKey,
    ObservabilityStatus,
)
from app.modules.observability.models import SystemHeartbeat
from app.modules.workflows.constants import WorkflowStatus
from app.modules.workflows.service import WorkflowService
from app.modules.subscriptions.constants import (
    PushDeliveryStatus,
    PushSubscriptionStatus,
)
from app.modules.subscriptions.models import PushDelivery, PushSubscription


class ObservabilityService:
    """Build health, readiness, and JSON metrics for V3 operations."""

    def __init__(self, db: Session):
        self.db = db

    def live(self) -> dict[str, str]:
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def ready(self) -> dict[str, Any]:
        checks = {
            ObservabilityKey.DATABASE: self._safe_call(self._database_check),
            ObservabilityKey.REDIS: self._safe_call(self._redis_check),
            ObservabilityKey.EVENTS: self._safe_call(self._events_check),
            ObservabilityKey.WORKFLOWS: self._safe_call(self._workflows_check),
            ObservabilityKey.HEARTBEATS: self._safe_call(self._heartbeats_check),
            "feishu_subscriptions": self._safe_call(
                self._feishu_subscriptions_check
            ),
        }
        statuses = {item[ObservabilityKey.STATUS] for item in checks.values()}
        if statuses & {ObservabilityStatus.ERROR, ObservabilityStatus.DEGRADED}:
            overall_status = ObservabilityStatus.DEGRADED
        else:
            overall_status = ObservabilityStatus.OK
        return {
            ObservabilityKey.STATUS: overall_status,
            ObservabilityKey.CHECKS: checks,
        }

    def metrics(self) -> dict[str, Any]:
        return {
            ObservabilityKey.METRICS: {
                ObservabilityKey.EVENTS: self._safe_call(
                    lambda: EventService(self.db).count_by_status()
                ),
                ObservabilityKey.WORKFLOWS: self._safe_call(
                    lambda: WorkflowService(self.db).count_by_status()
                ),
                ObservabilityKey.AI_MEMORY: self._safe_call(
                    lambda: AIMemoryService(self.db).count_by_status()
                ),
                ObservabilityKey.HEARTBEATS: self._safe_call(
                    self.heartbeat_summary
                ),
                "feishu_users": self._safe_call(self._feishu_user_metrics),
                "subscriptions": self._safe_call(self._subscription_metrics),
            }
        }

    def record_heartbeat(
        self,
        component: str,
        instance_id: str,
        status_value: str = HeartbeatStatus.OK,
        actor: str = ActorValue.SYSTEM,
    ) -> dict[str, Any]:
        now = utc_now()
        record = self._upsert_heartbeat(component, instance_id, status_value, now)
        AuditService(self.db).record(
            AuditLogCreate(
                actor=actor,
                source=AuditSource.OBSERVABILITY,
                action=AuditAction.HEARTBEAT,
                target_type=AuditTargetType.HEARTBEAT,
                target_id=f"{component}:{instance_id}",
                risk_level=AuditRiskLevel.LOW,
                response_payload={
                    ObservabilityMetricKey.COMPONENT: component,
                    ObservabilityMetricKey.INSTANCE_ID: instance_id,
                    ObservabilityKey.STATUS: status_value,
                    ObservabilityMetricKey.LAST_SEEN_AT: record.last_seen_at.isoformat(),
                },
            )
        )
        self.db.commit()
        self.db.refresh(record)
        return self._serialize_heartbeat(record)

    def heartbeat_summary(self) -> dict[str, Any]:
        records = list(
            self.db.execute(
                select(SystemHeartbeat)
                .where(SystemHeartbeat.last_seen_at >= self._heartbeat_retention_threshold())
                .order_by(SystemHeartbeat.component.asc(), SystemHeartbeat.instance_id.asc())
            ).scalars()
        )
        threshold = self._heartbeat_stale_threshold()
        stale = [item for item in records if item.last_seen_at < threshold]
        active = len(records) - len(stale)
        last_seen_at = max((item.last_seen_at for item in records), default=None)
        return {
            ObservabilityMetricKey.TOTAL: len(records),
            ObservabilityMetricKey.ACTIVE: active,
            ObservabilityMetricKey.STALE: len(stale),
            ObservabilityMetricKey.LAST_SEEN_AT: (
                last_seen_at.isoformat() if last_seen_at else None
            ),
            ObservabilityMetricKey.ITEMS: [
                self._serialize_heartbeat(item) for item in records
            ],
        }

    def _safe_call(self, operation: Callable[[], Any]) -> Any:
        try:
            return operation()
        except Exception:
            self.db.rollback()
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.ERROR,
                ObservabilityMetricKey.ERROR: "unavailable",
            }

    def _database_check(self) -> dict[str, Any]:
        self.db.execute(text("select 1")).scalar()
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def _redis_check(self) -> dict[str, Any]:
        settings = get_settings()
        if not settings.task_queue_enabled:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}
        from redis import Redis

        client = Redis.from_url(
            settings.redis_url,
            socket_connect_timeout=1,
            socket_timeout=1,
        )
        try:
            client.ping()
        finally:
            client.close()
        return {ObservabilityKey.STATUS: ObservabilityStatus.OK}

    def _upsert_heartbeat(
        self,
        component: str,
        instance_id: str,
        status_value: str,
        now: Any,
    ) -> SystemHeartbeat:
        dialect_name = self.db.get_bind().dialect.name
        insert_factory = {
            "postgresql": postgresql_insert,
            "sqlite": sqlite_insert,
        }.get(dialect_name)
        if insert_factory is None:
            raise RuntimeError(f"Unsupported heartbeat database dialect: {dialect_name}")
        statement = insert_factory(SystemHeartbeat).values(
            component=component,
            instance_id=instance_id,
            status=status_value,
            last_seen_at=now,
            created_at=now,
            updated_at=now,
        )
        statement = statement.on_conflict_do_update(
            index_elements=["component", "instance_id"],
            set_={
                "status": status_value,
                "last_seen_at": now,
                "updated_at": now,
            },
        )
        self.db.execute(statement)
        return self.db.execute(
            select(SystemHeartbeat).where(
                SystemHeartbeat.component == component,
                SystemHeartbeat.instance_id == instance_id,
            )
        ).scalar_one()

    def _events_check(self) -> dict[str, Any]:
        counts = EventService(self.db).count_by_status()
        failed = counts.get(EventStatus.FAILED, 0)
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if failed else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.PENDING: counts.get(EventStatus.PENDING, 0),
            ObservabilityMetricKey.FAILED: failed,
        }

    def _workflows_check(self) -> dict[str, Any]:
        counts = WorkflowService(self.db).count_by_status()
        failed = counts.get(WorkflowStatus.FAILED, 0)
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if failed else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.RUNNING: counts.get(WorkflowStatus.RUNNING, 0),
            ObservabilityMetricKey.FAILED: failed,
        }

    def _heartbeats_check(self) -> dict[str, Any]:
        summary = self.heartbeat_summary()
        total = summary[ObservabilityMetricKey.TOTAL]
        stale = summary[ObservabilityMetricKey.STALE]
        if total == 0:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if stale else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.TOTAL: total,
            ObservabilityMetricKey.STALE: stale,
            ObservabilityMetricKey.LAST_SEEN_AT: summary[
                ObservabilityMetricKey.LAST_SEEN_AT
            ],
        }

    def _feishu_subscriptions_check(self) -> dict[str, Any]:
        active = int(
            self.db.scalar(
                select(func.count())
                .select_from(PushSubscription)
                .where(PushSubscription.status == PushSubscriptionStatus.ACTIVE)
            )
            or 0
        )
        processable_delivery_filter = or_(
            and_(
                PushDelivery.status.in_(
                    [
                        PushDeliveryStatus.PENDING,
                        PushDeliveryStatus.RETRY,
                    ]
                ),
                PushDelivery.next_attempt_at.is_not(None),
            ),
            and_(
                PushDelivery.status == PushDeliveryStatus.PROCESSING,
                PushDelivery.locked_until.is_not(None),
            ),
        )
        processable_deliveries = int(
            self.db.scalar(
                select(func.count())
                .select_from(PushDelivery)
                .where(processable_delivery_filter)
            )
            or 0
        )
        if not active and not processable_deliveries:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED,
                "active": 0,
                "processable_deliveries": 0,
            }
        settings = get_settings()
        app_id = str(settings.feishu_app_id or "").strip()
        credentials_configured = bool(app_id and settings.feishu_app_secret)
        active_tenant_count = int(
            self.db.scalar(
                select(func.count(func.distinct(FeishuUser.tenant_key)))
                .select_from(PushSubscription)
                .join(FeishuUser, FeishuUser.id == PushSubscription.owner_id)
                .outerjoin(
                    PushDelivery,
                    PushDelivery.subscription_id == PushSubscription.id,
                )
                .where(
                    or_(
                        PushSubscription.status == PushSubscriptionStatus.ACTIVE,
                        processable_delivery_filter,
                    )
                )
            )
            or 0
        )
        ticket_configured = False
        default_tenant_configured = bool(
            str(settings.feishu_default_tenant_key or "").strip()
        )
        reasons: list[str] = []
        if not credentials_configured:
            reasons.append("credentials_missing")
        if settings.feishu_app_type == FeishuAppType.STORE:
            database_ticket = (
                FeishuAppTicketService(self.db).get_ticket(app_id)
                if app_id
                else None
            )
            ticket_configured = bool(
                str(database_ticket or settings.feishu_app_ticket or "").strip()
            )
            if not ticket_configured:
                reasons.append("app_ticket_missing")
            if not default_tenant_configured:
                reasons.append("default_tenant_missing")
        elif active_tenant_count > 1:
            reasons.append("self_app_multiple_tenants")
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.OK
                if not reasons
                else ObservabilityStatus.DEGRADED
            ),
            "active": active,
            "processable_deliveries": processable_deliveries,
            "app_type": settings.feishu_app_type,
            "credentials_configured": credentials_configured,
            "ticket_configured": ticket_configured,
            "default_tenant_configured": default_tenant_configured,
            "active_tenant_count": active_tenant_count,
            "reasons": reasons,
        }

    def _feishu_user_metrics(self) -> dict[str, int]:
        active = int(
            self.db.scalar(
                select(func.count())
                .select_from(FeishuUser)
                .where(FeishuUser.status == FeishuUserStatus.ACTIVE)
            )
            or 0
        )
        return {"active": active}

    def _subscription_metrics(self) -> dict[str, int]:
        active = int(
            self.db.scalar(
                select(func.count())
                .select_from(PushSubscription)
                .where(PushSubscription.status == PushSubscriptionStatus.ACTIVE)
            )
            or 0
        )
        delivery_rows = self.db.execute(
            select(PushDelivery.status, func.count()).group_by(PushDelivery.status)
        ).all()
        deliveries = {str(status_value): int(count) for status_value, count in delivery_rows}
        return {
            "active": active,
            "pending": deliveries.get(PushDeliveryStatus.PENDING, 0)
            + deliveries.get(PushDeliveryStatus.RETRY, 0),
            "failed": deliveries.get(PushDeliveryStatus.FAILED, 0),
        }

    @staticmethod
    def _serialize_heartbeat(record: SystemHeartbeat) -> dict[str, Any]:
        return {
            ObservabilityMetricKey.COMPONENT: record.component,
            ObservabilityMetricKey.INSTANCE_ID: record.instance_id,
            ObservabilityKey.STATUS: record.status,
            ObservabilityMetricKey.LAST_SEEN_AT: record.last_seen_at.isoformat(),
        }

    @staticmethod
    def _heartbeat_stale_threshold() -> Any:
        settings = get_settings()
        return utc_now() - timedelta(seconds=settings.heartbeat_interval_seconds * 3)

    @staticmethod
    def _heartbeat_retention_threshold() -> Any:
        settings = get_settings()
        retention_seconds = max(
            settings.heartbeat_retention_seconds,
            settings.heartbeat_interval_seconds * 3,
        )
        return utc_now() - timedelta(seconds=retention_seconds)