@opengeni/api-router 0.5.2 → 0.5.4

package/src/sandbox/metrics-ingestion.ts

@@ -21,6 +21,9 @@
 //     re-surfaced. Consuming the Hello's `capabilities.desktop` / `display` makes
 //     `has_display` track reality (both directions), which the desktop-capability
 //     gate (packages/runtime capabilities.ts) keys off.
+//     refreshEnrollmentOpStream — reconcile `enrollments.op_stream` to the LIVE
+//     runner capability the Hello reports, leaving legacy request/reply exec as the
+//     fallback unless the runner advertises the streaming engine.
 //
 // Both consumers are BEST-EFFORT and fail-soft: a decode/DB error for one message
 // is logged + swallowed (the bus subscription already swallows handler throws) so
@@ -28,16 +31,27 @@
 // back-pressures the agent, or breaks its connect.
 
 import {
+  clearEnrollmentWentOffline,
   getEnrollment,
   ingestMachineMetricsSample,
+  sessionsWithActiveOpOnEnrollment,
   setEnrollmentDisplayState,
+  setEnrollmentOpStreamState,
+  setEnrollmentWentOffline,
   touchEnrollmentLastSeen,
+  type AppendEventInput,
   type Database,
   type MachineMetricsSample,
 } from "@opengeni/db";
-import type { EventBus } from "@opengeni/events";
+import { appendAndPublishEvents, type EventBus } from "@opengeni/events";
 import type { Observability } from "@opengeni/observability";
-import { AgentEvent, Hello, type MetricsSample } from "@opengeni/agent-proto";
+import {
+  AgentEvent,
+  GoingOfflineReason,
+  Hello,
+  goingOfflineReasonToJSON,
+  type MetricsSample,
+} from "@opengeni/agent-proto";
 
 /** The wildcard subject the agent event plane publishes heartbeats on. */
 export const AGENT_EVENTS_SUBJECT = "agent.*.*.events";
@@ -50,7 +64,10 @@ export const AGENT_HELLO_SUBJECT = "agent.*.*.hello";
  * expected tail token. Returns null for a subject that does not match the shape
  * (defensive — the subscription pattern already constrains it).
  */
-function parseAgentSubject(subject: string, tail: "events" | "hello"): { workspaceId: string; agentId: string } | null {
+function parseAgentSubject(
+  subject: string,
+  tail: "events" | "hello",
+): { workspaceId: string; agentId: string } | null {
   const parts = subject.split(".");
   if (parts.length !== 4 || parts[0] !== "agent" || parts[3] !== tail) {
     return null;
@@ -59,12 +76,16 @@ function parseAgentSubject(subject: string, tail: "events" | "hello"): { workspa
 }
 
 /** Parse `agent.<ws>.<id>.events` → `{ workspaceId, agentId }` (heartbeat plane). */
-export function parseAgentEventSubject(subject: string): { workspaceId: string; agentId: string } | null {
+export function parseAgentEventSubject(
+  subject: string,
+): { workspaceId: string; agentId: string } | null {
   return parseAgentSubject(subject, "events");
 }
 
 /** Parse `agent.<ws>.<id>.hello` → `{ workspaceId, agentId }` (connect plane). */
-export function parseAgentHelloSubject(subject: string): { workspaceId: string; agentId: string } | null {
+export function parseAgentHelloSubject(
+  subject: string,
+): { workspaceId: string; agentId: string } | null {
   return parseAgentSubject(subject, "hello");
 }
 
@@ -96,7 +117,10 @@ export function wireSampleToDbSample(wire: MetricsSample): MachineMetricsSample
     contention: wire.runQueue,
     // The sample carries its own wall-clock stamp (epoch ms); fall back to now on
     // a missing/zero stamp so a series row is never NULL-dated.
-    sampledAt: wire.sampledAtMs && Number(wire.sampledAtMs) > 0 ? new Date(Number(wire.sampledAtMs)) : new Date(),
+    sampledAt:
+      wire.sampledAtMs && Number(wire.sampledAtMs) > 0
+        ? new Date(Number(wire.sampledAtMs))
+        : new Date(),
   };
 }
 
@@ -132,16 +156,61 @@ export async function ingestHeartbeat(
   return { ingested: true, seriesAppended: result.seriesAppended };
 }
 
+/**
+ * Fan out one or more machine-LINK session events to the sessions that had an
+ * active op running on the machine when its control link changed (per
+ * `sessionsWithActiveOpOnEnrollment`) — the announce-only failure-visibility
+ * plane. Each session's events are stamped on its OWN active turn. No matching
+ * session ⇒ nothing is emitted (an idle-machine blip must never spam idle /
+ * historical sessions). Called best-effort inside the handlers' fail-soft blocks.
+ *
+ * Each session's emission is ISOLATED: one session's append failing (a
+ * session-specific constraint like a sequence collision from a racing writer, a
+ * transient write error) is logged with that sessionId and skipped, never
+ * aborting the fan-out — one session's failure must never cost the OTHER matching
+ * sessions their events. A partial fan-out stays visible per-session in the logs.
+ */
+async function fanOutMachineLinkEvents(
+  db: Database,
+  bus: EventBus,
+  observability: Observability | undefined,
+  workspaceId: string,
+  enrollmentId: string,
+  build: (activeTurnId: string) => AppendEventInput[],
+): Promise<void> {
+  const sessions = await sessionsWithActiveOpOnEnrollment(db, { workspaceId, enrollmentId });
+  for (const session of sessions) {
+    try {
+      await appendAndPublishEvents(
+        db,
+        bus,
+        workspaceId,
+        session.sessionId,
+        build(session.activeTurnId),
+      );
+    } catch (error) {
+      observability?.warn?.("Failed to fan out a machine-link event to a session", {
+        workspaceId,
+        sessionId: session.sessionId,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+  }
+}
+
 /**
  * Decode a raw `AgentEvent` payload + ingest it (the per-message handler). A
- * heartbeat carrying a metrics sample is ingested; a going-offline (or a
- * heartbeat without metrics) is a no-op. Decode failures are reported + swallowed.
+ * heartbeat carrying a metrics sample is ingested; a going-offline records the
+ * machine-plane marker + fans out the link-plane session events. Decode failures
+ * are reported + swallowed. `bus` (when present) enables the session-event
+ * fan-out; the live consumer always supplies it, pure unit tests may omit it.
  */
 export async function handleAgentEventPayload(
   db: Database,
   observability: Observability | undefined,
   payload: Uint8Array,
   subject: string,
+  bus?: EventBus,
 ): Promise<void> {
   const ids = parseAgentEventSubject(subject);
   if (!ids) {
@@ -157,15 +226,84 @@ export async function handleAgentEventPayload(
     });
     return;
   }
+  // A clean GoingOffline is the machine-plane's typed shutdown signal. Two things
+  // happen, in this order:
+  //   1. Record it ALWAYS on the machine plane (a Prometheus counter keyed by the
+  //      typed reason) so a fleet operator can see clean stops / self-updates /
+  //      host shutdowns. This fires unconditionally, independent of the DB.
+  //   2. Stamp the enrollment's clean going-offline marker so the liveness
+  //      derivation reads the machine OFFLINE immediately instead of waiting out
+  //      the last_seen dead-detect window. Best-effort + fail-soft (like the rest
+  //      of this module): an unknown enrollment is a no-op and a DB error is
+  //      swallowed so a bad write never tears down the consumer. Deliberately does
+  //      NOT touch last-seen (a shutdown must not look "more recently alive").
+  if (event.event?.$case === "goingOffline") {
+    const reason = goingOfflineReasonToJSON(event.event.goingOffline.reason);
+    observability?.incrementCounter({
+      name: "opengeni_machine_going_offline_total",
+      help: "Total Connected Machine clean GoingOffline signals by typed reason.",
+      labels: { reason },
+    });
+    try {
+      const enrollment = await getEnrollment(db, ids.workspaceId, ids.agentId);
+      if (enrollment) {
+        await setEnrollmentWentOffline(db, {
+          accountId: enrollment.accountId,
+          workspaceId: ids.workspaceId,
+          enrollmentId: ids.agentId,
+          reason,
+        });
+        // Fan out the link-plane events to the sessions with an active op on this
+        // machine: machine.link.lost (its control link is going away) for every
+        // clean going-offline, PLUS machine.runner.restarted when the reason is a
+        // self-update restart specifically (link.lost fires for it too; this
+        // distinguishes a restart from a plain stop / host shutdown).
+        if (bus) {
+          const isSelfUpdate =
+            event.event.goingOffline.reason === GoingOfflineReason.GOING_OFFLINE_REASON_UPDATE;
+          await fanOutMachineLinkEvents(
+            db,
+            bus,
+            observability,
+            ids.workspaceId,
+            ids.agentId,
+            (activeTurnId) => {
+              const events: AppendEventInput[] = [
+                { type: "machine.link.lost", turnId: activeTurnId, payload: { reason } },
+              ];
+              if (isSelfUpdate) {
+                events.push({
+                  type: "machine.runner.restarted",
+                  turnId: activeTurnId,
+                  payload: {},
+                });
+              }
+              return events;
+            },
+          );
+        }
+      }
+    } catch (error) {
+      observability?.warn?.("Failed to record a machine clean going-offline", {
+        subject,
+        error: error instanceof Error ? error.message : String(error),
+      });
+    }
+    return;
+  }
   if (event.event?.$case !== "heartbeat") {
-    return; // going-offline / unknown → not a metrics point.
+    return; // an unknown event kind → not a metrics point.
   }
   const metrics = event.event.heartbeat.metrics;
   if (!metrics) {
     return; // a heartbeat without a sample → liveness already touched elsewhere.
   }
   try {
-    await ingestHeartbeat(db, { workspaceId: ids.workspaceId, agentId: ids.agentId, sample: metrics });
+    await ingestHeartbeat(db, {
+      workspaceId: ids.workspaceId,
+      agentId: ids.agentId,
+      sample: metrics,
+    });
   } catch (error) {
     observability?.warn?.("Failed to ingest a machine metrics heartbeat", {
       subject,
@@ -185,7 +323,7 @@ export function startMetricsIngestion(deps: {
   observability?: Observability;
 }): () => void {
   return deps.bus.subscribeAgentEvents(AGENT_EVENTS_SUBJECT, (payload, subject) =>
-    handleAgentEventPayload(deps.db, deps.observability, payload, subject),
+    handleAgentEventPayload(deps.db, deps.observability, payload, subject, deps.bus),
   );
 }
 
@@ -229,6 +367,11 @@ export function helloDesktopUnavailableReason(hello: Hello): string | null {
   return reason ? reason : null;
 }
 
+/** Whether the runner's current Hello advertises the op-stream engine. */
+export function helloReportsOpStream(hello: Hello): boolean {
+  return hello.capabilities?.opStream === true;
+}
+
 /**
  * Reconcile `enrollments.has_display` (+ the capture-blocked reason) to what a Hello
  * reports. Resolves the enrollment (the accountId is the RLS principal + the
@@ -239,7 +382,12 @@ export function helloDesktopUnavailableReason(hello: Hello): string | null {
  */
 export async function refreshEnrollmentDisplay(
   db: Database,
-  input: { workspaceId: string; agentId: string; hasDisplay: boolean; desktopUnavailableReason?: string | null },
+  input: {
+    workspaceId: string;
+    agentId: string;
+    hasDisplay: boolean;
+    desktopUnavailableReason?: string | null;
+  },
 ): Promise<{ updated: boolean }> {
   const desktopUnavailableReason = input.desktopUnavailableReason ?? null;
   const enrollment = await getEnrollment(db, input.workspaceId, input.agentId);
@@ -264,15 +412,50 @@ export async function refreshEnrollmentDisplay(
 }
 
 /**
- * Decode a raw `Hello` payload + refresh the enrollment's display cursor (the
- * per-message handler for the hello plane). Decode failures + write failures are
- * reported + swallowed — a display refresh must NEVER break the agent's connect.
+ * Reconcile `enrollments.op_stream` to what a Hello reports. Resolves the
+ * enrollment first so the accountId remains the RLS principal and so a no-change
+ * Hello short-circuits BEFORE issuing any write (the DB writer is itself
+ * change-guarded as a backstop). An unknown/cross-workspace agentId is a no-op.
+ */
+export async function refreshEnrollmentOpStream(
+  db: Database,
+  input: {
+    workspaceId: string;
+    agentId: string;
+    opStream: boolean;
+  },
+): Promise<{ updated: boolean }> {
+  const enrollment = await getEnrollment(db, input.workspaceId, input.agentId);
+  if (!enrollment) {
+    return { updated: false };
+  }
+  if (enrollment.opStream === input.opStream) {
+    // The capability is unchanged — do not even issue the UPDATE (no churn on a
+    // steady-state Hello).
+    return { updated: false };
+  }
+  return await setEnrollmentOpStreamState(db, {
+    accountId: enrollment.accountId,
+    workspaceId: input.workspaceId,
+    enrollmentId: input.agentId,
+    opStream: input.opStream,
+  });
+}
+
+/**
+ * Decode a raw `Hello` payload + refresh the enrollment's display cursor + clear
+ * any pending clean going-offline marker and, when the reconnect actually cleared
+ * one, fan out machine.link.restored to the sessions with an active op on the
+ * machine (the per-message handler for the hello plane). Decode failures + write
+ * failures are reported + swallowed — a Hello must NEVER break the agent's connect.
+ * `bus` (when present) enables the link.restored fan-out.
  */
 export async function handleHelloPayload(
   db: Database,
   observability: Observability | undefined,
   payload: Uint8Array,
   subject: string,
+  bus?: EventBus,
 ): Promise<void> {
   const ids = parseAgentHelloSubject(subject);
   if (!ids) {
@@ -295,8 +478,46 @@ export async function handleHelloPayload(
       hasDisplay: helloReportsDisplay(hello),
       desktopUnavailableReason: helloDesktopUnavailableReason(hello),
     });
+    await refreshEnrollmentOpStream(db, {
+      workspaceId: ids.workspaceId,
+      agentId: ids.agentId,
+      opStream: helloReportsOpStream(hello),
+    });
+  } catch (error) {
+    observability?.warn?.("Failed to refresh an enrollment's capabilities from a Hello", {
+      subject,
+      error: error instanceof Error ? error.message : String(error),
+    });
+  }
+  // A reconnect Hello re-announces the machine, so any pending clean going-offline
+  // marker no longer holds — clear it so the liveness derivation stops reading the
+  // machine offline. Best-effort + fail-soft, and change-guarded in the DB (a
+  // steady-state Hello with no marker writes nothing), so this never breaks the
+  // agent's connect and never churns. When a marker was ACTUALLY cleared (the
+  // machine had been reported link.lost), fan out machine.link.restored to the
+  // sessions with an active op on it — a restored only ever pairs a prior lost, so
+  // a routine connect Hello (no marker) emits nothing.
+  try {
+    const enrollment = await getEnrollment(db, ids.workspaceId, ids.agentId);
+    if (enrollment) {
+      const { cleared } = await clearEnrollmentWentOffline(db, {
+        accountId: enrollment.accountId,
+        workspaceId: ids.workspaceId,
+        enrollmentId: ids.agentId,
+      });
+      if (cleared && bus) {
+        await fanOutMachineLinkEvents(
+          db,
+          bus,
+          observability,
+          ids.workspaceId,
+          ids.agentId,
+          (activeTurnId) => [{ type: "machine.link.restored", turnId: activeTurnId, payload: {} }],
+        );
+      }
+    }
   } catch (error) {
-    observability?.warn?.("Failed to refresh an enrollment's display from a Hello", {
+    observability?.warn?.("Failed to clear a machine going-offline marker on a Hello", {
       subject,
       error: error instanceof Error ? error.message : String(error),
     });
@@ -315,6 +536,6 @@ export function startHelloIngestion(deps: {
   observability?: Observability;
 }): () => void {
   return deps.bus.subscribeAgentEvents(AGENT_HELLO_SUBJECT, (payload, subject) =>
-    handleHelloPayload(deps.db, deps.observability, payload, subject),
+    handleHelloPayload(deps.db, deps.observability, payload, subject, deps.bus),
   );
 }