diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index 5ad76e359c..a55ebffe67 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -583,6 +583,17 @@ const settleAbortFence = async (user, streamId) => { * enrolling them in the job manager so deletion can fence them like everything else. */ const quiesceInteractiveGenerations = async (user) => { + // MARKERS BEFORE THE ACTIVE SCAN — the read order is what makes the admission + // lease -> durable job handoff atomic from this side. Writers hold the lease + // strictly until the job is active-set visible, so with leases read FIRST, any + // interleaving either shows the lease here or the job below; jobs-first allowed + // a request to create its job after this scan and release its lease before the + // count, hiding both. The count itself gates settlement further down; this early + // read only has to ORDER before the job scan. + const earlyPendingOwnerWork = await GenerationJobManager.countUserFinalizations( + user.id, + user.tenantId, + ); const activeJobIds = await GenerationJobManager.getActiveJobIdsForUser(user.id, user.tenantId); for (const streamId of activeJobIds ?? []) { // FENCE BEFORE THE CAS. An abort transitions the shared job to terminal, which @@ -638,10 +649,10 @@ const quiesceInteractiveGenerations = async (user) => { // and durable (TTL-free, so an outage can never read as settlement) abort fences on // the user document. Multi-replica deployments require the Redis job store (see // isTopologySafeToArm), where the owner markers are cross-replica visible. - const pendingOwnerWork = await GenerationJobManager.countUserFinalizations( - user.id, - user.tenantId, - ); + const pendingOwnerWork = + earlyPendingOwnerWork > 0 + ? earlyPendingOwnerWork + : await GenerationJobManager.countUserFinalizations(user.id, user.tenantId); if (pendingOwnerWork > 0) { logger.info( `[quiesceInteractiveGenerations] ${pendingOwnerWork} owner finalization(s) still landing for ${user.id}`, diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 6b1897a5f0..382f6467fe 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -542,6 +542,36 @@ describe('deleteUserController - interactive generation quiesce', () => { expect(res.status).toHaveBeenCalledWith(503); }); + it('reads finalization leases BEFORE the active-job scan', async () => { + // The admission-lease -> durable-job handoff is only atomic against a quiesce + // that reads in the OPPOSITE order of the writer: writers hold the lease until + // the job is active-set visible, so leases-first shows every interleaving either + // the lease or the job. Jobs-first allowed a request to create its job after the + // scan and release its lease before the count — hiding both, cascading, and then + // letting the new generation persist into a deleted account. + const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; + const res = createRes(); + mockGetUserById.mockResolvedValue({ _id: 'user1', twoFactorEnabled: false }); + + await deleteUserController(req, res); + + expect(GenerationJobManager.countUserFinalizations.mock.invocationCallOrder[0]).toBeLessThan( + GenerationJobManager.getActiveJobIdsForUser.mock.invocationCallOrder[0], + ); + }); + + it('defers on an admission lease observed before the job scan', async () => { + const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; + const res = createRes(); + mockGetUserById.mockResolvedValue({ _id: 'user1', twoFactorEnabled: false }); + GenerationJobManager.countUserFinalizations.mockResolvedValueOnce(1); + + await deleteUserController(req, res); + + expect(res.status).toHaveBeenCalledWith(503); + expect(mockDeleteUserById).not.toHaveBeenCalled(); + }); + it('keeps a fence whose complete terminal is still persisting', async () => { const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; const res = createRes(); diff --git a/api/server/controllers/agents/__tests__/jobReplacement.spec.js b/api/server/controllers/agents/__tests__/jobReplacement.spec.js index a9de913fc8..af1cedff20 100644 --- a/api/server/controllers/agents/__tests__/jobReplacement.spec.js +++ b/api/server/controllers/agents/__tests__/jobReplacement.spec.js @@ -19,6 +19,7 @@ const mockGenerationJobManager = { // CAS whenever post-terminal title work is possible, so the facade mock must // mirror it or the turn stalls on an undefined call. registerUserFinalization: jest.fn(async () => undefined), + holdUserFinalization: jest.fn(async () => ({ leaseId: 'held-lease', release: jest.fn() })), clearUserFinalization: jest.fn(async () => undefined), countUserFinalizations: jest.fn(async () => 0), createJob: jest.fn(), diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index d8d5a6b0ab..e7f4b22a4f 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -26,6 +26,7 @@ const mockGenerationJobManager = { // CAS whenever post-terminal title work is possible, so the facade mock must // mirror it or the turn stalls on an undefined call. registerUserFinalization: jest.fn(async () => undefined), + holdUserFinalization: jest.fn(async () => ({ leaseId: 'held-lease', release: jest.fn() })), clearUserFinalization: jest.fn(async () => undefined), countUserFinalizations: jest.fn(async () => 0), createJob: jest.fn(), diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index 810ccc3ae0..5ec5c843d2 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -12,6 +12,10 @@ const mockGenerationJobManager = { // CAS whenever post-terminal title work is possible, so the facade mock must // mirror it or the turn stalls on an undefined call. registerUserFinalization: jest.fn(async () => undefined), + holdUserFinalization: jest.fn(async (u, s2, t, g, leaseId) => ({ + leaseId, + release: mockLeaseRelease, + })), clearUserFinalization: jest.fn(async () => undefined), countUserFinalizations: jest.fn(async () => 0), createJob: jest.fn(), @@ -137,6 +141,7 @@ jest.mock('@librechat/data-schemas', () => ({ logger: mockLogger, })); +const mockLeaseRelease = jest.fn(); jest.mock('@librechat/api', () => ({ sendEvent: jest.fn(), toPendingSteer: jest.fn((item) => item), @@ -267,6 +272,9 @@ describe('ResumableAgentController resume metadata', () => { jest.clearAllMocks(); mockGenerationJobManager.registerUserFinalization.mockResolvedValue(undefined); mockGenerationJobManager.clearUserFinalization.mockResolvedValue(undefined); + mockGenerationJobManager.holdUserFinalization.mockImplementation( + async (u, s2, t, g, leaseId) => ({ leaseId, release: mockLeaseRelease }), + ); mockMCPContexts = new WeakMap(); mockCheckAndIncrementPendingRequest.mockResolvedValue({ allowed: true }); mockDecrementPendingRequest.mockResolvedValue(undefined); @@ -3051,26 +3059,26 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); await nextTick(); - // Generation- and lease-qualified: a predecessor's (or a losing same-generation - // contender's) late clear must not drop this marker. The FIRST registration is - // the admission lease covering the auth-to-createJob window. - expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( + // Generation- and lease-qualified HOLDS (heartbeat-renewed): a predecessor's (or + // a losing same-generation contender's) late clear must not drop this marker. + // The FIRST hold is the admission lease covering the auth-to-createJob window. + expect(mockGenerationJobManager.holdUserFinalization).toHaveBeenCalledWith( 'user-123', 'conversation-123', undefined, undefined, expect.stringMatching(/^admission-/), ); - expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( + expect(mockGenerationJobManager.holdUserFinalization).toHaveBeenCalledWith( 'user-123', 'conversation-123', undefined, 1000, expect.any(String), ); - expect( - mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], - ).toBeLessThan(mockGenerationJobManager.claimTerminalJob.mock.invocationCallOrder[0]); + expect(mockGenerationJobManager.holdUserFinalization.mock.invocationCallOrder[0]).toBeLessThan( + mockGenerationJobManager.claimTerminalJob.mock.invocationCallOrder[0], + ); }); it('refuses admission when the deletion barrier is up, lease registered FIRST', async () => { @@ -3098,22 +3106,16 @@ describe('ResumableAgentController resume metadata', () => { expect(res.status).toHaveBeenCalledWith(403); expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); expect(initializeClient).not.toHaveBeenCalled(); - const admissionRegister = mockGenerationJobManager.registerUserFinalization.mock.calls.find( + const admissionRegister = mockGenerationJobManager.holdUserFinalization.mock.calls.find( (call) => String(call[4]).startsWith('admission-'), ); expect(admissionRegister).toBeDefined(); // Lease before barrier read. - expect( - mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], - ).toBeLessThan(models.isUserDeleting.mock.invocationCallOrder[0]); - // The refusal releases its own lease. - expect(mockGenerationJobManager.clearUserFinalization).toHaveBeenCalledWith( - 'user-123', - 'conversation-123', - undefined, - undefined, - admissionRegister[4], + expect(mockGenerationJobManager.holdUserFinalization.mock.invocationCallOrder[0]).toBeLessThan( + models.isUserDeleting.mock.invocationCallOrder[0], ); + // The refusal releases its own hold. + expect(mockLeaseRelease).toHaveBeenCalled(); }); it('releases the admission lease once createJob is durable', async () => { @@ -3146,15 +3148,12 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); await nextTick(); - const admissionRegister = mockGenerationJobManager.registerUserFinalization.mock.calls.find( + const admissionRegister = mockGenerationJobManager.holdUserFinalization.mock.calls.find( (call) => String(call[4]).startsWith('admission-'), ); expect(admissionRegister).toBeDefined(); - const admissionClear = mockGenerationJobManager.clearUserFinalization.mock.calls.find( - (call) => call[4] === admissionRegister[4], - ); // Held exactly until the job is durable (active-set visible), then released. - expect(admissionClear).toBeDefined(); + expect(mockLeaseRelease).toHaveBeenCalled(); }); it('fails CLOSED when the finalization marker cannot be registered: no terminal CAS at all', async () => { @@ -3162,9 +3161,9 @@ describe('ResumableAgentController resume metadata', () => { // fences: neither the complete-claim NOR the error-path completeJob may run // unfenced. The job stays active — deletion-visible by itself — for the stale // reaper to recover. - mockGenerationJobManager.registerUserFinalization.mockRejectedValue( - new Error('marker store down'), - ); + mockGenerationJobManager.holdUserFinalization + .mockResolvedValueOnce({ leaseId: 'admission-ok', release: mockLeaseRelease }) + .mockRejectedValue(new Error('marker store down')); mockGenerationJobManager.getJob.mockResolvedValue({ createdAt: 1000, status: 'running' }); const client = { options: {}, diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index 759696f1fa..47b7a6c724 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -42,6 +42,7 @@ const mockLogger = { info: jest.fn(), }; +const mockLeaseRelease = jest.fn(); const mockJobStore = { getJob: jest.fn(), updateJob: jest.fn(), @@ -52,6 +53,10 @@ const mockGenerationJobManager = { // CAS whenever post-terminal title work is possible, so the facade mock must // mirror it or the turn stalls on an undefined call. registerUserFinalization: jest.fn(async () => undefined), + holdUserFinalization: jest.fn(async (u, s2, t, g, leaseId) => ({ + leaseId, + release: mockLeaseRelease, + })), clearUserFinalization: jest.fn(async () => undefined), countUserFinalizations: jest.fn(async () => 0), getJob: jest.fn(), @@ -249,6 +254,9 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { mockJobStore.updateJob.mockResolvedValue(undefined); mockGenerationJobManager.registerUserFinalization.mockResolvedValue(undefined); mockGenerationJobManager.clearUserFinalization.mockResolvedValue(undefined); + mockGenerationJobManager.holdUserFinalization.mockImplementation( + async (u, s2, t, g, leaseId) => ({ leaseId, release: mockLeaseRelease }), + ); mockGenerationJobManager.getResumeState.mockResolvedValue({ aggregatedContent: [] }); mockGenerationJobManager.emitDone.mockResolvedValue(undefined); mockGenerationJobManager.emitError.mockResolvedValue(undefined); @@ -2196,35 +2204,39 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { // first-turn title) are still ahead — without the marker a deletion quiesce in // that window sees neither an active job nor a fence. // Generation-qualified: a predecessor's late clear must not drop this marker. - // Lease-qualified: a losing same-generation contender (a racing Stop) must - // only ever release its own lease, never this one. - expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( + // Lease-qualified HOLD (heartbeat-renewed): a losing same-generation contender + // (a racing Stop) must only ever release its own lease, never this one. + expect(mockGenerationJobManager.holdUserFinalization).toHaveBeenCalledWith( USER_ID, CONVO_ID, TENANT_ID, 1000, expect.any(String), ); - const resumeLeaseId = mockGenerationJobManager.registerUserFinalization.mock.calls[0][4]; expect( - mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], + mockGenerationJobManager.holdUserFinalization.mock.invocationCallOrder[0], ).toBeLessThan(mockGenerationJobManager.claimTerminalJob.mock.invocationCallOrder[0]); + // Released only after the persistence phase (the finish of the terminal claim), + // never between the CAS and the save — and the owner-lifecycle lease releases + // with it (see USER_FINALIZATION_OWNER_LEASE). + expect(mockLeaseRelease).toHaveBeenCalled(); + expect(mockLeaseRelease.mock.invocationCallOrder[0]).toBeGreaterThan( + mockSaveMessage.mock.invocationCallOrder[0], + ); expect(mockGenerationJobManager.clearUserFinalization).toHaveBeenCalledWith( USER_ID, CONVO_ID, TENANT_ID, 1000, - resumeLeaseId, + 'owner', ); - // Cleared only after the persistence phase (the finish of the terminal claim), - // never between the CAS and the save. - expect( - mockGenerationJobManager.clearUserFinalization.mock.invocationCallOrder[0], - ).toBeGreaterThan(mockSaveMessage.mock.invocationCallOrder[0]); }); it('fails CLOSED when the finalization marker cannot be registered on resume', async () => { mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); + mockGenerationJobManager.holdUserFinalization.mockRejectedValue( + new Error('marker store down'), + ); mockGenerationJobManager.registerUserFinalization.mockRejectedValue( new Error('marker store down'), ); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 70ebe2b349..db00bbdb9b 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -925,19 +925,14 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // deletion-visible. FAIL CLOSED on an unregistrable lease: nothing has been // admitted into generation yet, so refusing is free and plainly retryable. const admissionLeaseId = `admission-${crypto.randomUUID()}`; - let admissionLeaseHeld = false; + let admissionHold = null; const releaseAdmissionLease = () => { - if (!admissionLeaseHeld) { + if (admissionHold == null) { return; } - admissionLeaseHeld = false; - void GenerationJobManager.clearUserFinalization( - userId, - streamId, - req.user?.tenantId, - undefined, - admissionLeaseId, - ).catch(() => undefined); + const hold = admissionHold; + admissionHold = null; + hold.release(); }; const refuseAdmission = async (status, body) => { releaseAdmissionLease(); @@ -955,22 +950,17 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit startupTelemetry?.end('rejected'); return sendGenerationJson(res, status, body, generationProtocolVersion); }; - admissionLeaseHeld = await GenerationJobManager.registerUserFinalization( + admissionHold = await GenerationJobManager.holdUserFinalization( userId, streamId, req.user?.tenantId, undefined, admissionLeaseId, - ) - .then(() => true) - .catch((err) => { - logger.warn( - `[ResumableAgentController] Admission lease registration failed: ${streamId}`, - err, - ); - return false; - }); - if (!admissionLeaseHeld) { + ).catch((err) => { + logger.warn(`[ResumableAgentController] Admission lease registration failed: ${streamId}`, err); + return null; + }); + if (admissionHold == null) { res.set('Retry-After', '2'); return refuseAdmission(503, { error: 'Could not fence the request against account deletion. Please retry.', @@ -1434,47 +1424,54 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit * CAS — registering after it leaves a window where a deletion quiesce sees neither * an active job nor a fence and cascades mid-title — and cleared once the title * work settles (or immediately when no post-terminal title runs). */ - let userFinalizationRegistered = false; + let userFinalizationHold = null; /** Unique per contender: completion and Stop race the same generation, and a - * losing contender's cleanup must only ever release its own lease. */ + * losing contender's cleanup must only ever release its own lease. HELD + * (heartbeat-renewed) rather than one-shot: the store TTL only bounds a crashed + * holder, and live persistence — a stalled save, a long deferred title — must + * never outlive its own fence. */ const finalizationLeaseId = crypto.randomUUID(); const registerUserFinalizationFence = async () => { - if (userFinalizationRegistered) { + if (userFinalizationHold != null) { return true; } // Generation-qualified (jobCreatedAt): a predecessor's late clear on the same // conversation must never drop THIS generation's marker. One retry absorbs a // transient store blip; a persistent failure is the caller's fail-closed signal. - for (let attempt = 1; attempt <= 2 && !userFinalizationRegistered; attempt++) { - userFinalizationRegistered = await GenerationJobManager.registerUserFinalization( + for (let attempt = 1; attempt <= 2 && userFinalizationHold == null; attempt++) { + userFinalizationHold = await GenerationJobManager.holdUserFinalization( userId, streamId, req.user?.tenantId, jobCreatedAt, finalizationLeaseId, - ) - .then(() => true) - .catch((err) => { - logger.warn( - `[ResumableAgentController] Finalization registration failed (attempt ${attempt}/2): ${streamId}`, - err, - ); - return false; - }); + ).catch((err) => { + logger.warn( + `[ResumableAgentController] Finalization registration failed (attempt ${attempt}/2): ${streamId}`, + err, + ); + return null; + }); } - return userFinalizationRegistered; + return userFinalizationHold != null; }; const clearUserFinalization = () => { - if (!userFinalizationRegistered) { + if (userFinalizationHold == null) { return; } - userFinalizationRegistered = false; + const hold = userFinalizationHold; + userFinalizationHold = null; + hold.release(); + // The OWNER-lifecycle lease is registered out-of-band (at abort-signal time, by + // the manager / transport pre-ACK fence — see USER_FINALIZATION_OWNER_LEASE) and + // this controller, as the generation owner, releases it with its own: by this + // point every owner-side write has landed or failed for good. void GenerationJobManager.clearUserFinalization( userId, streamId, req.user?.tenantId, jobCreatedAt, - finalizationLeaseId, + 'owner', ).catch(() => undefined); }; diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 61dfb491e6..3e47e3bdd5 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -448,37 +448,41 @@ async function finalizeResumedTurn({ // this generation's marker. Throwing routes to the controller catch, whose own // fenced settlement either registers its marker or leaves the job paused/active — // deletion-visible either way. Cleared in the persistence try's finally. - let finalizationRegistered = false; + let finalizationHold = null; /** Unique per contender: this resume and a racing Stop contend on the same - * generation, and a losing contender's cleanup must only release its own lease. */ + * generation, and a losing contender's cleanup must only release its own lease. + * HELD (heartbeat-renewed): live persistence must never outlive its fence. */ const finalizationLeaseId = crypto.randomUUID(); - for (let attempt = 1; attempt <= 2 && !finalizationRegistered; attempt++) { - finalizationRegistered = await GenerationJobManager.registerUserFinalization( + for (let attempt = 1; attempt <= 2 && finalizationHold == null; attempt++) { + finalizationHold = await GenerationJobManager.holdUserFinalization( userId, streamId, req.user?.tenantId, job.createdAt, finalizationLeaseId, - ) - .then(() => true) - .catch((err) => { - logger.warn( - `[ResumeAgentController] Finalization registration failed (attempt ${attempt}/2): ${streamId}`, - err, - ); - return false; - }); + ).catch((err) => { + logger.warn( + `[ResumeAgentController] Finalization registration failed (attempt ${attempt}/2): ${streamId}`, + err, + ); + return null; + }); } - if (!finalizationRegistered) { + if (finalizationHold == null) { throw new Error('Finalization marker unavailable before resumed terminal persistence'); } const clearFinalization = () => { + finalizationHold?.release(); + finalizationHold = null; + // Release the owner-lifecycle lease too: this resumed controller IS the + // generation owner, and by this point its writes have landed or failed for good + // (see USER_FINALIZATION_OWNER_LEASE in the manager). void GenerationJobManager.clearUserFinalization( userId, streamId, req.user?.tenantId, job.createdAt, - finalizationLeaseId, + 'owner', ).catch(() => undefined); }; @@ -1460,6 +1464,15 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) if (wasAborted) { logger.debug(`[ResumeAgentController] Resume aborted; settling the run: ${streamId}`); await settleAbortedScheduledResume(job, streamId, conversationId); + // This resumed controller is the generation owner; its abort-path writes have + // now landed, so the signal-time owner lease releases here. + void GenerationJobManager.clearUserFinalization( + userId, + streamId, + req.user?.tenantId, + job.createdAt, + 'owner', + ).catch(() => undefined); return; } logger.error('[ResumeAgentController] Resume failed', err); diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 7e7a3a98f7..fd39e6914f 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -65,6 +65,17 @@ const PREEMPT_TOMBSTONE_MAX = STEER_QUEUE_MAX_DEPTH * 2; * would wedge the abort path after the terminal CAS already landed. */ const ABORT_PUBLISH_TIMEOUT_MS = 5_000; + +/** Deterministic owner-lifecycle lease token: exactly one generation owner exists per + * (stream, generation), so its lease needs no contender uniqueness — and determinism + * is what lets the signal-time registrant (the abort path / transport ACK hook) and + * the owner's catch-side release agree on the same lease across processes. */ +export const USER_FINALIZATION_OWNER_LEASE = 'owner'; +/** Deterministic retention lease for an abort whose delivery is still unconfirmed; + * renewed by every resignal attempt, cleared when one finally leaves. */ +const USER_FINALIZATION_STOP_LEASE = 'stop'; +/** Heartbeat cadence for held leases (store TTL is five minutes). */ +const USER_FINALIZATION_RENEWAL_MS = 60_000; import { SteeringLifecycle, toPendingSteer, @@ -712,6 +723,7 @@ class GenerationJobManagerClass { this._approvals = this.createApprovalLifecycle(this.jobStore); this._steering = new SteeringLifecycle(this.jobStore); this.eventTransport = options?.eventTransport ?? new InMemoryEventTransport(); + this.installAbortAcknowledgementFence(); this._cleanupOnComplete = options?.cleanupOnComplete ?? true; } @@ -804,6 +816,7 @@ class GenerationJobManagerClass { this._approvals = this.createApprovalLifecycle(this.jobStore); this._steering = new SteeringLifecycle(this.jobStore); this.eventTransport = services.eventTransport; + this.installAbortAcknowledgementFence(); this._isRedis = services.isRedis ?? false; /** Coalescing needs BOTH configured services to actually batch: the flush * capabilities are how implementations advertise it. A custom transport @@ -1611,6 +1624,23 @@ class GenerationJobManagerClass { acknowledged.push(receipt.createdAt); } else { allReceiptsDelivered = false; + if (hadRunningProvider) { + // The atomic replacement already removed this predecessor from active + // storage, and an unconfirmed handoff terminalizes the replacement too — + // leaving NOTHING a deletion quiesce could discover while the + // predecessor's provider may still be generating. Retain its owner lease; + // the owner replica renews it through the pre-ACK fence when the signal + // finally lands, and the TTL bounds a dead owner. + await this.jobStore + .registerUserFinalization?.( + job.userId, + streamId, + job.tenantId, + receipt.createdAt, + USER_FINALIZATION_OWNER_LEASE, + ) + .catch(() => undefined); + } } } @@ -3772,8 +3802,8 @@ class GenerationJobManagerClass { } } if (!delivered && !published) { - // HEARTBEAT the fence: the winning abort's lease is TTL-bounded and never - // renewed on its own, so each retry re-registers a fresh lease while the stop + // HEARTBEAT the fence: the stop lease is TTL-bounded and never renewed on its + // own, so each retry re-registers the SAME deterministic lease while the stop // remains undelivered — the deletion quiesce keeps deferring for as long as // someone is still trying to deliver it. await this.jobStore @@ -3782,7 +3812,19 @@ class GenerationJobManagerClass { streamId, jobData.tenantId, jobData.createdAt, - randomUUID(), + USER_FINALIZATION_STOP_LEASE, + ) + .catch(() => undefined); + } else { + // Delivery finally left (or this replica owns the generation): the retained + // stop lease has served its purpose — reap it instead of waiting out the TTL. + await this.jobStore + .clearUserFinalization?.( + jobData.userId, + streamId, + jobData.tenantId, + jobData.createdAt, + USER_FINALIZATION_STOP_LEASE, ) .catch(() => undefined); } @@ -3995,55 +4037,62 @@ class GenerationJobManagerClass { let transitionFrom = abortableStatus; let drainedSteers: SteerQueueItem[] | null = null; let currentAfterConflict: SerializableJobData | null = jobData; - // An approval decision can move the same epoch between running and - // requires_action after our read. Retry that legal same-generation state - // change instead of telling the caller it stopped a run that is still live. - for (let attempt = 0; attempt < 3 && drainedSteers == null; attempt++) { - const expectedActionId = - transitionFrom === 'requires_action' ? currentAfterConflict?.pendingActionId : undefined; - drainedSteers = await this.jobStore.transitionStatusAndDrainSteers(streamId, { - from: transitionFrom, - to: 'aborted', - expectCreatedAt: jobData.createdAt, - ...(expectedActionId != null && { expectActionId: expectedActionId }), - patch: { - // preserveForReconcile omits completedAt so the finished-job sweep can't - // reap the abort before the schedules reconciler observes it — the store - // holds completedAt-less terminals on the retained-evidence TTL. - ...(options?.preserveForReconcile ? {} : { completedAt: terminalPersistenceStartedAt }), - terminalPersistencePending: true, - terminalPersistenceStartedAt, - }, - }); - if (drainedSteers != null) { - break; + try { + // An approval decision can move the same epoch between running and + // requires_action after our read. Retry that legal same-generation state + // change instead of telling the caller it stopped a run that is still live. + for (let attempt = 0; attempt < 3 && drainedSteers == null; attempt++) { + const expectedActionId = + transitionFrom === 'requires_action' ? currentAfterConflict?.pendingActionId : undefined; + drainedSteers = await this.jobStore.transitionStatusAndDrainSteers(streamId, { + from: transitionFrom, + to: 'aborted', + expectCreatedAt: jobData.createdAt, + ...(expectedActionId != null && { expectActionId: expectedActionId }), + patch: { + // preserveForReconcile omits completedAt so the finished-job sweep can't + // reap the abort before the schedules reconciler observes it — the store + // holds completedAt-less terminals on the retained-evidence TTL. + ...(options?.preserveForReconcile ? {} : { completedAt: terminalPersistenceStartedAt }), + terminalPersistencePending: true, + terminalPersistenceStartedAt, + }, + }); + if (drainedSteers != null) { + break; + } + currentAfterConflict = await this.jobStore.getJob(streamId); + if ( + currentAfterConflict?.createdAt === jobData.createdAt && + currentAfterConflict.status === 'requires_action' && + currentAfterConflict.terminalPersistencePending === true + ) { + currentAfterConflict = await this._approvals.waitForPausePersistence( + streamId, + jobData.createdAt, + ); + } + if ( + currentAfterConflict?.createdAt !== jobData.createdAt || + (currentAfterConflict.status !== 'running' && + currentAfterConflict.status !== 'requires_action') + ) { + break; + } + if ( + currentAfterConflict.status === transitionFrom && + (transitionFrom !== 'requires_action' || + currentAfterConflict.pendingActionId === expectedActionId) + ) { + break; + } + transitionFrom = currentAfterConflict.status; } - currentAfterConflict = await this.jobStore.getJob(streamId); - if ( - currentAfterConflict?.createdAt === jobData.createdAt && - currentAfterConflict.status === 'requires_action' && - currentAfterConflict.terminalPersistencePending === true - ) { - currentAfterConflict = await this._approvals.waitForPausePersistence( - streamId, - jobData.createdAt, - ); - } - if ( - currentAfterConflict?.createdAt !== jobData.createdAt || - (currentAfterConflict.status !== 'running' && - currentAfterConflict.status !== 'requires_action') - ) { - break; - } - if ( - currentAfterConflict.status === transitionFrom && - (transitionFrom !== 'requires_action' || - currentAfterConflict.pendingActionId === expectedActionId) - ) { - break; - } - transitionFrom = currentAfterConflict.status; + } catch (transitionError) { + // A THROWN transition (store outage, not a lost CAS) aborted nothing: release + // the contender lease instead of leaking it to its TTL. + clearAbortFinalization(); + throw transitionError; } if (drainedSteers == null) { // Lost the CAS: this call will persist nothing, so its fence releases now. @@ -4122,6 +4171,22 @@ class GenerationJobManagerClass { if (runtime) { this.releaseAbortSubscription(runtime); } + // OWNER-LIFECYCLE LEASE before the trip, mirroring the transport's pre-ACK + // fence: this abort's own lease clears when the call returns, but the owned + // generation's catch persists asynchronously after — the deterministic owner + // lease bridges that gap and the owner's catch releases it once its writes + // land. Best-effort: the contender lease still covers this call's window. + if (abortSignalDelivered && runtime?.createdAt === jobData.createdAt) { + await this.jobStore + .registerUserFinalization?.( + jobData.userId, + streamId, + jobData.tenantId, + jobData.createdAt, + USER_FINALIZATION_OWNER_LEASE, + ) + .catch(() => undefined); + } runtime?.abortController.abort(); if (this.eventTransport.emitAbort) { diff --git a/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts b/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts index 4e2e9e5c87..f2a3ac74d3 100644 --- a/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts +++ b/packages/api/src/stream/__tests__/RedisEventTransport.spec.ts @@ -157,6 +157,53 @@ describe('RedisEventTransport', () => { transport.destroy(); }); + it('publishes the abort acknowledgement only AFTER the pre-ACK fence resolves', async () => { + const mockPublisher = createMockPublisher(); + const mockSubscriber = createMockSubscriber(); + const transport = new RedisEventTransport( + mockPublisher as unknown as Redis, + mockSubscriber as unknown as Redis, + ); + const streamId = 'abort-ack-fence'; + // The stopping side releases its own lease the moment this ACK arrives; the + // owner's abort-catch persistence is still ahead. If the ACK can outrun the + // owner-lifecycle lease registration, a deletion quiesce lands in the gap. + let releaseFence!: () => void; + const fencePending = new Promise((resolve) => { + releaseFence = resolve; + }); + const fence = jest.fn(() => fencePending); + transport.beforeAbortAcknowledged = fence; + await transport.onAbort(streamId, () => true); + const messageHandler = getMessageHandler(mockSubscriber); + + messageHandler( + `stream:{${streamId}}:events`, + JSON.stringify({ type: 'abort', generationId: 7, abortRequestId: 'req-1' }), + ); + await Promise.resolve(); + await Promise.resolve(); + + expect(fence).toHaveBeenCalledWith(streamId, 7); + // Neither the durable ack proof nor the ack publication may precede the fence. + expect(mockPublisher.set).not.toHaveBeenCalled(); + const ackPublishes = mockPublisher.publish.mock.calls.filter(([, payload]) => + String(payload).includes('abort_ack'), + ); + expect(ackPublishes).toHaveLength(0); + + releaseFence(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(mockPublisher.set).toHaveBeenCalledWith( + expect.stringContaining('abort-ack'), + '1', + 'EX', + expect.any(Number), + ); + transport.destroy(); + }); + it('ignores sequenced events while only internal abort listeners are attached', async () => { jest.useFakeTimers(); const mockPublisher = createMockPublisher(); diff --git a/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts b/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts index d9f9837373..afb0476d30 100644 --- a/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts +++ b/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts @@ -145,9 +145,13 @@ describe('abort signal honesty across replicas', () => { expect(register.mock.invocationCallOrder[0]).toBeLessThan( transition.mock.invocationCallOrder[0], ); - // Every write this abort owns has landed by return: the marker is released. + // An OWNED abort additionally bridges to the deterministic owner-lifecycle lease + // before tripping the local generation: the abort call's own contender lease is + // released on return, while the owner lease survives for the generation's catch + // (its asynchronous persistence) to release. + expect(register).toHaveBeenCalledWith('user-1', 'conv-6', undefined, createdAt, 'owner'); await new Promise((resolve) => setImmediate(resolve)); - await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(0); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); }); it('releases the abort marker when the terminal CAS is lost', async () => { @@ -253,6 +257,81 @@ describe('abort signal honesty across replicas', () => { await winner; }); + it('clears the retained stop lease once a resignal finally leaves', async () => { + const { owner, peer, peerTransport, jobStore } = await makeManagers(); + await owner.createJob('conv-11', 'user-1', 'conv-11'); + Object.assign(peerTransport, { + emitAbort: jest.fn((): void => undefined), + emitAbortConfirmed: jest.fn(async () => false), + }); + await peer.abortJob('conv-11'); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); + // A failed resignal renews the fence… + await peer.resignalAbort('conv-11'); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); + + // …and the first one that leaves reaps it instead of waiting out the TTL. + Object.assign(peerTransport, { emitAbortConfirmed: jest.fn(async () => true) }); + const resignal = await peer.resignalAbort('conv-11'); + + expect(resignal.published).toBe(true); + await new Promise((resolve) => setImmediate(resolve)); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(0); + }); + + it('retains the predecessor owner lease when a replacement handoff fails', async () => { + const { GenerationJobManagerClass } = await import('../GenerationJobManager'); + const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore'); + const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport'); + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 0 }); + const owner = new GenerationJobManagerClass(); + owner.configure({ jobStore, eventTransport: new InMemoryEventTransport(), isRedis: false }); + owner.initialize(); + // Redis mode is what requires the durable owner proof for a replacement handoff; + // this transport provides none, so the acknowledgement provably fails. + const peer = new GenerationJobManagerClass(); + const peerTransport = new InMemoryEventTransport(); + Object.assign(peerTransport, { emitAbortConfirmed: jest.fn(async () => false) }); + peer.configure({ jobStore, eventTransport: peerTransport, isRedis: true }); + peer.initialize(); + + await owner.createJob('conv-12', 'user-1', 'conv-12'); + const predecessor = await jobStore.getJob('conv-12'); + expect(predecessor?.status).toBe('running'); + + // The peer replaces the stream: the atomic create removes the predecessor from + // active storage, but its owner never acknowledges the handoff — the replacement + // is terminalized and createJob throws. Without the retained lease, a deletion + // quiesce could discover NEITHER generation while the predecessor's provider may + // still be generating. + await expect(peer.createJob('conv-12', 'user-1', 'conv-12')).rejects.toThrow( + 'predecessor handoff', + ); + + await new Promise((resolve) => setImmediate(resolve)); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBeGreaterThanOrEqual(1); + }); + + it('heartbeats a held lease past the store TTL until released', async () => { + const { owner, jobStore } = await makeManagers(); + jest.useFakeTimers({ now: Date.now() }); + try { + const hold = await owner.holdUserFinalization('user-1', 'conv-13', undefined, 1000, 'held'); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); + + // The store TTL is five minutes and only bounds a crashed holder: live + // persistence (a stalled save, a long title) must never outlive its fence. + await jest.advanceTimersByTimeAsync(6 * 60 * 1000); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); + + hold.release(); + await jest.advanceTimersByTimeAsync(1); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(0); + } finally { + jest.useRealTimers(); + } + }); + it('reports republication as succeeded when the owner acknowledges', async () => { const { owner, peer, peerTransport, jobStore } = await makeManagers(); await owner.createJob('conv-4', 'user-1', 'conv-4'); diff --git a/packages/api/src/stream/implementations/RedisEventTransport.ts b/packages/api/src/stream/implementations/RedisEventTransport.ts index c03d966eef..38aeb4163f 100644 --- a/packages/api/src/stream/implementations/RedisEventTransport.ts +++ b/packages/api/src/stream/implementations/RedisEventTransport.ts @@ -258,6 +258,8 @@ interface StreamSubscribers { * ``` */ export class RedisEventTransport implements IEventTransport { + /** Awaited before an owned abort is acknowledged; see IEventTransport. */ + beforeAbortAcknowledged?: (streamId: string, generationId: number) => Promise; /** Redis client for publishing events */ private publisher: Redis | Cluster; /** Redis client for subscribing to events (separate connection required) */ @@ -941,11 +943,24 @@ export class RedisEventTransport implements IEventTransport { } } if (ownerAcknowledged && message.abortRequestId != null && message.generationId != null) { - void this.publishAbortAcknowledgement( - streamId, - message.generationId, - message.abortRequestId, - ); + const generationId = message.generationId; + const abortRequestId = message.abortRequestId; + void (async () => { + // The owner-lifecycle lease must be durable BEFORE the acknowledgement: + // the stopping side releases its own lease the moment this ACK arrives, + // and the owner's abort-catch persistence is still ahead. Best-effort — + // an unregistrable lease must not suppress the ACK forever (the stopping + // side's retention lease and resignal heartbeat cover that outage). + try { + await this.beforeAbortAcknowledged?.(streamId, generationId); + } catch (error) { + logger.error( + `[RedisEventTransport] Failed to fence abort acknowledgement for ${streamId}:`, + error, + ); + } + await this.publishAbortAcknowledgement(streamId, generationId, abortRequestId); + })(); } } diff --git a/packages/api/src/stream/index.ts b/packages/api/src/stream/index.ts index febb20e08a..bb48d3b742 100644 --- a/packages/api/src/stream/index.ts +++ b/packages/api/src/stream/index.ts @@ -1,6 +1,7 @@ export { GenerationJobManager, GenerationJobManagerClass, + USER_FINALIZATION_OWNER_LEASE, type CreateGenerationJobOptions, type GenerationJobManagerOptions, type TerminalJobClaim, diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index e5da8b4f79..f16efe5bb3 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -487,6 +487,9 @@ export function createUserMethods( { $addToSet: { deletionAbortFences: streamId } }, { timestamps: false }, ); + // Every user-document mutation invalidates the auth user-doc cache, or a cached + // hit keeps serving the pre-mutation document for the cache TTL. + await invalidateAuthUserDocCache(userId); } /** Clears an acknowledged abort fence (signal left the replica, or the job is gone). */ @@ -497,6 +500,7 @@ export function createUserMethods( { $pull: { deletionAbortFences: streamId } }, { timestamps: false }, ); + await invalidateAuthUserDocCache(userId); } /**