@relaycast/engine 10.2.0 → 10.3.1

package/dist/engine/node.js

@@ -12,7 +12,7 @@ import { generateId } from './snowflake.js';
 import { assertRegistrableAgentName, agentLivenessLoss, agentLivenessRefresh, RELEASED_AGENT_STATUS } from './agent.js';
 import { rotateAgentIdentity } from './agentIdentity.js';
 import { isNodeLive, isReusableForMachineMatch, nodeHasCapacity, NODE_LIVENESS_TTL_MS } from './placement.js';
-import { DEFAULT_PROVIDER_NAME, assertActionTaskMode, capabilityKind, heartbeatProvider, isProviderLive, markProviderOffline, materializeProviderActions, providerActionsMaterialized, recomputeNodeAggregate, removeProvider, upsertProvider, } from './nodeProvider.js';
+import { DEFAULT_PROVIDER_NAME, assertActionTaskMode, capabilityKind, heartbeatProvider, isProviderLive, markProviderOffline, materializeProviderActions, providerActionsMaterialized, recomputeNodeAggregate, refreshNodeAggregate, removeProvider, upsertProvider, } from './nodeProvider.js';
 import { serializeNodeOp } from './nodeLock.js';
 import { completeNodeInvocation, dispatchCapacitySpawn, rescheduleInvocationsForLostNode, rescheduleNodeInvocation, } from './action.js';
 import { emitInvocationCompletionEffects, emitAgentExitedEffects, emitNodeStatusEffects, fleetInvocationId } from './invocationCompletion.js';
@@ -704,7 +704,11 @@ export async function registerNode(db, workspaceId, authenticatedNodeId, message
     }));
     return { node: publicNode(updated), acceptance, provider };
 }
-export async function heartbeatNode(db, workspaceId, nodeId, providerName, message) {
+export async function heartbeatNode(db, workspaceId, nodeId, providerName, message, 
+// A caller that already read the node row may pass it to skip the role
+// probe: `null` means "the row was missing", an object supplies its role.
+// `undefined` keeps the historical own-read behavior.
+priorNode) {
     // A heartbeat is provider-scoped by connection. It may carry the provider's
     // roster snapshot so the engine can refresh the descriptor between (or in the
     // absence of a fresh) node.register — e.g. after an engine restart where the
@@ -713,10 +717,12 @@ export async function heartbeatNode(db, workspaceId, nodeId, providerName, messa
     if (message.node_id !== undefined && message.node_id !== nodeId) {
         throw codedError('node_id does not match the authenticated node token', 'node_id_mismatch', 403);
     }
-    const [nodeState] = await db
-        .select({ role: nodes.role })
-        .from(nodes)
-        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
+    const nodeState = priorNode === undefined
+        ? (await db
+            .select({ role: nodes.role })
+            .from(nodes)
+            .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId))))[0]
+        : priorNode;
     if (!nodeState)
         return null;
     if (nodeState.role === 'direct') {
@@ -818,22 +824,55 @@ export async function heartbeatNode(db, workspaceId, nodeId, providerName, messa
         // ordering the separate stamp used to guarantee: the aggregate derives
         // lastHeartbeatAt from provider rows, which registration also writes,
         // whereas this column is only ever set by an arriving frame.
-        updated = await recomputeNodeAggregate(tx, workspaceId, nodeId, {
-            ...(message.version !== undefined ? { version: message.version } : {}),
-            provenLiveAt: now,
-        });
+        // A capability-free beat cannot change the manifest union, so it takes the
+        // single-statement aggregate path; only a beat carrying `capabilities`
+        // pays the provider-listing recompute.
+        updated = message.capabilities === undefined
+            ? await refreshNodeAggregate(tx, workspaceId, nodeId, {
+                ...(message.version !== undefined ? { version: message.version } : {}),
+                provenLiveAt: now,
+            })
+            : await recomputeNodeAggregate(tx, workspaceId, nodeId, {
+                ...(message.version !== undefined ? { version: message.version } : {}),
+                provenLiveAt: now,
+            });
     });
     return updated ? publicNode(updated) : null;
 }
-export async function markNodeOffline(db, registry, workspaceId, nodeId, effects) {
-    // Capture the node's name + prior status BEFORE flipping it so a durable
-    // node.status.offline is emitted only on a real online -> offline transition
-    // (never on a node that was already offline).
-    const [before] = await db
-        .select({ status: nodes.status, name: nodes.name })
-        .from(nodes)
-        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
-    await db
+/**
+ * Flip a node's row offline in ONE conditional UPDATE — the only part of an
+ * offline transition a caller must serialize against reconnects. `expected`
+ * carries the heartbeat the caller's stale-read observed: any heartbeat or
+ * registration that landed since rewrote `last_heartbeat_at`, the CAS misses,
+ * and the caller must abandon the transition rather than clobber the live
+ * node. Returns the claimed row's name, timestamp, and monotonic revision; the
+ * revision is the bound the settle tail scopes its writes to. Returns undefined
+ * when the claim was lost.
+ */
+export async function claimNodeOffline(db, workspaceId, nodeId, expected) {
+    const claimAt = new Date();
+    const nextRevision = sql `MAX(
+    ${nodes.livenessRevision},
+    COALESCE((
+      SELECT MAX(np.liveness_revision)
+      FROM node_providers np
+      WHERE np.workspace_id = ${workspaceId} AND np.node_id = ${nodeId}
+    ), 0),
+    COALESCE((
+      SELECT MAX(a.node_liveness_revision)
+      FROM agents a
+      WHERE a.workspace_id = ${workspaceId} AND a.location_node_id = ${nodeId}
+        AND a.location_type = 'via_node'
+        AND a.status NOT IN ('offline', 'released')
+    ), 0),
+    COALESCE((
+      SELECT MAX(ai.node_liveness_revision)
+      FROM action_invocations ai
+      WHERE ai.workspace_id = ${workspaceId} AND ai.dispatched_node_id = ${nodeId}
+        AND ai.status IN ('pending', 'dispatched', 'invoked')
+    ), 0)
+  ) + 1`;
+    const [claimed] = await db
         .update(nodes)
         .set({
         status: 'offline',
@@ -841,29 +880,86 @@ export async function markNodeOffline(db, registry, workspaceId, nodeId, effects
         load: 0,
         loadReported: false,
         activeAgents: 0,
-        lastHeartbeatAt: new Date(),
+        // The written heartbeat IS the claim epoch: settle tails and later
+        // transitions detect a post-claim liveness write by observing a
+        // `last_heartbeat_at` newer than `claimAt`, so both values must be the
+        // same instant (second-granular column — two `new Date()` calls could
+        // straddle a boundary).
+        lastHeartbeatAt: claimAt,
+        // This is the durable claim boundary. Triggers stamp every later child
+        // liveness mutation above it, including writes in the same wall-clock
+        // second, so the settle tail never relies on timestamp ordering.
+        livenessRevision: nextRevision,
     })
-        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
+        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId), eq(nodes.status, 'online'), eq(nodes.livenessRevision, expected.livenessRevision), expected.lastHeartbeatAt === null
+        ? isNull(nodes.lastHeartbeatAt)
+        : eq(nodes.lastHeartbeatAt, expected.lastHeartbeatAt)))
+        .returning({ name: nodes.name, claimRevision: nodes.livenessRevision });
+    return claimed ? { name: claimed.name, claimAt, claimRevision: claimed.claimRevision } : undefined;
+}
+/**
+ * The resumable tail of an offline transition: provider/agent liveness loss,
+ * invocation rescheduling, and the durable event. Safe to run after the
+ * node's claim row is released and concurrently with a reconnect, because
+ * every write is claim-scoped — `claimRevision` bounds the transition to rows
+ * stamped before the claim, so a provider, hosted agent, or invocation changed
+ * after it is never clobbered even when both writes share a timestamp second.
+ * Callers that cannot interleave (in-process serialized paths) omit the claim
+ * fields for the legacy unconditional sweep.
+ */
+export async function settleNodeOffline(db, registry, workspaceId, nodeId, opts = {}) {
+    const claimRevision = opts.claimRevision;
     // Provider capability sets persist (offline nodes still show their manifest);
     // only their liveness flips. Zero activeAgents too, so a later
     // recomputeNodeAggregate never resurrects a dropped provider's agent count.
     await db
         .update(nodeProviders)
         .set({ status: 'offline', handlersLive: false, load: 0, loadReported: false, activeAgents: 0, lastHeartbeatAt: new Date() })
-        .where(and(eq(nodeProviders.workspaceId, workspaceId), eq(nodeProviders.nodeId, nodeId)));
+        .where(and(eq(nodeProviders.workspaceId, workspaceId), eq(nodeProviders.nodeId, nodeId), ...(claimRevision !== undefined
+        ? [lt(nodeProviders.livenessRevision, claimRevision)]
+        : [])));
     await db
         .update(agents)
         .set({ ...agentLivenessLoss(), lastSeen: new Date() })
-        .where(and(eq(agents.workspaceId, workspaceId), ne(agents.status, 'offline'), ne(agents.status, RELEASED_AGENT_STATUS), eq(agents.locationType, 'via_node'), eq(agents.locationNodeId, nodeId)));
-    await rescheduleInvocationsForLostNode(db, registry, workspaceId, nodeId);
-    if (effects?.deps && before && before.status === 'online') {
-        await emitNodeStatusEffects(effects.deps, workspaceId, {
+        .where(and(eq(agents.workspaceId, workspaceId), ne(agents.status, 'offline'), ne(agents.status, RELEASED_AGENT_STATUS), eq(agents.locationType, 'via_node'), eq(agents.locationNodeId, nodeId), ...(claimRevision !== undefined
+        ? [lt(agents.nodeLivenessRevision, claimRevision)]
+        : [])));
+    await rescheduleInvocationsForLostNode(db, registry, workspaceId, nodeId, { claimRevision });
+    if (opts.deps && opts.before && opts.before.status === 'online') {
+        await emitNodeStatusEffects(opts.deps, workspaceId, {
             status: 'offline',
             nodeId,
-            nodeName: before.name,
-            reason: effects.reason,
+            nodeName: opts.before.name,
+            reason: opts.reason,
+            ...(claimRevision !== undefined ? { claimRevision } : {}),
         });
     }
+    return true;
+}
+export async function markNodeOffline(db, registry, workspaceId, nodeId, effects) {
+    // Capture the node's name + prior status BEFORE flipping it so a durable
+    // node.status.offline is emitted only on a real online -> offline transition
+    // (never on a node that was already offline).
+    const [before] = await db
+        .select({ status: nodes.status, name: nodes.name })
+        .from(nodes)
+        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
+    await db
+        .update(nodes)
+        .set({
+        status: 'offline',
+        handlersLive: false,
+        load: 0,
+        loadReported: false,
+        activeAgents: 0,
+        lastHeartbeatAt: new Date(),
+    })
+        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
+    await settleNodeOffline(db, registry, workspaceId, nodeId, {
+        before,
+        deps: effects?.deps,
+        reason: effects?.reason,
+    });
 }
 /**
  * A provider's control connection dropped. Its capability set persists (marked
@@ -2295,6 +2391,7 @@ async function readHeartbeatDrainState(db, workspaceId, nodeId, providerName) {
         db
             .select({
             status: nodes.status,
+            role: nodes.role,
             handlersLive: nodes.handlersLive,
             maxAgents: nodes.maxAgents,
             activeAgents: nodes.activeAgents,
@@ -2474,7 +2571,7 @@ export async function handleNodeControlMessage(args) {
                 }
                 case 'node.heartbeat': {
                     const prior = await readHeartbeatDrainState(args.db, args.workspaceId, args.nodeId, frameProviderName);
-                    const beat = await heartbeatNode(args.db, args.workspaceId, args.nodeId, frameProviderName, message);
+                    const beat = await heartbeatNode(args.db, args.workspaceId, args.nodeId, frameProviderName, message, prior.node ?? null);
                     // The node row is already persisted, so emit the durable online
                     // transition before draining: a drainNode rejection must not skip the
                     // node.status.online event. Isolated so its own failure can't either.