diff --git a/internal/cloud/hwpod-node-ops.test.ts b/internal/cloud/hwpod-node-ops.test.ts index 348bf997..f21f0749 100644 --- a/internal/cloud/hwpod-node-ops.test.ts +++ b/internal/cloud/hwpod-node-ops.test.ts @@ -522,6 +522,63 @@ test("WebSocket registry clears pre-accept timeout and disconnect for same-proce } }); +test("WebSocket registry replays an accepted plan after disconnect without rebuilding registry", async () => { + const registry = createHwpodNodeWsRegistry({ operationLedger: createMemoryHwpodOperationLedger() }); + const plan = { ...samplePlan(), nodeId: "pc-host-1", planId: "accepted-disconnect-replay-plan" }; + const executedPlans = new Set(); + const dispatchedPlanIds: string[] = []; + let authoritativeExecutionCount = 0; + let dispatchCount = 0; + const firstSocket = { + send(value: string) { + const message = JSON.parse(value); + if (message.type !== "hwpod-node-ops") return; + dispatchCount += 1; + dispatchedPlanIds.push(message.plan.planId); + if (!executedPlans.has(message.plan.planId)) { + executedPlans.add(message.plan.planId); + authoritativeExecutionCount += 1; + } + queueMicrotask(() => { + registry.handleBunMessage(firstSocket, JSON.stringify({ type: "hwpod-node-ops-accepted", requestId: message.requestId, nodeId: "pc-host-1" })); + registry.closeBunSocket(firstSocket); + }); + }, + close() {} + }; + registry.openBunSocket(firstSocket); + registry.handleBunMessage(firstSocket, JSON.stringify({ type: "register", nodeId: "pc-host-1", maxInFlight: 1 })); + const first = await registry.dispatch(plan, { requestId: "req_accepted_disconnected" }, { timeoutMs: 1000 }); + assert.equal(first.status, "blocked"); + + const replaySocket = { + send(value: string) { + const message = JSON.parse(value); + if (message.type !== "hwpod-node-ops") return; + dispatchCount += 1; + dispatchedPlanIds.push(message.plan.planId); + const replayed = executedPlans.has(message.plan.planId); + if (!replayed) { + executedPlans.add(message.plan.planId); + authoritativeExecutionCount += 1; + } + queueMicrotask(() => { + registry.handleBunMessage(replaySocket, JSON.stringify({ type: "hwpod-node-ops-accepted", requestId: message.requestId, nodeId: "pc-host-1" })); + registry.handleBunMessage(replaySocket, JSON.stringify({ type: "hwpod-node-ops-result", requestId: message.requestId, nodeId: "pc-host-1", result: { ok: true, status: "completed", planId: message.plan.planId, replayed, results: [] } })); + }); + }, + close() {} + }; + registry.openBunSocket(replaySocket); + registry.handleBunMessage(replaySocket, JSON.stringify({ type: "register", nodeId: "pc-host-1", maxInFlight: 1 })); + const recovered = await registry.dispatch(plan, { requestId: "req_accepted_replay" }, { timeoutMs: 1000 }); + assert.equal(recovered.status, "completed"); + assert.equal(recovered.replayed, true); + assert.deepEqual(dispatchedPlanIds, [plan.planId, plan.planId]); + assert.equal(authoritativeExecutionCount, 1); + assert.equal(dispatchCount, 2); +}); + test("WebSocket registry bounds ledger write failures and permits node replay recovery", async () => { const records = new Map(); const memoryLedger = createMemoryHwpodOperationLedger({ records }); diff --git a/internal/cloud/hwpod-node-ws-registry.ts b/internal/cloud/hwpod-node-ws-registry.ts index 2c073221..016ae205 100644 --- a/internal/cloud/hwpod-node-ws-registry.ts +++ b/internal/cloud/hwpod-node-ws-registry.ts @@ -38,6 +38,7 @@ type PendingDispatch = { timer: ReturnType; accepted: boolean; ledgerWrite: Promise; + promise: Promise; resolve: (value: any) => void; }; @@ -159,6 +160,7 @@ export function createHwpodNodeWsRegistry(options: any = {}) { timer: null as unknown as ReturnType, accepted: false, ledgerWrite: Promise.resolve(true), + promise, resolve: resolveDispatch } satisfies PendingDispatch; if (planId) operations.set(planId, { planId, nodeId, operation, promise, result: null }); @@ -173,7 +175,7 @@ export function createHwpodNodeWsRegistry(options: any = {}) { } void promise.then((result) => { const authoritativeOperation = planId ? operations.get(planId) : null; - if (authoritativeOperation) { + if (authoritativeOperation?.promise === promise) { authoritativeOperation.result = result; authoritativeOperation.operation = operationFinished( authoritativeOperation.operation, @@ -327,7 +329,7 @@ export function createHwpodNodeWsRegistry(options: any = {}) { clearTimeout(waiter.timer); pending.delete(waiter.requestId); finishPendingOperation(waiter, status, now(), "hwpod_node_unavailable"); - if (waiter.planId && !waiter.accepted) operations.delete(waiter.planId); + if (waiter.planId && operations.get(waiter.planId)?.promise === waiter.promise) operations.delete(waiter.planId); waiter.resolve(result); }