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.database.migrations import (
    current_alembic_revisions,
    expected_alembic_heads,
)
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.models import DomainEvent
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.services import FeishuInboundService
from app.modules.feishu_users.constants import (
    FeishuUserStatus,
    parse_admin_identities,
)
from app.modules.feishu_users.models import FeishuUser
from app.modules.observability.constants import (
    HeartbeatComponent,
    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.SCHEMA: self._safe_call(self._schema_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),
            ObservabilityKey.API: self._safe_call(
                lambda: self._component_heartbeat_check(
                    HeartbeatComponent.API,
                    required=self._api_required(),
                )
            ),
            ObservabilityKey.SCHEDULER: self._safe_call(
                lambda: self._component_heartbeat_check(
                    HeartbeatComponent.SCHEDULER,
                    required=self._scheduler_required(),
                )
            ),
            ObservabilityKey.WORKER: self._safe_call(
                lambda: self._component_heartbeat_check(
                    HeartbeatComponent.WORKER,
                    required=get_settings().task_queue_enabled,
                )
            ),
            ObservabilityKey.FEISHU_EVENTS: self._safe_call(
                lambda: self._component_heartbeat_check(
                    HeartbeatComponent.FEISHU_EVENTS,
                    required=self._feishu_events_required(),
                )
            ),
            "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),
                "feishu_inbound": self._safe_call(self._feishu_inbound_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 _schema_check(self) -> dict[str, Any]:
        settings = get_settings()
        required = (
            settings.app_env.lower() in {"prod", "production"}
            or settings.scheduler_enabled
            or settings.feishu_user_features_enabled
            or settings.feishu_event_transport != "disabled"
        )
        if not required:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}

        expected = expected_alembic_heads()
        current = current_alembic_revisions(self.db.connection())
        matches = current == expected and len(expected) == 1
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.OK if matches else ObservabilityStatus.DEGRADED
            ),
            "current": [] if current is None else list(current),
            "expected": list(expected),
            "reason": None if matches else "schema_revision_mismatch",
        }

    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()
        current = utc_now()
        failed = counts.get(EventStatus.FAILED, 0)
        processable = int(
            self.db.scalar(
                select(func.count())
                .select_from(DomainEvent)
                .where(
                    DomainEvent.status == EventStatus.PENDING,
                    DomainEvent.next_attempt_at.is_not(None),
                    DomainEvent.next_attempt_at <= current,
                    or_(
                        DomainEvent.locked_until.is_(None),
                        DomainEvent.locked_until <= current,
                    ),
                )
            )
            or 0
        )
        expired_leases = int(
            self.db.scalar(
                select(func.count())
                .select_from(DomainEvent)
                .where(
                    DomainEvent.status == EventStatus.PENDING,
                    DomainEvent.locked_by.is_not(None),
                    DomainEvent.locked_until.is_not(None),
                    DomainEvent.locked_until <= current,
                )
            )
            or 0
        )
        reasons: list[str] = []
        if processable and not get_settings().event_dispatch_enabled:
            reasons.append("dispatch_disabled")
        if expired_leases:
            reasons.append("expired_leases")
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.DEGRADED if reasons else ObservabilityStatus.OK
            ),
            ObservabilityMetricKey.PENDING: counts.get(EventStatus.PENDING, 0),
            ObservabilityMetricKey.FAILED: failed,
            "processable": processable,
            "expired_leases": expired_leases,
            "reasons": reasons,
        }

    def _workflows_check(self) -> dict[str, Any]:
        counts = WorkflowService(self.db).count_by_status()
        failed = counts.get(WorkflowStatus.FAILED, 0)
        return {
            ObservabilityKey.STATUS: 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]
        active = summary[ObservabilityMetricKey.ACTIVE]
        stale = summary[ObservabilityMetricKey.STALE]
        if total == 0:
            return {ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED}
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.OK if active else ObservabilityStatus.DEGRADED
            ),
            ObservabilityMetricKey.TOTAL: total,
            ObservabilityMetricKey.ACTIVE: active,
            ObservabilityMetricKey.STALE: stale,
            ObservabilityMetricKey.LAST_SEEN_AT: summary[
                ObservabilityMetricKey.LAST_SEEN_AT
            ],
        }

    def _component_heartbeat_check(
        self,
        component: str,
        *,
        required: bool,
    ) -> dict[str, Any]:
        if not required:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED,
                ObservabilityMetricKey.COMPONENT: component,
            }
        records = list(
            self.db.execute(
                select(SystemHeartbeat.status, SystemHeartbeat.last_seen_at).where(
                    SystemHeartbeat.component == component
                )
            ).all()
        )
        if not records:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.DEGRADED,
                ObservabilityMetricKey.COMPONENT: component,
                "reason": "heartbeat_missing",
            }
        last_seen_at = max(item.last_seen_at for item in records)
        fresh_records = [
            item
            for item in records
            if item.last_seen_at >= self._heartbeat_stale_threshold()
        ]
        if not fresh_records:
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.DEGRADED,
                ObservabilityMetricKey.COMPONENT: component,
                ObservabilityMetricKey.LAST_SEEN_AT: last_seen_at.isoformat(),
                "reason": "heartbeat_stale",
            }
        healthy = any(item.status == HeartbeatStatus.OK for item in fresh_records)
        return {
            ObservabilityKey.STATUS: (
                ObservabilityStatus.OK if healthy else ObservabilityStatus.DEGRADED
            ),
            ObservabilityMetricKey.COMPONENT: component,
            ObservabilityMetricKey.LAST_SEEN_AT: last_seen_at.isoformat(),
            "reason": None if healthy else "heartbeat_degraded",
        }

    @staticmethod
    def _api_required() -> bool:
        return get_settings().app_env.lower() in {"prod", "production"}

    @staticmethod
    def _scheduler_required() -> bool:
        settings = get_settings()
        return (
            settings.feishu_user_features_enabled
            or settings.app_env.lower() in {"prod", "production"}
        )

    @staticmethod
    def _feishu_events_required() -> bool:
        settings = get_settings()
        return settings.feishu_event_transport == "long_connection"

    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
        )
        settings = get_settings()
        if (
            not settings.feishu_user_features_enabled
            and not active
            and not processable_deliveries
        ):
            return {
                ObservabilityKey.STATUS: ObservabilityStatus.SKIPPED,
                "active": 0,
                "processable_deliveries": 0,
            }
        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 settings.feishu_user_features_enabled
            and (active or processable_deliveries)
        ):
            reasons.append("features_disabled")
        if (
            settings.app_env.lower() in {"prod", "production"}
            and settings.feishu_user_features_enabled
        ):
            try:
                configured_admins = parse_admin_identities(
                    settings.feishu_admin_identities
                )
            except ValueError:
                reasons.append("admin_identities_invalid")
            else:
                if not configured_admins:
                    reasons.append("admin_identities_missing")
        if settings.feishu_event_transport == "disabled":
            reasons.append("event_transport_disabled")
        elif (
            settings.feishu_event_transport == "webhook"
            and not settings.feishu_verification_token
        ):
            reasons.append("verification_token_missing")
        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,
            "features_enabled": settings.feishu_user_features_enabled,
            "event_transport": settings.feishu_event_transport,
            "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 _feishu_inbound_metrics(self) -> dict[str, int]:
        return FeishuInboundService(self.db).status_counts()

    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)