Harden production workflows and agent admission

server/monlet_server/services/detector.py

@@ -202,14 +202,16 @@ async def _apply_liveness(session, agent: Agent, next_status: str, now: datetime
                 summary=summary,
             )
         )
-    else:
+    elif changed:
+        # PH-009: update liveness check row only on status transitions to avoid
+        # O(N) per-tick rewrites of every agent_liveness row. last_observed_at
+        # then reflects the meaningful transition timestamp, not detector ticks.
+        cur.status = liveness_status
+        cur.exit_code = exit_code
+        cur.last_event_id = event_id
+        cur.incident_key = incident_key
         cur.last_observed_at = now
         cur.summary = summary
-        if changed:
-            cur.status = liveness_status
-            cur.exit_code = exit_code
-            cur.last_event_id = event_id
-            cur.incident_key = incident_key
 
     if next_status in ("stale", "dead"):
         await _open_liveness_incident(session, agent, next_status, now, event_id)
@@ -219,6 +221,34 @@ async def _apply_liveness(session, agent: Agent, next_status: str, now: datetime
 
 async def _prune_history(session) -> None:
     s = get_settings()
+    if s.pending_agent_ttl_days > 0:
+        cutoff = datetime.now(UTC) - timedelta(days=s.pending_agent_ttl_days)
+        rows = (
+            (
+                await session.execute(
+                    select(Agent)
+                    .where(
+                        Agent.pending_key_hash.is_not(None),
+                        Agent.pending_seen_at < cutoff,
+                    )
+                    .order_by(Agent.pending_seen_at)
+                    .limit(max(s.pending_agent_prune_batch_size, 1))
+                )
+            )
+            .scalars()
+            .all()
+        )
+        for agent in rows:
+            if agent.accepted_key_hash is None:
+                await session.delete(agent)
+            else:
+                agent.pending_key_hash = None
+                agent.pending_key_fingerprint = None
+                agent.pending_hostname = None
+                agent.pending_seen_at = None
+                agent.pending_features = None
+                agent.pending_labels = None
+
     if s.event_dedup_retention_days > 0:
         cutoff = datetime.now(UTC) - timedelta(days=s.event_dedup_retention_days)
         batch = max(s.event_dedup_prune_batch_size, 1)
@@ -233,17 +263,29 @@ async def _prune_history(session) -> None:
         await session.execute(delete(EventIngestDedup).where(EventIngestDedup.event_id.in_(sub)))
 
     if s.outbox_retention_max_rows > 0:
+        # PH-012: keep the N newest terminal rows. Delete in a bounded batch to
+        # avoid one huge DELETE under failure storms. Active rows ('pending',
+        # 'sending', 'retry') are never eligible — terminal-only filter
+        # protects them.
         keep_outbox = (
             select(NotificationOutbox.id)
             .where(NotificationOutbox.state.in_(OUTBOX_TERMINAL_STATES))
             .order_by(NotificationOutbox.created_at.desc(), NotificationOutbox.id.desc())
             .limit(s.outbox_retention_max_rows)
         )
-        await session.execute(
-            delete(NotificationOutbox)
+        prune_ids = (
+            select(NotificationOutbox.id)
             .where(NotificationOutbox.state.in_(OUTBOX_TERMINAL_STATES))
             .where(NotificationOutbox.id.not_in(keep_outbox))
+            .limit(max(s.outbox_prune_batch_size, 1))
+            .scalar_subquery()
         )
+        await session.execute(
+            delete(NotificationOutbox).where(NotificationOutbox.id.in_(prune_ids))
+        )
+
+
+DETECTOR_BATCH_SIZE = 500
 
 
 async def _tick(sm: async_sessionmaker) -> None:
@@ -252,23 +294,50 @@ async def _tick(sm: async_sessionmaker) -> None:
     stale_cut = now - timedelta(seconds=s.stale_after_sec)
     dead_cut = now - timedelta(seconds=s.dead_after_sec)
     async with sm() as session:
-        agents = (await session.execute(select(Agent))).scalars().all()
-        for agent in agents:
-            if agent.last_seen_at <= dead_cut:
-                next_status = "dead"
-            elif agent.last_seen_at <= stale_cut:
-                next_status = "stale"
-            else:
-                next_status = "alive"
-            if agent.status != next_status:
-                agent.status = next_status
-            await _apply_liveness(session, agent, next_status, now)
+        # PH-009: keyset-paginate agents by agent_id to bound per-tick memory
+        # and per-batch transaction work. Each batch is committed separately so
+        # large inventories do not produce one huge transaction.
+        # Race trade-off: a new agent inserted with agent_id < last_seen_id
+        # mid-tick is skipped this tick and processed next tick. That delay is
+        # bounded by detector_tick_sec (5s by default) and is acceptable —
+        # liveness detection has a stale_after_sec / dead_after_sec horizon
+        # measured in tens of seconds, so a one-tick lag is invisible.
+        last_seen_id: str | None = None
+        while True:
+            q = (
+                select(Agent)
+                .where(Agent.accepted_key_hash.is_not(None))
+                .order_by(Agent.agent_id)
+                .limit(DETECTOR_BATCH_SIZE)
+            )
+            if last_seen_id is not None:
+                q = q.where(Agent.agent_id > last_seen_id)
+            agents = (await session.execute(q)).scalars().all()
+            if not agents:
+                break
+            for agent in agents:
+                if agent.last_seen_at <= dead_cut:
+                    next_status = "dead"
+                elif agent.last_seen_at <= stale_cut:
+                    next_status = "stale"
+                else:
+                    next_status = "alive"
+                if agent.status != next_status:
+                    agent.status = next_status
+                await _apply_liveness(session, agent, next_status, now)
+            last_seen_id = agents[-1].agent_id
+            await session.commit()
+            if len(agents) < DETECTOR_BATCH_SIZE:
+                break
+
         await _prune_history(session)
         await session.commit()
 
         for status in ("alive", "stale", "dead"):
             r = await session.execute(
-                select(func.count()).select_from(Agent).where(Agent.status == status)
+                select(func.count())
+                .select_from(Agent)
+                .where(Agent.accepted_key_hash.is_not(None), Agent.status == status)
             )
             metrics.agents_gauge.labels(status=status).set(r.scalar_one())
         r = await session.execute(

server/monlet_server/services/ingestion.py

@@ -2,13 +2,22 @@ from datetime import UTC, datetime
 from uuid import UUID, uuid4
 
 from fastapi import HTTPException
-from sqlalchemy import select, text
+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, Check, Event, EventIngestDedup, Incident, NotificationOutbox
+from ..models import (
+    Agent,
+    AgentBlacklist,
+    Check,
+    Event,
+    EventIngestDedup,
+    Incident,
+    NotificationOutbox,
+)
 from ..redaction import redact
 from ..schemas import CheckResultEvent, HeartbeatRequest
 from ..settings import get_settings
@@ -37,28 +46,143 @@ def _incident_key(agent_id: str, ev: CheckResultEvent) -> str:
     return key
 
 
-async def upsert_agent_heartbeat(session: AsyncSession, hb: HeartbeatRequest) -> None:
+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,
+    agent_id: str,
+    hostname: str,
+    observed: datetime | None = None,
+) -> 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)
-    stmt = pg_insert(Agent).values(
-        agent_id=hb.agent_id,
-        hostname=hb.hostname,
-        status="alive",
-        last_seen_at=observed,
-        features=hb.features.model_dump(),
-        labels=hb.labels,
+    await reject_if_blocked(session, cred, hb.agent_id, hb.hostname, observed)
+
+    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,
+        )
     )
-    stmt = stmt.on_conflict_do_update(
-        index_elements=[Agent.agent_id],
-        set_={
-            "hostname": stmt.excluded.hostname,
-            "status": "alive",
-            "last_seen_at": stmt.excluded.last_seen_at,
-            "features": stmt.excluded.features,
-            "labels": stmt.excluded.labels,
-        },
-    )
-    await session.execute(stmt)
-    metrics.heartbeats_total.labels(result="ok").inc()
+    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(

server/monlet_server/services/notifier_worker.py

@@ -25,6 +25,43 @@ def _backoff(attempts: int) -> timedelta:
     return timedelta(seconds=_BACKOFF_SCHEDULE_SEC[idx])
 
 
+def _apply_output_policy(payload: dict) -> dict:
+    """PH-020: enforce a single notifier output policy.
+
+    Per-notifier code keeps its own format-specific shortening (e.g. Telegram
+    300 chars), but a hard cap and include-flag apply consistently before any
+    delivery. Truncation uses a clear marker so receivers can see it happened.
+    """
+    settings = get_settings()
+    out = dict(payload)
+    text_value = out.get("output")
+    if text_value is None:
+        return out
+    if not settings.notifier_include_output:
+        out["output"] = None
+        out["output_truncated"] = True
+        return out
+    max_bytes = max(settings.notifier_max_output_bytes, 0)
+    encoded = text_value.encode("utf-8")
+    if len(encoded) <= max_bytes:
+        return out
+    # Reserve room for the marker so the final string never exceeds max_bytes.
+    # If max_bytes is smaller than even the marker, drop the body entirely —
+    # this keeps the absolute cap a hard guarantee under any setting.
+    dropped = len(encoded) - max_bytes
+    marker = f"...[truncated {dropped} bytes]"
+    marker_bytes = len(marker.encode("utf-8"))
+    if marker_bytes >= max_bytes:
+        out["output"] = ""
+        out["output_truncated"] = True
+        return out
+    prefix_budget = max_bytes - marker_bytes
+    truncated = encoded[:prefix_budget].decode("utf-8", errors="ignore")
+    out["output"] = f"{truncated}{marker}"
+    out["output_truncated"] = True
+    return out
+
+
 async def _claim_batch(session, limit: int, lease_sec: int) -> list[dict]:
     res = await session.execute(
         text(
@@ -107,8 +144,10 @@ async def _process_row(
             metrics.notifications_total.labels(notifier=row["notifier"], result="discarded").inc()
             return
 
+        # PH-020: apply central output policy before handing to any notifier.
+        delivery_payload = _apply_output_policy(row["payload"])
         try:
-            result = await notifier.deliver(row["event_type"], row["payload"])
+            result = await notifier.deliver(row["event_type"], delivery_payload)
         except Exception as exc:  # defensive: notifier should not raise
             result_ok = False
             result_perm = False