@relaycast/engine 11.1.1 → 12.1.0

package/dist/engine/action.js

@@ -9,10 +9,12 @@ import { D1_SAFE_IN_QUERY_CHUNK_SIZE } from '../lib/queryChunks.js';
 import { toFleetWireJson } from './deliveryWire.js';
 import { emitAgentExitedEffects, emitInvocationCompletionEffects, fleetInvocationId, } from './invocationCompletion.js';
 import { runAtomic, runAtomicWrites } from '../ports/database.js';
+import { exclusiveClaimWrite, retryAfterAtMatches } from './invocationClaim.js';
 import { claimSpawnNode, chooseNodeForAction, isNodeLive, releaseNodeCapacity, reserveNodeCapacity } from './placement.js';
 import { DEFAULT_PROVIDER_NAME, capacityProviderName, getProvider, isProviderLive } from './nodeProvider.js';
 import { AGENT_TOKEN_HASH_PATTERN } from '@relaycast/types';
 import { createTaskState, expireTaskInvocation, expireTaskInvocations, taskExecution } from './taskInvocation.js';
+import { withEngineOperation } from '../lib/engineDiagnostics.js';
 const OPEN_INVOCATION_STATUSES = ['pending', 'dispatched', 'invoked'];
 const REPLAYED_INVOCATION = Symbol('replayed-action-invocation');
 export const ACTION_DISPATCH_TIMEOUT_MS = 30_000;
@@ -40,14 +42,13 @@ export const ACTION_HANDLER_UNREACHABLE_TTL_MS = 120_000;
 export const PENDING_INVOCATION_MAX_AGE_MS = 72 * 60 * 60 * 1000;
 /**
  * Sleep between the never-dispatched sweep's candidate SELECT and its atomic
- * UPDATE. `dispatchNodeInvocation` sends the frame before recording
- * `dispatchAttempts += 1`, so a sweep landing between those two dispatcher
- * calls could otherwise mark an in-flight dispatch `never_dispatched_expired`
- * while the handler already has the frame. The atomic UPDATE re-checks
- * `dispatch_attempts = 0`, so once the dispatcher's UPDATE lands, the sweep's
- * WHERE excludes the row. This grace window (default 5s) covers the send→record
- * gap by a wide margin — the dispatcher's UPDATE is one D1 round-trip — while
- * still bounding sweep latency. Set to 0 in tests to keep them fast.
+ * UPDATE. Builtin spawn and agent-hosted sends claim before the frame, so
+ * `dispatch_attempts` is already set. Release and other unclaimed sends still
+ * record the attempt after the frame. A sweep landing in that gap could mark
+ * the row `never_dispatched_expired` while the handler already has the frame.
+ * The atomic UPDATE re-checks `dispatch_attempts = 0`, so once the attempt is
+ * recorded the row is excluded. This grace window (default 5s) covers that
+ * record gap. Set to 0 in tests to keep them fast.
  */
 export const NEVER_DISPATCHED_SWEEP_GRACE_MS = 5_000;
 const ACTION_RETRY_BACKOFF_MS = 5_000;
@@ -130,11 +131,10 @@ function nextRetryAfter(attempts) {
     return new Date(Date.now() + backoff);
 }
 /**
- * Field set that moves an invocation into the live `dispatched` state once its
- * `action.invoke` frame has actually been delivered to the node. Shared by the
- * live dispatch path (`dispatchNodeAttempt`) and exported offline-queue drain
- * path (`drainNodeInvocations`) so the dispatch-timeout sweep — which keys off
- * `dispatchedAt` — and the reschedule path cover drained invocations too.
+ * Live `dispatched` fields (`dispatchedAt`, `retryAfterAt`). The timeout sweep
+ * keys off `dispatchedAt`. A builtin spawn or agent-hosted claim writes them
+ * before the frame; offline-queue drain writes them when a queued frame is
+ * delivered.
  */
 function dispatchedStateFields(opts = {}) {
     return {
@@ -763,7 +763,33 @@ async function dispatchNodeProviderInvocation(args, options = {}) {
         }
         return { accepted, pending: true };
     }
-    return dispatchNodeInvocation({ ...args });
+    // A retry passes the generation it selected. Re-reading here would adopt a
+    // claim that won while this retry was placing, and both would send. Callers
+    // without a snapshot still read the row so a dispatched spawn does not fall
+    // back to the fresh-row defaults.
+    const suppliedClaim = args.claimStatus
+        ? {
+            claimStatus: args.claimStatus,
+            observedRetryAfterAt: args.observedRetryAfterAt ?? null,
+            observedDispatchAttempts: args.observedDispatchAttempts ?? 0,
+            observedDispatchedNodeId: args.observedDispatchedNodeId ?? null,
+        }
+        : null;
+    const [claimRow] = suppliedClaim
+        ? []
+        : await args.db
+            .select({
+            status: actionInvocations.status,
+            retryAfterAt: actionInvocations.retryAfterAt,
+            dispatchAttempts: actionInvocations.dispatchAttempts,
+            dispatchedNodeId: actionInvocations.dispatchedNodeId,
+        })
+            .from(actionInvocations)
+            .where(and(eq(actionInvocations.workspaceId, args.workspaceId), eq(actionInvocations.id, args.invocationId)));
+    return dispatchNodeInvocation({
+        ...args,
+        ...(suppliedClaim ?? (claimRow ? observedInvocationClaim(claimRow) : {})),
+    });
 }
 async function dispatchRelease(args) {
     const input = recordInput(args.data.input);
@@ -1268,11 +1294,12 @@ async function dispatchSpawn(args) {
     // Placement reserved capacity for an attempt that never landed. Return it
     // first, then mark the row unsent: only the mark makes it drainable, so a
     // drain never sees it while this send is in flight or still holding the
-    // reservation.
-    if (!dispatched.accepted && !placement.queued) {
+    // reservation. A thrown send has an unknown outcome, so it is not marked
+    // and only the never-dispatched sweep ends it.
+    if (!dispatched.accepted && !placement.queued && !dispatched.reservationStillHeld) {
         await releaseNodeCapacity(args.db, args.workspaceId, nodeId);
     }
-    if (!dispatched.sent) {
+    if (!dispatched.sent && !dispatched.sendOutcomeUnknown) {
         await markUnsentNativeSpawn(args.db, args.workspaceId, invocation.id, nodeId);
     }
     return spawnResult(invocation, nodeId, dispatched);
@@ -1399,7 +1426,10 @@ export async function invokeAction(db, workspaceId, actionName, data, options =
             invocationId,
         });
     }
-    if (!action && actionName === 'release') {
+    // Destructive roster mutation. The input shape is not the guard: a payload
+    // that already carries `name` and `delete_agent` still does not release
+    // unless this caller set `authorizeBuiltinRelease`.
+    if (!action && actionName === 'release' && options.authorizeBuiltinRelease) {
         return dispatchRelease({
             db,
             registry: options.nodeConnections,
@@ -1557,6 +1587,21 @@ export async function invokeAction(db, workspaceId, actionName, data, options =
             });
         }
     }
+    // The agent CAS recorded dispatch before the frame. A handler can complete
+    // that row before this response is built, and the ack returns that terminal row.
+    if (dispatched.sent && dispatched.accepted) {
+        const [current] = await db
+            .select()
+            .from(actionInvocations)
+            .where(and(eq(actionInvocations.workspaceId, workspaceId), eq(actionInvocations.id, invocation.id)));
+        if (current && !OPEN_INVOCATION_STATUSES.some((status) => status === current.status)) {
+            return invocationAck(current, {
+                actionName,
+                handlerAgentId: action.handlerAgentId,
+                handlerNodeId: handlerAgent.locationNodeId,
+            });
+        }
+    }
     if (!dispatched.sent) {
         if (invocationId !== undefined) {
             // A keyed claim must stay consistent with its replays. Re-read it so a
@@ -1839,7 +1884,16 @@ async function dispatchNodeAttempt(db, workspaceId, invocationId, nodeId, opts)
         ...(opts.targetActionId ? { actionId: opts.targetActionId } : {}),
         ...attemptFields,
     })
-        .where(and(eq(actionInvocations.workspaceId, workspaceId), eq(actionInvocations.id, invocationId), inArray(actionInvocations.status, OPEN_INVOCATION_STATUSES), ...(opts.expectedNodeLivenessRevision !== undefined
+        .where(and(eq(actionInvocations.workspaceId, workspaceId), eq(actionInvocations.id, invocationId), ...(opts.claim
+        ? [
+            eq(actionInvocations.status, opts.claim.status),
+            retryAfterAtMatches(opts.claim.retryAfterAt),
+            eq(actionInvocations.dispatchAttempts, opts.claim.dispatchAttempts),
+            ...(opts.claim.spawnReservedAt
+                ? [eq(actionInvocations.spawnReservedAt, opts.claim.spawnReservedAt)]
+                : []),
+        ]
+        : [inArray(actionInvocations.status, OPEN_INVOCATION_STATUSES)]), ...(opts.expectedNodeLivenessRevision !== undefined
         ? [eq(actionInvocations.nodeLivenessRevision, opts.expectedNodeLivenessRevision)]
         : []), ...(opts.expectedAction ? [
         eq(actionInvocations.invocationOrigin, 'registered_action'),
@@ -1863,6 +1917,25 @@ async function dispatchNodeAttempt(db, workspaceId, invocationId, nodeId, opts)
         .returning();
     return !!updated;
 }
+/**
+ * Drop a spawn claim whose frame was not accepted. The compare matches the
+ * attempt count and deadline this claim wrote, so a failed older send cannot
+ * reset a newer winner. Restore the observed attempt count, route, deadline,
+ * and reservation token so the pending row is due for another sender.
+ */
+async function restoreUnsentSpawnClaim(db, workspaceId, invocationId, claim) {
+    await db
+        .update(actionInvocations)
+        .set({
+        status: 'pending',
+        dispatchedAt: null,
+        dispatchedNodeId: claim.restoreDispatchedNodeId,
+        spawnReservedAt: claim.restoreSpawnReservedAt,
+        retryAfterAt: claim.restoreRetryAfterAt,
+        dispatchAttempts: claim.restoreDispatchAttempts,
+    })
+        .where(and(eq(actionInvocations.workspaceId, workspaceId), eq(actionInvocations.id, invocationId), inArray(actionInvocations.status, ['pending', 'dispatched']), eq(actionInvocations.dispatchAttempts, claim.dispatchAttempts), retryAfterAtMatches(claim.retryAfterAt)));
+}
 function registeredActionDispatchPredicate(args) {
     const handlerPredicate = args.handlerAgentId
         ? sql `AND ${actions.handlerAgentId} = ${args.handlerAgentId} AND ${actions.handlerNodeId} IS NULL`
@@ -2068,6 +2141,17 @@ async function dispatchNodeInvocation(args) {
         return { accepted: false, pending: false, sent: false };
     const connectedBefore = args.registry.isProviderConnected(args.workspaceId, args.nodeId, args.providerName);
     const pending = !!args.pending || !connectedBefore;
+    const claimStatus = args.claimStatus ?? 'pending';
+    const observedRetryAfterAt = args.observedRetryAfterAt ?? null;
+    const observedDispatchAttempts = args.observedDispatchAttempts ?? 0;
+    const claimWrite = exclusiveClaimWrite({
+        observedStatus: claimStatus,
+        observedRetryAfterAt,
+        observedDispatchAttempts,
+        nextStatus: pending ? 'pending' : 'dispatched',
+        nextRetryAfterAt: args.retryAfterAt ?? null,
+        incrementAttempts: !args.skipIncrementAttempts,
+    });
     const expectedActionId = args.invocationOrigin === 'registered_action'
         ? (args.expectedActionId ?? args.actionId)
         : null;
@@ -2097,38 +2181,127 @@ async function dispatchNodeInvocation(args) {
     if (registeredNodeClaim?.taskState) {
         frame.task_execution = taskExecution({ id: args.invocationId, ...registeredNodeClaim });
     }
-    const sent = args.agent && args.actionId
-        ? await (args.registry.sendAuthorizedActionToProvider?.(args.workspaceId, args.nodeId, args.providerName, frame, {
-            kind: 'agent-action-v1',
-            invocationId: args.invocationId,
-            actionId: args.actionId,
-            handlerAgentId: args.agent.id,
-            recordAttempt: !args.skipIncrementAttempts,
-            ...(args.expectedNodeLivenessRevision !== undefined
-                ? { expectedNodeLivenessRevision: args.expectedNodeLivenessRevision }
-                : {}),
-        }) ?? false)
-        : registeredNodeClaim
+    // Native spawn has no registered handoff. Claim before the frame so a
+    // concurrent drain and sweep cannot both send. The claim matches the
+    // observed status, so a dispatched timeout retry can win, and it changes
+    // the attempt count or the deadline so a second scheduler loses. Success
+    // does not update the row again.
+    const nativeSpawnSend = isSpawnInvocation(args.action) && !args.agent && !registeredNodeClaim;
+    let nativeSpawnClaimed = false;
+    if (nativeSpawnSend) {
+        try {
+            nativeSpawnClaimed = await dispatchNodeAttempt(args.db, args.workspaceId, args.invocationId, args.nodeId, {
+                pending,
+                providerName: args.providerName,
+                retryAfterAt: claimWrite.retryAfterAt,
+                reservationHeld: args.reservationHeld,
+                skipIncrementAttempts: args.skipIncrementAttempts,
+                expectedNodeLivenessRevision: args.expectedNodeLivenessRevision,
+                claim: {
+                    status: claimStatus,
+                    retryAfterAt: observedRetryAfterAt,
+                    dispatchAttempts: observedDispatchAttempts,
+                    ...(args.preserveSpawnReservation && args.observedSpawnReservedAt
+                        ? { spawnReservedAt: args.observedSpawnReservedAt }
+                        : {}),
+                },
+            });
+        }
+        catch {
+            // The claim write did not commit: no frame left, and the row has no
+            // reservation marker for a later sweep to release. A false result lets
+            // the caller return the slot it already reserved and leave the
+            // invocation pending for another sender.
+            return { accepted: false, pending: false, sent: false };
+        }
+    }
+    if (nativeSpawnSend && !nativeSpawnClaimed) {
+        return { accepted: false, pending: false, sent: false };
+    }
+    const restoreSpawnClaim = () => restoreUnsentSpawnClaim(args.db, args.workspaceId, args.invocationId, {
+        dispatchAttempts: claimWrite.dispatchAttempts,
+        retryAfterAt: claimWrite.retryAfterAt,
+        restoreRetryAfterAt: observedRetryAfterAt,
+        restoreDispatchAttempts: observedDispatchAttempts,
+        restoreDispatchedNodeId: args.observedDispatchedNodeId ?? null,
+        restoreSpawnReservedAt: args.preserveSpawnReservation ? args.observedSpawnReservedAt ?? null : null,
+    });
+    // A failed restore leaves the marker set. Returning the slot as well lets
+    // the timeout sweep release it a second time.
+    const restoreSpawnClaimOrHold = async () => {
+        try {
+            await restoreSpawnClaim();
+            return false;
+        }
+        catch {
+            return true;
+        }
+    };
+    let sent = false;
+    try {
+        sent = args.agent && args.actionId
             ? await (args.registry.sendAuthorizedActionToProvider?.(args.workspaceId, args.nodeId, args.providerName, frame, {
-                kind: 'registered-node-action-v2',
+                kind: 'agent-action-v1',
                 invocationId: args.invocationId,
-                actionId: registeredNodeClaim.actionId,
-                dispatchAttempt: registeredNodeClaim.dispatchAttempts,
-                invocationActionName: registeredNodeClaim.actionName,
-                actionName: args.action,
+                actionId: args.actionId,
+                handlerAgentId: args.agent.id,
+                recordAttempt: !args.skipIncrementAttempts,
+                pending,
+                retryAfterAt: claimWrite.retryAfterAt,
+                claimStatus,
+                observedRetryAfterAt,
+                observedDispatchAttempts,
+                reservationHeld: !!args.reservationHeld,
+                ...(args.expectedNodeLivenessRevision !== undefined
+                    ? { expectedNodeLivenessRevision: args.expectedNodeLivenessRevision }
+                    : {}),
             }) ?? false)
-            : guardedReleaseHash || guardedReleaseAgentId
+            : registeredNodeClaim
                 ? await (args.registry.sendAuthorizedActionToProvider?.(args.workspaceId, args.nodeId, args.providerName, frame, {
-                    kind: 'release-generation-v1',
+                    kind: 'registered-node-action-v2',
                     invocationId: args.invocationId,
-                    agentName: typeof args.input.name === 'string' ? args.input.name : '',
-                    ...(guardedReleaseHash && guardedReleaseAgentId
-                        ? { expectedTokenHash: guardedReleaseHash, expectedAgentId: guardedReleaseAgentId }
-                        : guardedReleaseHash
-                            ? { expectedTokenHash: guardedReleaseHash }
-                            : { expectedAgentId: guardedReleaseAgentId }),
+                    actionId: registeredNodeClaim.actionId,
+                    dispatchAttempt: registeredNodeClaim.dispatchAttempts,
+                    invocationActionName: registeredNodeClaim.actionName,
+                    actionName: args.action,
                 }) ?? false)
-                : await args.registry.sendToProvider(args.workspaceId, args.nodeId, args.providerName, frame);
+                : guardedReleaseHash || guardedReleaseAgentId
+                    ? await (args.registry.sendAuthorizedActionToProvider?.(args.workspaceId, args.nodeId, args.providerName, frame, {
+                        kind: 'release-generation-v1',
+                        invocationId: args.invocationId,
+                        agentName: typeof args.input.name === 'string' ? args.input.name : '',
+                        ...(guardedReleaseHash && guardedReleaseAgentId
+                            ? { expectedTokenHash: guardedReleaseHash, expectedAgentId: guardedReleaseAgentId }
+                            : guardedReleaseHash
+                                ? { expectedTokenHash: guardedReleaseHash }
+                                : { expectedAgentId: guardedReleaseAgentId }),
+                    }) ?? false)
+                    : await args.registry.sendToProvider(args.workspaceId, args.nodeId, args.providerName, frame);
+    }
+    catch (error) {
+        // A thrown send never reached the false-result path. Restore only this
+        // claim, then report a false result so the caller releases the capacity
+        // it reserved. The outcome is unknown, so callers must not queue another
+        // send. Registered-action crashes still propagate: their claim has to stay
+        // in place for the keyed replay.
+        if (nativeSpawnClaimed) {
+            const reservationStillHeld = await restoreSpawnClaimOrHold();
+            return {
+                accepted: false,
+                pending: false,
+                sent: false,
+                sendOutcomeUnknown: true,
+                ...(reservationStillHeld ? { reservationStillHeld: true } : {}),
+            };
+        }
+        throw error;
+    }
+    if (!sent && nativeSpawnClaimed) {
+        const reservationStillHeld = await restoreSpawnClaimOrHold();
+        if (reservationStillHeld) {
+            return { accepted: false, pending: false, sent: false, reservationStillHeld: true };
+        }
+    }
     if (!sent && (guardedReleaseHash || guardedReleaseAgentId)) {
         // An adapter compiled against the older owner-authorization contract may
         // expose the method but reject the new proof kind. Settle any still-open
@@ -2157,10 +2330,11 @@ async function dispatchNodeInvocation(args) {
         }
         return { accepted: false, pending: false, sent: false };
     }
-    // Registered node-action state was committed by the pre-send handoff CAS;
-    // agent-action state was committed by the socket owner's authorization CAS.
+    // Registered node-action and native-spawn state was committed by the
+    // pre-send claim. Agent-action state was committed by the socket owner's
+    // authorization CAS, then re-read below in case the handler already finished.
     // Do not write either generation again after the frame is accepted.
-    if (registeredNodeClaim) {
+    if (registeredNodeClaim || nativeSpawnClaimed) {
         return { accepted: true, pending, sent: true };
     }
     if (args.agent && args.actionId) {
@@ -2313,6 +2487,7 @@ export async function drainNodeInvocations(db, registry, workspaceId, nodeId, op
         input: actionInvocations.input,
         providerAcceptedAttempt: actionInvocations.providerAcceptedAttempt,
         dispatchAttempts: actionInvocations.dispatchAttempts,
+        retryAfterAt: actionInvocations.retryAfterAt,
         spawnReservedAt: actionInvocations.spawnReservedAt,
         dispatchedNodeId: actionInvocations.dispatchedNodeId,
         dispatchedProvider: actionInvocations.dispatchedProvider,
@@ -2434,15 +2609,23 @@ export async function drainNodeInvocations(db, registry, workspaceId, nodeId, op
                 retryAfterAt: new Date(Date.now() + ACTION_DISPATCH_TIMEOUT_MS),
                 reservationHeld,
                 skipIncrementAttempts: row.dispatchedNodeId === nodeId,
+                claimStatus: 'pending',
+                observedRetryAfterAt: row.retryAfterAt,
+                observedDispatchAttempts: row.dispatchAttempts,
+                observedDispatchedNodeId: row.dispatchedNodeId,
+                observedSpawnReservedAt: row.spawnReservedAt,
+                preserveSpawnReservation: reservationHeld && !reservedForDrain,
             });
             if (dispatched.accepted) {
                 drained++;
             }
-            else if (reservedForDrain) {
+            else if (reservedForDrain && !dispatched.reservationStillHeld) {
                 await releaseNodeCapacity(db, workspaceId, nodeId);
                 reservedForDrain = false;
             }
-            if (!dispatched.sent)
+            // A thrown send may have delivered the frame. Leave an unsent spawn
+            // unmarked so a later drain does not send it again.
+            if (!dispatched.sent && !dispatched.sendOutcomeUnknown)
                 await remarkUnsent();
         }
         catch {
@@ -2459,6 +2642,16 @@ export async function drainNodeInvocations(db, registry, workspaceId, nodeId, op
     }
     return drained;
 }
+function observedInvocationClaim(invocation) {
+    if (invocation.status !== 'pending' && invocation.status !== 'dispatched' && invocation.status !== 'invoked')
+        return {};
+    return {
+        claimStatus: invocation.status,
+        observedRetryAfterAt: invocation.retryAfterAt ?? null,
+        observedDispatchAttempts: invocation.dispatchAttempts ?? 0,
+        observedDispatchedNodeId: invocation.dispatchedNodeId ?? null,
+    };
+}
 export async function rescheduleNodeInvocation(db, registry, invocation, opts = {}) {
     // Once the exact provider accepted this generation, pruning only clears its
     // materialized action FK. It must not make timeout, disconnect, or inventory
@@ -2473,9 +2666,36 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
     }
     // Only release capacity the invocation actually holds. A shadowed spawn
     // dropped its native reservation at dispatch (spawnReservedAt is null), so it
-    // must not release capacity that now belongs to another spawn.
+    // must not release capacity that now belongs to another spawn. Clear the
+    // marker once so overlapping retries release that one hold, not every hold
+    // on the node.
     if (invocation.dispatchedNodeId && isSpawnInvocation(invocation.actionName) && invocation.spawnReservedAt) {
-        await releaseNodeCapacity(db, invocation.workspaceId, invocation.dispatchedNodeId);
+        const heldAt = invocation.spawnReservedAt;
+        const nodeId = invocation.dispatchedNodeId;
+        const holdWhere = and(eq(actionInvocations.workspaceId, invocation.workspaceId), eq(actionInvocations.id, invocation.id), eq(actionInvocations.status, invocation.status), eq(actionInvocations.dispatchAttempts, invocation.dispatchAttempts ?? 0), retryAfterAtMatches(invocation.retryAfterAt ?? null), ...(opts.expectedNodeLivenessRevision !== undefined
+            ? [eq(actionInvocations.nodeLivenessRevision, opts.expectedNodeLivenessRevision)]
+            : []));
+        try {
+            // Clear the marker and decrement together. A transaction rolls both
+            // back when the counter update throws. A handle without one can commit
+            // the clear first; the catch puts that marker back.
+            await runAtomic(db, async (tx) => {
+                const [releasedHold] = await tx
+                    .update(actionInvocations)
+                    .set({ spawnReservedAt: null })
+                    .where(and(holdWhere, eq(actionInvocations.spawnReservedAt, heldAt)))
+                    .returning({ id: actionInvocations.id });
+                if (releasedHold)
+                    await releaseNodeCapacity(tx, invocation.workspaceId, nodeId);
+            });
+        }
+        catch (error) {
+            await db
+                .update(actionInvocations)
+                .set({ spawnReservedAt: heldAt })
+                .where(and(holdWhere, isNull(actionInvocations.spawnReservedAt)));
+            throw error;
+        }
     }
     // A shadowed spawn stays pinned to its node: while a `spawn:<harness>` shadow
     // is registered there, native capacity for that harness is unreachable — even
@@ -2523,6 +2743,7 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
             invocationOrigin: invocation.invocationOrigin,
             retryAfterAt: opts.retryAfterAt ?? null,
             expectedNodeLivenessRevision: opts.expectedNodeLivenessRevision,
+            ...observedInvocationClaim(invocation),
         });
         return dispatched.accepted;
     }
@@ -2572,6 +2793,7 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
                         queue: target.queue,
                         reservationHeld: false,
                         expectedNodeLivenessRevision: opts.expectedNodeLivenessRevision,
+                        ...observedInvocationClaim(invocation),
                     }, { failIfUnavailable: false });
                     if (dispatched.providerUnavailable) {
                         providerRejectedNodeIds.add(placement.node.id);
@@ -2589,15 +2811,32 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
                     const expectedActionId = invocation.invocationOrigin === 'registered_action'
                         ? invocation.actionId
                         : null;
+                    const queuedClaim = invocation.status === 'pending' || invocation.status === 'dispatched' || invocation.status === 'invoked'
+                        ? exclusiveClaimWrite({
+                            observedStatus: invocation.status,
+                            observedRetryAfterAt: invocation.retryAfterAt ?? null,
+                            observedDispatchAttempts: invocation.dispatchAttempts ?? 0,
+                            nextStatus: 'pending',
+                            nextRetryAfterAt: opts.retryAfterAt ?? null,
+                            incrementAttempts: true,
+                        })
+                        : null;
                     const accepted = await dispatchNodeAttempt(db, invocation.workspaceId, invocation.id, placement.node.id, {
                         providerName,
                         pending: true,
-                        retryAfterAt: opts.retryAfterAt ?? null,
+                        retryAfterAt: queuedClaim?.retryAfterAt ?? opts.retryAfterAt ?? null,
                         reservationHeld: false,
                         expectedAction: expectedActionId
                             ? { id: expectedActionId, name: actionToSend }
                             : undefined,
                         targetActionId: target?.id,
+                        ...(queuedClaim ? {
+                            claim: {
+                                status: invocation.status,
+                                retryAfterAt: invocation.retryAfterAt ?? null,
+                                dispatchAttempts: invocation.dispatchAttempts ?? 0,
+                            },
+                        } : {}),
                         expectedNodeLivenessRevision: opts.expectedNodeLivenessRevision,
                     });
                     if (!accepted && invocation.invocationOrigin === 'registered_action') {
@@ -2605,21 +2844,41 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
                     }
                     return accepted;
                 }
-                const dispatched = await dispatchNodeInvocation({
-                    db,
-                    registry,
-                    workspaceId: invocation.workspaceId,
-                    invocationId: invocation.id,
-                    nodeId: placement.node.id,
-                    providerName,
-                    action: actionToSend,
-                    actionId: target?.id ?? null,
-                    expectedActionId: invocation.actionId,
-                    invocationOrigin: invocation.invocationOrigin,
-                    input,
-                    reservationHeld,
-                    expectedNodeLivenessRevision: opts.expectedNodeLivenessRevision,
-                });
+                let dispatched;
+                try {
+                    dispatched = await dispatchNodeInvocation({
+                        db,
+                        registry,
+                        workspaceId: invocation.workspaceId,
+                        invocationId: invocation.id,
+                        nodeId: placement.node.id,
+                        providerName,
+                        action: actionToSend,
+                        actionId: target?.id ?? null,
+                        expectedActionId: invocation.actionId,
+                        invocationOrigin: invocation.invocationOrigin,
+                        input,
+                        reservationHeld,
+                        retryAfterAt: opts.retryAfterAt ?? null,
+                        expectedNodeLivenessRevision: opts.expectedNodeLivenessRevision,
+                        ...observedInvocationClaim(invocation),
+                    });
+                }
+                catch (error) {
+                    // A throw before the claim commits has no marker, so return the
+                    // new slot. A restore failure returns reservationStillHeld instead
+                    // of throwing, and the check below leaves that slot in place.
+                    if (reservationHeld) {
+                        await releaseNodeCapacity(db, invocation.workspaceId, placement.node.id);
+                    }
+                    throw error;
+                }
+                // claimSpawnNode incremented reservedAgents for this attempt. Release
+                // it when the pending claim loses or the send never starts. Leave it
+                // when the reservation marker is still set.
+                if (!dispatched.accepted && !dispatched.sent && reservationHeld && !dispatched.reservationStillHeld) {
+                    await releaseNodeCapacity(db, invocation.workspaceId, placement.node.id);
+                }
                 return dispatched.accepted;
             }
             catch {
@@ -2642,7 +2901,7 @@ export async function rescheduleNodeInvocation(db, registry, invocation, opts =
         spawnReservedAt: null,
         retryAfterAt: opts.retryAfterAt ?? nextRetryAfter(invocation.dispatchAttempts + 1),
     })
-        .where(and(eq(actionInvocations.workspaceId, invocation.workspaceId), eq(actionInvocations.id, invocation.id), inArray(actionInvocations.status, OPEN_INVOCATION_STATUSES), ...(opts.expectedNodeLivenessRevision !== undefined
+        .where(and(eq(actionInvocations.workspaceId, invocation.workspaceId), eq(actionInvocations.id, invocation.id), eq(actionInvocations.status, invocation.status), eq(actionInvocations.dispatchAttempts, invocation.dispatchAttempts ?? 0), retryAfterAtMatches(invocation.retryAfterAt ?? null), ...(opts.expectedNodeLivenessRevision !== undefined
         ? [eq(actionInvocations.nodeLivenessRevision, opts.expectedNodeLivenessRevision)]
         : [])));
     return false;
@@ -2750,6 +3009,7 @@ export async function completeNodeInvocation(db, registry, workspaceId, nodeId,
         spawnReservedAt: actionInvocations.spawnReservedAt,
         attemptedNodeIds: actionInvocations.attemptedNodeIds,
         dispatchAttempts: actionInvocations.dispatchAttempts,
+        retryAfterAt: actionInvocations.retryAfterAt,
         createdAt: actionInvocations.createdAt,
         completedAt: actionInvocations.completedAt,
     })
@@ -2866,11 +3126,9 @@ async function failNeverDispatchedExpiredInvocations(db, maxAgeMs, graceMs, comp
         .where(and(eq(actionInvocations.status, 'pending'), eq(actionInvocations.dispatchAttempts, 0), lte(actionInvocations.createdAt, cutoff)));
     if (rows.length === 0)
         return;
-    // Give any concurrent dispatcher's send→record window (`dispatchNodeAttempt`
-    // UPDATE, one D1 round-trip) time to close before the atomic UPDATE fires.
-    // Combined with the UPDATE's `dispatch_attempts = 0` re-check, this makes an
-    // in-flight dispatch reliably invisible to the age sweep instead of racing
-    // with it. See `NEVER_DISPATCHED_SWEEP_GRACE_MS`.
+    // Unclaimed sends still record `dispatchAttempts` after the frame. Wait out
+    // that gap, then the UPDATE's `dispatch_attempts = 0` re-check skips any row
+    // the dispatcher has claimed. See `NEVER_DISPATCHED_SWEEP_GRACE_MS`.
     if (graceMs > 0) {
         await new Promise((resolve) => setTimeout(resolve, graceMs));
     }
@@ -2916,6 +3174,9 @@ async function failNeverDispatchedInvocationRows(db, workspaceId, invocationIds,
         .returning();
 }
 export async function sweepTimedOutInvocations(db, registry, opts = {}) {
+    return withEngineOperation('invocations.sweep', undefined, () => sweepTimedOutInvocationsUnlabeled(db, registry, opts));
+}
+async function sweepTimedOutInvocationsUnlabeled(db, registry, opts) {
     await expireTaskInvocations(db);
     const sweepOpts = typeof opts === 'number' || opts === null ? {} : opts;
     const timeoutMs = typeof opts === 'number' ? opts : sweepOpts.timeoutMs ?? ACTION_DISPATCH_TIMEOUT_MS;