@relaycast/engine 8.16.0 → 9.0.0

package/dist/engine/node.js

@@ -16,6 +16,7 @@ import { serializeNodeOp } from './nodeLock.js';
 import { completeNodeInvocation, dispatchCapacitySpawn, rescheduleInvocationsForLostNode, rescheduleNodeInvocation, } from './action.js';
 import { emitInvocationCompletionEffects, emitAgentExitedEffects, emitNodeStatusEffects, fleetInvocationId } from './invocationCompletion.js';
 import { ackDeliveriesUpToSeq, deliverPendingToNode } from './delivery.js';
+const NODE_ROTATION_RECOVERY_TTL_MS = 24 * 60 * 60 * 1_000;
 function resolveProviderIdentity(message, connectionId) {
     if (message.provider)
         return message.provider;
@@ -278,8 +279,8 @@ function normalizeLegacyAdapter(adapter) {
         return 'ws.node.v1';
     return adapter;
 }
-export async function createNodeToken(db, workspaceId, data) {
-    const token = `nt_live_${randomHex(24)}`;
+export async function createNodeToken(db, workspaceId, data, authorization = {}) {
+    const token = authorization.rotationRecovery?.replacementToken ?? `nt_live_${randomHex(24)}`;
     const tokenHash = await sha256Hex(token);
     const name = data.name.startsWith('#') ? data.name.slice(1) : data.name;
     // Reject names that are empty or still #-prefixed after normalization: an
@@ -292,7 +293,7 @@ export async function createNodeToken(db, workspaceId, data) {
     // lookup is only a fallback for callers without a stable id. A name held by
     // a *different* node id is a conflict — mirroring node.register — never a
     // silent rewrite of the other node.
-    const existing = await resolveNodeForEnroll(db, workspaceId, data);
+    const existing = await resolveNodeForEnroll(db, workspaceId, data, authorization.expectedTokenHash);
     if (data.node_id !== undefined) {
         const nameHolder = await getNodeByName(db, workspaceId, name);
         if (nameHolder && nameHolder.id !== data.node_id) {
@@ -320,11 +321,42 @@ export async function createNodeToken(db, workspaceId, data) {
         throw codedError('Direct nodes can bind at most one agent', 'direct_node_capacity_exceeded', 400);
     }
     if (existing) {
+        const hasPresentedTokenProof = authorization.expectedTokenHash !== undefined;
+        const authorizedDirectNode = authorization.directAgentId !== undefined
+            && existing.id === directNodeIdForAgent(authorization.directAgentId)
+            && existing.role === 'direct';
+        const recovery = authorization.rotationRecovery;
+        const isIdempotentRotationReplay = (candidate) => hasPresentedTokenProof
+            && recovery !== undefined
+            && candidate.tokenHash === tokenHash
+            && candidate.previousTokenHash === authorization.expectedTokenHash
+            && candidate.rotationIdempotencyKeyHash === recovery.idempotencyKeyHash
+            && candidate.rotationRequestDigest === recovery.requestDigest
+            && candidate.previousTokenExpiresAt !== null
+            && candidate.previousTokenExpiresAt.getTime() > now.getTime();
+        if (isIdempotentRotationReplay(existing)) {
+            return { ...publicNode(existing), token };
+        }
+        if (!hasPresentedTokenProof && !authorizedDirectNode) {
+            throw codedError('Node already exists; its current node token is required to rotate or modify it', 'node_token_proof_required', 409);
+        }
         const [updated] = await db
             .update(nodes)
             .set({
             name,
             tokenHash,
+            previousTokenHash: hasPresentedTokenProof && !authorizedDirectNode && recovery
+                ? authorization.expectedTokenHash
+                : null,
+            previousTokenExpiresAt: hasPresentedTokenProof && !authorizedDirectNode && recovery
+                ? new Date(now.getTime() + NODE_ROTATION_RECOVERY_TTL_MS)
+                : null,
+            rotationIdempotencyKeyHash: hasPresentedTokenProof && !authorizedDirectNode && recovery
+                ? recovery.idempotencyKeyHash
+                : null,
+            rotationRequestDigest: hasPresentedTokenProof && !authorizedDirectNode && recovery
+                ? recovery.requestDigest
+                : null,
             kind,
             role,
             deliveryAdapter,
@@ -344,8 +376,23 @@ export async function createNodeToken(db, workspaceId, data) {
             load: 0,
             loadReported: false,
         })
-            .where(eq(nodes.id, existing.id))
+            // Token proof is a generation fence, not just a read-time check. If a
+            // concurrent rotation changes the credential after `existing` was read,
+            // this request must not overwrite the winner with stale authority.
+            .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, existing.id), ...(hasPresentedTokenProof && !authorizedDirectNode
+            ? [eq(nodes.tokenHash, authorization.expectedTokenHash)]
+            : [])))
             .returning();
+        if (!updated) {
+            // Two identical retries can both read the previous generation before
+            // either update runs. If the other request committed this exact keyed
+            // rotation, recover its result instead of returning a spurious conflict.
+            const winner = recovery ? await getNodeById(db, workspaceId, existing.id) : null;
+            if (winner && isIdempotentRotationReplay(winner)) {
+                return { ...publicNode(winner), token };
+            }
+            throw codedError('Node token changed before enrollment completed; retry with the current node token', 'node_token_proof_required', 409);
+        }
         return { ...publicNode(updated), token };
     }
     const [created] = await db
@@ -417,21 +464,18 @@ const MACHINE_MATCH_SCAN_LIMIT = 20;
  * separate nodes. Declining the match costs an extra roster row; rejecting the
  * enrollment outright would break a whole fleet booted from one image.
  *
- * Ordered oldest-first so a roster that already holds several rows for one
- * machine converges onto its earliest reusable row instead of picking
- * arbitrarily.
+ * Current-token proof selects its exact reusable generation when several
+ * brokers share one machine. Without matching proof, the fallback remains
+ * oldest-first so a roster converges instead of picking arbitrarily.
  */
-export async function getBrokerNodeByMachineId(db, workspaceId, machineId) {
+export async function getBrokerNodeByMachineId(db, workspaceId, machineId, expectedTokenHash) {
     // Exclude live rows in SQL so the bound below counts only reusable
     // candidates. Bounding first and filtering after would let a machine with
     // enough live brokers hide its one offline row past the limit, and enrollment
     // would insert a new row instead of reusing it.
     const now = new Date();
     const liveCutoff = new Date(now.getTime() - NODE_LIVENESS_TTL_MS);
-    const candidates = await db
-        .select()
-        .from(nodes)
-        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.machineId, machineId), eq(nodes.role, 'broker'), 
+    const reusableWhere = and(eq(nodes.workspaceId, workspaceId), eq(nodes.machineId, machineId), eq(nodes.role, 'broker'), 
     // Mirrors isReusableForMachineMatch in full, so the scan bound below
     // counts only rows that are actually reusable. A stale proof alone is not
     // enough: a broker that registered recently has a fresh lastHeartbeatAt
@@ -441,7 +485,25 @@ export async function getBrokerNodeByMachineId(db, workspaceId, machineId) {
     // not live — mirrors isNodeLive
     or(ne(nodes.status, 'online'), isNull(nodes.lastHeartbeatAt), lt(nodes.lastHeartbeatAt, liveCutoff)), 
     // no future-dated heartbeat — server clock rollback
-    or(isNull(nodes.lastHeartbeatAt), lte(nodes.lastHeartbeatAt, now))))
+    or(isNull(nodes.lastHeartbeatAt), lte(nodes.lastHeartbeatAt, now)));
+    // Several legitimate brokers may share one machine id (VM clones and
+    // containers are common). When the caller retains its token, select that
+    // exact reusable generation before the oldest-row fallback so one broker's
+    // proof is never checked against a sibling row.
+    if (expectedTokenHash !== undefined) {
+        const [provenCandidate] = await db
+            .select()
+            .from(nodes)
+            .where(and(reusableWhere, eq(nodes.tokenHash, expectedTokenHash)))
+            .limit(1);
+        if (provenCandidate && isReusableForMachineMatch(provenCandidate)) {
+            return provenCandidate;
+        }
+    }
+    const candidates = await db
+        .select()
+        .from(nodes)
+        .where(reusableWhere)
         .orderBy(asc(nodes.createdAt), asc(nodes.id))
         .limit(MACHINE_MATCH_SCAN_LIMIT);
     // The SQL above is a pre-filter for `isReusableForMachineMatch`, which is the
@@ -473,8 +535,10 @@ export function requestedNodeRole(data) {
  * The machine_id step is what stops the roster refilling. A host that enrolls
  * without a persisted node_id arrives under a fresh name on every boot, and
  * each of those names used to mint a brand-new row that nothing ever reclaimed.
- * Falling back to the machine's existing broker turns re-enrollment into a
- * token rotation on the row that is already there.
+ * Falling back to the machine's existing broker lets an authorized holder
+ * rotate the row that is already there. createNodeToken separately requires
+ * current-node-token proof for enrollment, or authenticated-agent authorization
+ * for that agent's canonical direct node, before an existing row can change.
  *
  * node_id and name still win, so a caller that pins either keeps the exact
  * identity it asked for; passing node_id is the way to opt out of machine
@@ -484,17 +548,41 @@ export function requestedNodeRole(data) {
  * getBrokerNodeByMachineId — so a running broker is never adopted by another
  * caller presenting its machine_id.
  */
-export async function resolveNodeForEnroll(db, workspaceId, data) {
+export async function resolveNodeForEnroll(db, workspaceId, data, expectedTokenHash) {
     if (data.node_id !== undefined) {
-        return getNodeById(db, workspaceId, data.node_id);
+        const byId = await getNodeById(db, workspaceId, data.node_id);
+        if (byId)
+            return byId;
+        // A caller retrying a committed rotation may change the request body by
+        // mistake. Resolve the superseded proof below so that request is rejected
+        // against the rotated row rather than treated as a new-node enrollment.
+        return expectedTokenHash === undefined
+            ? null
+            : getNodeByPreviousTokenHash(db, workspaceId, expectedTokenHash);
     }
     const byName = await getNodeByName(db, workspaceId, data.name);
     if (byName)
         return byName;
     if (data.machine_id !== undefined && requestedNodeRole(data) === 'broker') {
-        return getBrokerNodeByMachineId(db, workspaceId, data.machine_id);
+        const byMachine = await getBrokerNodeByMachineId(db, workspaceId, data.machine_id, expectedTokenHash);
+        if (byMachine)
+            return byMachine;
     }
-    return null;
+    // Machine reuse clears provenLiveAt, so a committed keyed rotation no
+    // longer qualifies for the normal reusable-row lookup. Falling back to the
+    // bounded superseded proof makes an exact retry recoverable and a changed
+    // body/key fail closed instead of inserting a second node.
+    return expectedTokenHash === undefined
+        ? null
+        : getNodeByPreviousTokenHash(db, workspaceId, expectedTokenHash);
+}
+async function getNodeByPreviousTokenHash(db, workspaceId, tokenHash) {
+    const [node] = await db
+        .select()
+        .from(nodes)
+        .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.previousTokenHash, tokenHash)))
+        .limit(1);
+    return node ?? null;
 }
 export async function registerNode(db, workspaceId, authenticatedNodeId, message, provider) {
     if (message.node_id !== authenticatedNodeId) {