diff --git a/server/src/__tests__/question-response-delivery.test.ts b/server/src/__tests__/question-response-delivery.test.ts index 8c5cc35b72b..62b956c1a82 100644 --- a/server/src/__tests__/question-response-delivery.test.ts +++ b/server/src/__tests__/question-response-delivery.test.ts @@ -308,6 +308,111 @@ describeEmbeddedPostgres("question response delivery", () => { await expect(deliveryPromise).resolves.toBeNull(); }); + it("fences a stale worker after a newer claim generation takes ownership", async () => { + const seeded = await seed({ sourceStatus: "running" }); + let releaseFirstWake!: (value: { id: string; driverKind: string }) => void; + const firstWakeup = vi.fn(() => new Promise<{ id: string; driverKind: string }>((resolve) => { + releaseFirstWake = resolve; + })); + const firstService = questionResponseDeliveryService(db, { + heartbeat: { wakeup: firstWakeup } as never, + steer: vi.fn(), + claimStaleMs: 40, + claimRefreshMs: 5, + }); + + const firstDelivery = firstService.deliver(seeded.interaction.id); + await vi.waitFor(() => expect(firstWakeup).toHaveBeenCalledTimes(1)); + + // Simulate recovery after the first worker stopped renewing. The next + // claim increments attemptCount, which is the fencing generation. + await db.update(issueQuestionResponseDeliveries).set({ + status: "pending", + lastAttemptAt: new Date(0), + }).where(eq(issueQuestionResponseDeliveries.interactionId, seeded.interaction.id)); + + const secondRunId = randomUUID(); + const secondWakeup = vi.fn().mockImplementation(async () => db.insert(heartbeatRuns).values({ + id: secondRunId, + companyId: seeded.companyId, + agentId: seeded.agentId, + invocationSource: "automation", + status: "queued", + runtimeMode: "legacy", + driverKind: "codex", + contextSnapshot: { issueId: seeded.issueId }, + }).returning().then((rows) => rows[0]!)); + const secondOutcome = await questionResponseDeliveryService(db, { + heartbeat: { wakeup: secondWakeup } as never, + steer: vi.fn(), + }).deliver(seeded.interaction.id); + + expect(secondOutcome).toMatchObject({ + status: "fallback_queued", + targetRunId: secondRunId, + duplicate: false, + }); + releaseFirstWake({ id: randomUUID(), driverKind: "codex" }); + await expect(firstDelivery).resolves.toMatchObject({ + status: "fallback_queued", + targetRunId: secondRunId, + duplicate: true, + }); + + const [delivery] = await db.select().from(issueQuestionResponseDeliveries); + expect(delivery).toMatchObject({ + status: "fallback_queued", + targetRunId: secondRunId, + attemptCount: 2, + }); + const deliveryEvents = await db.select().from(activityLog) + .where(eq(activityLog.action, "issue.question_response_delivered")); + expect(deliveryEvents).toHaveLength(1); + }); + + it.each(DIRECT_ADAPTER_TYPES)( + "keeps %s on the existing wake path without invoking native steering", + async (adapterType) => { + const seeded = await seed({ + adapterType, + runtimeMode: "legacy", + successorStatus: "running", + }); + const fallbackRunId = randomUUID(); + const wakeup = vi.fn().mockImplementation(async () => db.insert(heartbeatRuns).values({ + id: fallbackRunId, + companyId: seeded.companyId, + agentId: seeded.agentId, + invocationSource: "automation", + status: "queued", + runtimeMode: "legacy", + contextSnapshot: { issueId: seeded.issueId }, + }).returning().then((rows) => rows[0]!)); + const steer = vi.fn(); + + const outcome = await questionResponseDeliveryService(db, { + heartbeat: { wakeup } as never, + steer, + }).deliver(seeded.interaction.id); + + expect(outcome).toMatchObject({ + status: "fallback_queued", + mode: "wake_fallback", + targetRunId: fallbackRunId, + }); + expect(steer).not.toHaveBeenCalled(); + expect(wakeup).toHaveBeenCalledTimes(1); + expect(wakeup.mock.calls[0]?.[1]).toMatchObject({ + idempotencyKey: `interaction:${seeded.interaction.id}:answered`, + contextSnapshot: { + issueId: seeded.issueId, + interactionId: seeded.interaction.id, + interactionStatus: "answered", + }, + }); + }, + ); + it("falls back once when successor steering is unsupported", async () => { const seeded = await seed({ successorStatus: "running" }); const fallbackRunId = randomUUID(); diff --git a/server/src/services/question-response-delivery.ts b/server/src/services/question-response-delivery.ts index fac667cb171..b7911e42188 100644 --- a/server/src/services/question-response-delivery.ts +++ b/server/src/services/question-response-delivery.ts @@ -29,6 +29,13 @@ const DELIVERY_CLAIM_STALE_MS = 30_000; const DELIVERY_CLAIM_REFRESH_MS = 10_000; const MAX_DELIVERY_ATTEMPTS = 5; const DELIVERY_CORRELATION_PREFIX = "question-response:"; + +class DeliveryClaimUnavailableError extends Error { + constructor() { + super("question_response_delivery_claim_unavailable"); + this.name = "DeliveryClaimUnavailableError"; + } +} const DURABLE_WAKE_REQUEST_STATUSES = [ "queued", "running", @@ -306,6 +313,7 @@ export function questionResponseDeliveryService( }).where(and( eq(issueQuestionResponseDeliveries.id, input.delivery.id), eq(issueQuestionResponseDeliveries.status, "delivering"), + eq(issueQuestionResponseDeliveries.attemptCount, input.delivery.attemptCount), )).returning().then((rows) => rows[0] ?? null); if (!row) return null; await logActivity(tx as unknown as Db, { @@ -368,6 +376,7 @@ export function questionResponseDeliveryService( }).where(and( eq(issueQuestionResponseDeliveries.id, delivery.id), eq(issueQuestionResponseDeliveries.status, "delivering"), + eq(issueQuestionResponseDeliveries.attemptCount, delivery.attemptCount), )); } return exhausted; @@ -386,6 +395,7 @@ export function questionResponseDeliveryService( }).where(and( eq(issueQuestionResponseDeliveries.id, delivery.id), eq(issueQuestionResponseDeliveries.status, "delivering"), + eq(issueQuestionResponseDeliveries.attemptCount, delivery.attemptCount), )).returning({ id: issueQuestionResponseDeliveries.id }); if (renewed.length === 0) stopped = true; }).catch((error) => { @@ -393,13 +403,41 @@ export function questionResponseDeliveryService( }); }, claimRefreshMs); timer.unref?.(); + let result: T | undefined; + let operationError: unknown; try { - return await operation(); + result = await operation(); + } catch (error) { + operationError = error; } finally { stopped = true; clearInterval(timer); await renewal; } + + // `attemptCount` is the claim generation. A stale worker must not continue + // after a sweep has reclaimed the row for a newer attempt, even if its + // external side effect eventually resolves. Confirm ownership after the + // side effect so a rejected native steer cannot fall through to a second + // wake after losing its claim. + let ownsClaim = false; + try { + ownsClaim = await db.select({ id: issueQuestionResponseDeliveries.id }) + .from(issueQuestionResponseDeliveries) + .where(and( + eq(issueQuestionResponseDeliveries.id, delivery.id), + eq(issueQuestionResponseDeliveries.status, "delivering"), + eq(issueQuestionResponseDeliveries.attemptCount, delivery.attemptCount), + )) + .limit(1) + .then((rows) => rows.length === 1); + } catch (error) { + logger.warn({ err: error, deliveryId: delivery.id }, "question response claim ownership check failed"); + throw new DeliveryClaimUnavailableError(); + } + if (!ownsClaim) throw new DeliveryClaimUnavailableError(); + if (operationError !== undefined) throw operationError; + return result as T; } async function findDurableWakeRequest(input: { @@ -537,6 +575,7 @@ export function questionResponseDeliveryService( adapter: successorRunning.driverKind ?? adapter, }); } catch (error) { + if (error instanceof DeliveryClaimUnavailableError) return terminalOutcome(interactionId); steeringErrorCode = error instanceof NativeSessionSteeringError ? error.code : "steering_rejected"; @@ -615,6 +654,7 @@ export function questionResponseDeliveryService( errorCode: steeringErrorCode, }); } catch (error) { + if (error instanceof DeliveryClaimUnavailableError) return terminalOutcome(interactionId); const errorCode = error instanceof Error && compactLine(error.message) ? compactLine(error.message)!.slice(0, 160) : "question_response_wake_failed";