monlet/server/monlet_server/services/ingestion.py

from datetime import UTC, datetime
from uuid import UUID, uuid4

from fastapi import HTTPException
from sqlalchemy import func, select, text
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from .. import metrics
from ..agent_auth import AgentCredential
from ..logging_config import get_logger
from ..models import (
    Agent,
    AgentBlacklist,
    Check,
    Event,
    EventIngestDedup,
    Incident,
    NotificationOutbox,
)
from ..redaction import redact
from ..schemas import CheckResultEvent, HeartbeatRequest
from ..settings import get_settings
from ..timeutil import to_utc

INCIDENT_KEY_MAX = 256

log = get_logger("monlet.ingestion")

_SEVERITY_RANK = {"unknown": 1, "warning": 2, "critical": 3}


def _severity_for_status(status: str) -> str:
    if status in ("warning", "critical", "unknown"):
        return status
    return "unknown"


def _incident_key(agent_id: str, ev: CheckResultEvent) -> str:
    key = ev.incident_key or f"{agent_id}:{ev.check_id}"
    if len(key) > INCIDENT_KEY_MAX:
        raise HTTPException(
            status_code=400,
            detail=f"incident_key exceeds {INCIDENT_KEY_MAX} chars",
        )
    return key


async def _blacklist_row(session: AsyncSession, cred: AgentCredential) -> AgentBlacklist | None:
    res = await session.execute(
        select(AgentBlacklist).where(AgentBlacklist.key_hash == cred.key_hash)
    )
    return res.scalar_one_or_none()


async def reject_if_blocked(session: AsyncSession, cred: AgentCredential) -> None:
    row = await _blacklist_row(session, cred)
    if row is None:
        return
    raise HTTPException(status_code=403, detail="agent blocked")


async def _ensure_pending_capacity(session: AsyncSession) -> None:
    max_pending = get_settings().max_pending_agents
    if max_pending <= 0:
        return
    count = (
        await session.execute(
            select(func.count()).select_from(Agent).where(Agent.pending_key_hash.is_not(None))
        )
    ).scalar_one()
    if count >= max_pending:
        raise HTTPException(status_code=429, detail="too many pending agents")


async def _reject_hostname_collision(session: AsyncSession, agent_id: str, hostname: str) -> None:
    row = (
        await session.execute(
            select(Agent).where(Agent.hostname == hostname, Agent.agent_id != agent_id)
        )
    ).scalar_one_or_none()
    if row is not None:
        raise HTTPException(status_code=409, detail="hostname already belongs to another agent")


def _set_pending(
    agent: Agent,
    hb: HeartbeatRequest,
    cred: AgentCredential,
    observed: datetime,
    features: dict,
) -> None:
    agent.pending_key_hash = cred.key_hash
    agent.pending_key_fingerprint = cred.key_fingerprint
    agent.pending_hostname = hb.hostname
    agent.pending_seen_at = observed
    agent.pending_features = features
    agent.pending_labels = hb.labels


async def upsert_agent_heartbeat(
    session: AsyncSession, hb: HeartbeatRequest, cred: AgentCredential
) -> str:
    observed = to_utc(hb.observed_at)
    await reject_if_blocked(session, cred)

    features = hb.features.model_dump()
    res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash))
    accepted = res.scalar_one_or_none()
    if accepted is not None:
        await _reject_hostname_collision(session, accepted.agent_id, hb.hostname)
        accepted.hostname = hb.hostname
        accepted.status = "alive"
        accepted.last_seen_at = observed
        accepted.features = features
        accepted.labels = hb.labels
        metrics.heartbeats_total.labels(result="ok").inc()
        return "accepted"

    res = await session.execute(select(Agent).where(Agent.pending_key_hash == cred.key_hash))
    pending = res.scalar_one_or_none()
    if pending is not None:
        await _reject_hostname_collision(session, pending.agent_id, hb.hostname)
        _set_pending(pending, hb, cred, observed, features)
        if pending.accepted_key_hash is None:
            pending.hostname = hb.hostname
            pending.status = "alive"
            pending.last_seen_at = observed
            pending.features = features
            pending.labels = hb.labels
        metrics.heartbeats_total.labels(result="pending").inc()
        return "pending"

    hostname_row = (
        await session.execute(select(Agent).where(Agent.hostname == hb.hostname))
    ).scalar_one_or_none()
    if hostname_row is not None:
        if hostname_row.pending_key_hash is not None:
            raise HTTPException(
                status_code=409,
                detail="agent already has a different pending key",
            )
        await _ensure_pending_capacity(session)
        _set_pending(hostname_row, hb, cred, observed, features)
        metrics.heartbeats_total.labels(result="pending").inc()
        return "pending"

    agent_id_row = (
        await session.execute(select(Agent).where(Agent.agent_id == hb.agent_id))
    ).scalar_one_or_none()
    if agent_id_row is not None:
        raise HTTPException(status_code=409, detail="agent_id already exists")

    await _ensure_pending_capacity(session)
    session.add(
        Agent(
            agent_id=hb.agent_id,
            hostname=hb.hostname,
            status="alive",
            last_seen_at=observed,
            features=features,
            labels=hb.labels,
            pending_key_hash=cred.key_hash,
            pending_key_fingerprint=cred.key_fingerprint,
            pending_hostname=hb.hostname,
            pending_seen_at=observed,
            pending_features=features,
            pending_labels=hb.labels,
        )
    )
    metrics.heartbeats_total.labels(result="pending").inc()
    return "pending"


async def accepted_agent_for_key(session: AsyncSession, cred: AgentCredential) -> Agent | None:
    if await _blacklist_row(session, cred) is not None:
        raise HTTPException(status_code=403, detail="agent blocked")
    res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash))
    return res.scalar_one_or_none()


async def _apply_incident(
    session: AsyncSession,
    agent_id: str,
    ev: CheckResultEvent,
    incident_key: str,
) -> tuple[bool, Incident | None]:
    """Return (transition, incident). transition True if open->resolved or none->open."""
    observed = to_utc(ev.observed_at)
    if ev.status == "ok":
        res = await session.execute(
            select(Incident).where(Incident.incident_key == incident_key, Incident.state == "open")
        )
        inc = res.scalar_one_or_none()
        if inc is None:
            return False, None
        inc.state = "resolved"
        inc.resolved_at = observed
        inc.last_event_id = UUID(ev.event_id)
        metrics.incidents_resolved_total.inc()
        return True, inc

    severity = _severity_for_status(ev.status)
    new_id = uuid4()
    summary = (redact(ev.output) or "")[:200] if ev.output else None
    stmt = pg_insert(Incident).values(
        id=new_id,
        incident_key=incident_key,
        agent_id=agent_id,
        check_id=ev.check_id,
        state="open",
        severity=severity,
        opened_at=observed,
        last_event_id=UUID(ev.event_id),
        summary=summary,
    )
    stmt = stmt.on_conflict_do_nothing(
        index_elements=[Incident.incident_key],
        index_where=text("state = 'open'"),
    ).returning(Incident.id)
    ins = await session.execute(stmt)
    inserted_id = ins.scalar_one_or_none()
    if inserted_id is not None:
        await session.flush()
        res = await session.execute(select(Incident).where(Incident.id == inserted_id))
        metrics.incidents_opened_total.inc()
        return True, res.scalar_one()

    res = await session.execute(
        select(Incident).where(Incident.incident_key == incident_key, Incident.state == "open")
    )
    inc = res.scalar_one()
    if inc.agent_id != agent_id or inc.check_id != ev.check_id:
        raise HTTPException(
            status_code=400,
            detail="incident_key already bound to another (agent_id, check_id)",
        )
    cur = _SEVERITY_RANK.get(inc.severity, 0)
    new = _SEVERITY_RANK.get(severity, 0)
    if new > cur:
        inc.severity = severity
    inc.last_event_id = UUID(ev.event_id)
    return False, inc


async def _enqueue_outbox(
    session: AsyncSession,
    incident: Incident,
    event_type: str,
    payload: dict,
) -> None:
    notifiers = get_settings().enabled_notifiers
    if not notifiers:
        return
    now = datetime.now(UTC)
    for name in notifiers:
        session.add(
            NotificationOutbox(
                id=uuid4(),
                incident_id=incident.id,
                notifier=name,
                event_type=event_type,
                state="pending",
                attempts=0,
                next_attempt_at=now,
                payload=payload,
            )
        )


async def ingest_event(
    session: AsyncSession,
    agent_id: str,
    ev: CheckResultEvent,
) -> bool:
    """Return True if a new change row was written, False if not (duplicate/unchanged/late).

    Idempotency is enforced first via the standalone event_ingest_dedup table — a
    duplicate event_id never touches checks/incidents/outbox/events. The full pipeline
    (dedup + events + check + incident + outbox) commits as a single transaction by
    the caller, so a partial failure cannot leave a dedup row without its side-effects.
    """
    observed = to_utc(ev.observed_at)
    output = redact(ev.output)
    check_summary = (output or "")[:200] if output else None
    incident_key = _incident_key(agent_id, ev)
    event_uuid = UUID(ev.event_id)

    # Order matters: late check before dedup insert so that repeated late events stay
    # observable as `late`, not silently collapsed into `deduplicated`.
    chk_res = await session.execute(
        select(Check).where(Check.agent_id == agent_id, Check.check_id == ev.check_id)
    )
    cur = chk_res.scalar_one_or_none()
    is_late = cur is not None and cur.last_observed_at > observed
    if is_late:
        metrics.events_total.labels(result="late").inc()
        return False

    dedup_stmt = (
        pg_insert(EventIngestDedup)
        .values(event_id=event_uuid, received_at=datetime.now(UTC))
        .on_conflict_do_nothing(index_elements=[EventIngestDedup.event_id])
        .returning(EventIngestDedup.event_id)
    )
    if (await session.execute(dedup_stmt)).scalar_one_or_none() is None:
        metrics.events_total.labels(result="deduplicated").inc()
        return False

    changed = cur is None or cur.status != ev.status or cur.exit_code != ev.exit_code

    if changed:
        stmt = pg_insert(Event).values(
            event_id=event_uuid,
            agent_id=agent_id,
            check_id=ev.check_id,
            observed_at=observed,
            status=ev.status,
            exit_code=ev.exit_code,
            duration_ms=ev.duration_ms,
            output=output,
            output_truncated=ev.output_truncated,
            notifications_enabled=ev.notifications_enabled,
            incident_key=incident_key,
        )
        await session.execute(stmt)
    else:
        metrics.events_total.labels(result="unchanged").inc()

    # Upsert Check to be safe under concurrent first-write race for the same (agent_id, check_id).
    # When changed=False keep the existing last_event_id and incident_key — incident_key changes
    # are only meaningful at transitions and would otherwise desync from any open Incident.
    check_values = {
        "agent_id": agent_id,
        "check_id": ev.check_id,
        "status": ev.status,
        "exit_code": ev.exit_code,
        "last_observed_at": observed,
        "last_event_id": UUID(ev.event_id),
        "incident_key": incident_key,
        "summary": check_summary,
    }
    update_set = {
        "status": ev.status,
        "exit_code": ev.exit_code,
        "last_observed_at": observed,
        "summary": check_summary,
    }
    if changed:
        update_set["last_event_id"] = event_uuid
        update_set["incident_key"] = incident_key
    chk_stmt = pg_insert(Check).values(**check_values)
    chk_stmt = chk_stmt.on_conflict_do_update(
        index_elements=[Check.agent_id, Check.check_id], set_=update_set
    )
    await session.execute(chk_stmt)

    if not changed:
        return False

    transition, inc = await _apply_incident(session, agent_id, ev, incident_key)
    # Count accepted only after _apply_incident: an incident_key collision raises
    # 400 and the savepoint rollback discards the Event insert above.
    metrics.events_total.labels(result="accepted").inc()
    if transition and inc is not None and ev.notifications_enabled:
        event_type = "resolved" if ev.status == "ok" else "firing"
        await _enqueue_outbox(
            session,
            inc,
            event_type,
            {
                "incident_id": str(inc.id),
                "agent_id": agent_id,
                "check_id": ev.check_id,
                "status": ev.status,
                "severity": inc.severity,
                "incident_key": incident_key,
                "observed_at": observed.isoformat(),
                "opened_at": inc.opened_at.isoformat() if inc.opened_at else None,
                "resolved_at": inc.resolved_at.isoformat() if inc.resolved_at else None,
                "summary": inc.summary,
                "output": output,
                "event_type": event_type,
            },
        )
    return True