From 047fe6766f61585c339409b37faa5057fa8d95d7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 17 Aug 2026 13:00:37 -0400 Subject: [PATCH] fix: fence scheduled resume handoffs --- .../agents/__tests__/resume.spec.js | 80 ++++++++++++- api/server/controllers/agents/resume.js | 100 +++++++++++++++- api/server/services/Schedules/index.js | 2 + packages/api/src/schedules/engine.spec.ts | 4 +- packages/api/src/schedules/fire.spec.ts | 19 +++ packages/api/src/schedules/fire.ts | 22 ++-- packages/api/src/schedules/service.spec.ts | 57 ++++++++- packages/api/src/schedules/service.ts | 110 ++++++++++++++---- .../src/methods/schedule.methods.spec.ts | 31 +++++ packages/data-schemas/src/methods/schedule.ts | 101 +++++++++++++++- 10 files changed, 482 insertions(+), 44 deletions(-) diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index fd6fa4eeb6..034c53189e 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -87,6 +87,8 @@ const mockRecordScheduleOutcome = jest.fn(); const mockIsScheduleLive = jest.fn(); const mockClaimScheduleResume = jest.fn(); const mockReleaseScheduleResumeClaim = jest.fn(); +const mockFinalizeScheduleResumeClaim = jest.fn(); +const mockReleaseScheduleResumeFence = jest.fn(); jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), @@ -113,6 +115,8 @@ jest.mock('~/server/services/Schedules', () => ({ recordScheduleOutcome: (...args) => mockRecordScheduleOutcome(...args), claimScheduleResume: (...args) => mockClaimScheduleResume(...args), releaseScheduleResumeClaim: (...args) => mockReleaseScheduleResumeClaim(...args), + finalizeScheduleResumeClaim: (...args) => mockFinalizeScheduleResumeClaim(...args), + releaseScheduleResumeFence: (...args) => mockReleaseScheduleResumeFence(...args), isScheduleLive: (...args) => mockIsScheduleLive(...args), })); @@ -269,8 +273,14 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { mockGenerationJobManager.approvals.finishPausePersistence.mockResolvedValue(true); mockRecordScheduleOutcome.mockResolvedValue(true); mockIsScheduleLive.mockResolvedValue(true); - mockClaimScheduleResume.mockResolvedValue({ capacitySlot: 0 }); + mockClaimScheduleResume.mockResolvedValue({ + capacitySlot: 0, + claimToken: 'resume-token', + leaseBy: 'resume:resume-token', + }); mockReleaseScheduleResumeClaim.mockResolvedValue(true); + mockFinalizeScheduleResumeClaim.mockResolvedValue(true); + mockReleaseScheduleResumeFence.mockResolvedValue(undefined); // `decrementPendingRequest` runs in the controller's `finally` on every // post-ACK path, so resolving on it signals the async continuation is done. @@ -393,6 +403,12 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expectedConfigRevision: 4, automatic: true, }); + expect(mockFinalizeScheduleResumeClaim).toHaveBeenCalledWith( + 'schedule-1', + 'resume-token', + 'resume:resume-token', + { expectedConfigRevision: 4, automatic: true }, + ); expect(mockRecordScheduleOutcome).toHaveBeenCalledWith({ scheduleId: 'schedule-1', @@ -429,6 +445,35 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { expect(res.status).toBe(409); expect(mockReleaseScheduleResumeClaim).toHaveBeenCalledWith('schedule-1', scheduledFor, 0); + expect(mockReleaseScheduleResumeFence).toHaveBeenCalledWith( + 'schedule-1', + 'resume:resume-token', + ); + }); + + it('stops before provider execution when an edit wins the final resume handoff', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(makeScheduledJob()); + mockFinalizeScheduleResumeClaim.mockResolvedValue(false); + + const res = await post(approveBody()); + + expect(res.status).toBe(409); + expect(res.body).toMatchObject({ code: 'SCHEDULE_NO_LONGER_ACTIVE' }); + expect(mockGenerationJobManager.approvals.resolve).toHaveBeenCalled(); + expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith(CONVO_ID, { + expectedCreatedAt: 1000, + awaitProviderDrain: true, + }); + expect(mockRecordScheduleOutcome).toHaveBeenCalledWith({ + scheduleId: 'schedule-1', + scheduledFor, + streamId: CONVO_ID, + jobCreatedAt: 1000, + status: 'interrupted', + conversationId: CONVO_ID, + error: 'Schedule was disabled, changed, or deleted before approval', + }); + expect(mockInitializeClient).not.toHaveBeenCalled(); }); it('settles a scheduled continuation stopped during its resumed segment', async () => { @@ -455,6 +500,39 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { }); }); + it('records an empty-preempt resumed segment as interrupted, not successful', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(makeScheduledJob()); + mockInitializeClient.mockResolvedValue({ + client: makeClient({ + run: { + getPreemptStats: () => ({ emptyBoundaries: 1 }), + getHaltReason: () => 'preempt_incomplete', + }, + }), + userMCPAuthMap: {}, + }); + + const res = await post(approveBody()); + expect(res.status).toBe(200); + await settled; + await flush(); + + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ unfinished: true }), + expect.anything(), + ); + expect(mockRecordScheduleOutcome).toHaveBeenCalledWith({ + scheduleId: 'schedule-1', + scheduledFor, + streamId: CONVO_ID, + jobCreatedAt: 1000, + status: 'interrupted', + conversationId: CONVO_ID, + error: 'Scheduled run was interrupted before completion', + }); + }); + it('does not overwrite the terminal winner when failed-resume finalization loses its CAS', async () => { mockGenerationJobManager.getJob.mockResolvedValue(makeScheduledJob()); mockGenerationJobManager.completeJob.mockResolvedValue(false); diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index e561a33d10..e4eb76e295 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -34,6 +34,8 @@ const { recordScheduleOutcome, claimScheduleResume, releaseScheduleResumeClaim, + finalizeScheduleResumeClaim, + releaseScheduleResumeFence, isScheduleLive, } = require('~/server/services/Schedules'); const { @@ -434,8 +436,11 @@ async function finalizeResumedTurn({ scheduledFor: meta.scheduledFor, streamId, jobCreatedAt: job.createdAt, - status: 'success', + status: preemptIncomplete ? 'interrupted' : 'success', conversationId, + ...(preemptIncomplete && { + error: 'Scheduled run was interrupted before completion', + }), }); } @@ -807,13 +812,16 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) // The database's partial unique indexes arbitrate both deployment capacity and a // concurrent active occurrence of the same schedule. let scheduleCapacitySlot; + let scheduleResumeClaimToken; + let scheduleResumeLeaseBy; + const scheduleResumeOptions = { + expectedConfigRevision: job.metadata?.scheduleConfigRevision, + automatic: job.metadata?.scheduleManual !== true, + }; if (scheduleId) { let scheduleClaim; try { - scheduleClaim = await claimScheduleResume(scheduleId, scheduledFor, { - expectedConfigRevision: job.metadata?.scheduleConfigRevision, - automatic: job.metadata?.scheduleManual !== true, - }); + scheduleClaim = await claimScheduleResume(scheduleId, scheduledFor, scheduleResumeOptions); } catch (err) { await decrementPendingRequest(userId); logger.error('[ResumeAgentController] Failed to claim scheduled resume capacity', err); @@ -858,8 +866,21 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) ); } scheduleCapacitySlot = scheduleClaim.capacitySlot; + scheduleResumeClaimToken = scheduleClaim.claimToken; + scheduleResumeLeaseBy = scheduleClaim.leaseBy; } + const releaseScheduleFence = async () => { + if (scheduleId == null || scheduleResumeLeaseBy == null) { + return; + } + try { + await releaseScheduleResumeFence(scheduleId, scheduleResumeLeaseBy); + } catch (releaseError) { + logger.warn('[ResumeAgentController] Failed to release scheduled resume fence', releaseError); + } + }; + /** Release only when the exact generation demonstrably remains paused. If the * approval CAS reply is ambiguous and the job cannot be read, retaining the slot * until reconciliation is the safe direction: releasing it could exceed the cap @@ -911,6 +932,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) } catch (err) { const currentJob = await GenerationJobManager.getJob(streamId).catch(() => null); await rollbackUnconsumedScheduleClaim(currentJob); + await releaseScheduleFence(); await decrementPendingRequest(userId); logger.error('[ResumeAgentController] Failed to claim resume', err); return sendGenerationJson(res, 500, { error: 'Failed to resume' }, generationProtocolVersion); @@ -919,6 +941,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) await decrementPendingRequest(userId); const currentJob = await GenerationJobManager.getJob(streamId).catch(() => null); await rollbackUnconsumedScheduleClaim(currentJob); + await releaseScheduleFence(); if (currentJob != null && currentJob.createdAt !== job.createdAt) { return sendGenerationJson(res, 409, { code: 'RUN_REPLACED' }, generationProtocolVersion); } @@ -930,6 +953,73 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) ); } + // Linearize the consumed approval against the schedule's live config. The schedule + // document fence was acquired only after all async policy reads, and this atomic + // consume checks its token/revision/enabled state immediately after the approval CAS. + // An edit/disable that won first makes this fail; one that lands afterward is ordered + // after the continuation has started. Never begin provider execution on a stale claim. + if (scheduleId) { + let scheduleClaimCurrent = false; + try { + scheduleClaimCurrent = await finalizeScheduleResumeClaim( + scheduleId, + scheduleResumeClaimToken, + scheduleResumeLeaseBy, + scheduleResumeOptions, + ); + } catch (error) { + logger.error('[ResumeAgentController] Failed to finalize scheduled resume fence', error); + await releaseScheduleFence(); + } + if (!scheduleClaimCurrent) { + await decrementPendingRequest(userId); + let stopped = false; + try { + const abortResult = await GenerationJobManager.abortJob(streamId, { + expectedCreatedAt: job.createdAt, + awaitProviderDrain: true, + }); + stopped = abortResult != null && abortResult.failureReason == null; + } catch (error) { + logger.warn('[ResumeAgentController] Failed to stop stale scheduled resume', error); + } + if (!stopped) { + res.set('Retry-After', '1'); + return sendGenerationJson( + res, + 503, + { + code: 'SCHEDULE_STOP_UNCONFIRMED', + error: 'The stale scheduled resume could not be confirmed stopped.', + }, + generationProtocolVersion, + ); + } + await recordScheduleOutcome({ + scheduleId, + scheduledFor, + streamId, + jobCreatedAt: job.createdAt, + status: 'interrupted', + conversationId, + error: 'Schedule was disabled, changed, or deleted before approval', + }); + if (checkpointNamespace !== '') { + await deleteAgentCheckpoint(conversationId, checkpointerCfg, undefined, { + checkpointNamespace, + }).catch((error) => { + logger.warn('[ResumeAgentController] Failed to prune stale schedule checkpoint', error); + }); + } + return sendGenerationJson( + res, + 409, + { code: 'SCHEDULE_NO_LONGER_ACTIVE', error: 'This schedule can no longer be resumed' }, + generationProtocolVersion, + ); + } + } + /** * An interrupt steer enqueued just before the pause survives durably with * its `preempt` flag, but the ARM lived only in the previous owner's diff --git a/api/server/services/Schedules/index.js b/api/server/services/Schedules/index.js index 77456a671b..ba051f34fb 100644 --- a/api/server/services/Schedules/index.js +++ b/api/server/services/Schedules/index.js @@ -76,6 +76,8 @@ module.exports = { recordScheduleOutcome: invoke('recordScheduleOutcome'), claimScheduleResume: invoke('claimScheduleResume'), releaseScheduleResumeClaim: invoke('releaseScheduleResumeClaim'), + finalizeScheduleResumeClaim: invoke('finalizeScheduleResumeClaim'), + releaseScheduleResumeFence: invoke('releaseScheduleResumeFence'), isScheduleLive: invoke('isScheduleLive'), deleteScheduleForOwner: invoke('deleteScheduleForOwner'), quiesceUserSchedules: invoke('quiesceUserSchedules'), diff --git a/packages/api/src/schedules/engine.spec.ts b/packages/api/src/schedules/engine.spec.ts index 572cdedef7..58eec69202 100644 --- a/packages/api/src/schedules/engine.spec.ts +++ b/packages/api/src/schedules/engine.spec.ts @@ -47,10 +47,10 @@ function makeMethods(schedule: FireableSchedule) { claimDueSchedule: jest.fn(async () => (claims++ === 0 ? schedule : null)), countActiveRuns: jest.fn(async () => 0), advanceSchedule: jest.fn( - async (_id: string, _next: Date | null, _from?: Date, _token?: string) => undefined, + async (_id: string, _next: Date | null, _from?: Date, _token?: string) => true, ), disableSchedule: jest.fn(async (_id: string, _reason: string, _token?: string) => undefined), - releaseLease: jest.fn(async () => undefined), + releaseLease: jest.fn(async () => true), holdsLease: jest.fn(async () => true), scheduleExists: jest.fn(async () => true), deleteScheduleRun: jest.fn(async () => undefined), diff --git a/packages/api/src/schedules/fire.spec.ts b/packages/api/src/schedules/fire.spec.ts index 7fb1de58b4..82eaf4e4b4 100644 --- a/packages/api/src/schedules/fire.spec.ts +++ b/packages/api/src/schedules/fire.spec.ts @@ -53,9 +53,11 @@ function makeMethods() { const methods = { releaseLease: jest.fn(async () => { calls.releaseLease += 1; + return true; }), advanceSchedule: jest.fn(async () => { calls.advance += 1; + return true; }), disableSchedule: jest.fn(async (_id: string, reason: string) => { calls.disable.push(reason); @@ -265,6 +267,23 @@ describe('fireSchedule', () => { expect([...runs.values()][0].status).toBe('started'); }); + it('releases the old holder when a post-enqueue owner edit fences the advance', async () => { + const { methods } = makeMethods(); + const enqueueTrigger = jest.fn(async () => undefined); + methods.advanceSchedule.mockResolvedValueOnce(false); + + const result = await fireSchedule( + makeDeps(methods, { enqueueTrigger }), + makeSchedule({ leaseBy: 'claim-holder-1' }), + LIMITS, + dueAt(), + ); + + expect(result.fired).toBe(true); + expect(enqueueTrigger).toHaveBeenCalledTimes(1); + expect(methods.releaseLeaseByHolder).toHaveBeenCalledWith('sched-1', 'claim-holder-1'); + }); + it('marks a Run Now trigger as manual for downstream limiter policy', async () => { const { methods } = makeMethods(); const enqueueTrigger = jest.fn< diff --git a/packages/api/src/schedules/fire.ts b/packages/api/src/schedules/fire.ts index 583d20c61b..c6b8c9d1ca 100644 --- a/packages/api/src/schedules/fire.ts +++ b/packages/api/src/schedules/fire.ts @@ -128,13 +128,6 @@ export async function fireSchedule( // worker (lease expired + re-claimed, or the schedule edited/re-enabled/deleted // — all of which rotate the token) cannot clobber the newer authoritative state. const claimToken = schedule.claimToken; - // A manual run-now must never reschedule the next automatic occurrence; it - // only releases the lease it acquired for serialization. - const advance = options?.manual - ? () => methods.releaseLease(schedule.id, claimToken) - : // Predicate the advance on the claimed occurrence AND the claim token so a - // concurrent owner edit or a lease-expiry re-claim isn't clobbered. - () => methods.advanceSchedule(schedule.id, nextRunAt, scheduledFor, claimToken); /** * Rolls back the `started` row this fire reserved. Fenced on the conversation id it @@ -163,6 +156,21 @@ export async function fireSchedule( } }; + // A manual run-now must never reschedule the next automatic occurrence; it only + // releases its serialization lease. Both writes are token-fenced, so an owner edit + // racing after durable enqueue can make them miss while deliberately preserving the + // old holder fields. Release by the unique OLD holder on a miss: an edit is unwedged + // immediately, while a takeover changed leaseBy and remains untouched. + const advance = async () => { + const advanced = options?.manual + ? await methods.releaseLease(schedule.id, claimToken) + : await methods.advanceSchedule(schedule.id, nextRunAt, scheduledFor, claimToken); + if (!advanced) { + await releaseSupersededLease(); + } + return advanced; + }; + /** * Steps aside after a failed revalidation WITHOUT advancing. `advance()` here is a * designed no-op in every true supersession (takeover, owner edit — both rotate the diff --git a/packages/api/src/schedules/service.spec.ts b/packages/api/src/schedules/service.spec.ts index 5a26e92a47..b4118a58d1 100644 --- a/packages/api/src/schedules/service.spec.ts +++ b/packages/api/src/schedules/service.spec.ts @@ -911,6 +911,13 @@ describe('scheduled resume capacity', () => { })), getRoleByName: jest.fn(async () => ({ permissions: { SCHEDULES: { USE: true } } })), getCapacityOccupancy: jest.fn(async () => occupancy), + acquireResumeLease: jest.fn(async () => ({ + id: 's1', + claimToken: 'resume-token', + leaseBy: 'resume:resume-token', + })), + consumeResumeLease: jest.fn(async () => true), + releaseLeaseByHolder: jest.fn(async () => undefined), markRunResumeClaimed: jest.fn(async (_id, _scheduledFor, capacitySlot) => ({ capacitySlot, })), @@ -945,8 +952,24 @@ describe('scheduled resume capacity', () => { await expect(service.claimScheduleResume('s1', scheduledFor)).resolves.toEqual({ capacitySlot: 0, + claimToken: 'resume-token', + leaseBy: 'resume:resume-token', }); + expect(methods.acquireResumeLease).toHaveBeenCalledWith('s1', undefined, true, 60_000); expect(methods.markRunResumeClaimed).toHaveBeenCalledWith('s1', new Date(scheduledFor), 0); + await expect( + service.finalizeScheduleResumeClaim('s1', 'resume-token', 'resume:resume-token', { + expectedConfigRevision: 3, + automatic: true, + }), + ).resolves.toBe(true); + expect(methods.consumeResumeLease).toHaveBeenCalledWith( + 's1', + 'resume-token', + 'resume:resume-token', + true, + 3, + ); await expect(service.releaseScheduleResumeClaim('s1', scheduledFor, 0)).resolves.toBe(true); expect(methods.releaseRunResumeClaim).toHaveBeenCalledWith('s1', new Date(scheduledFor), 0); }); @@ -958,6 +981,7 @@ describe('scheduled resume capacity', () => { conflict: 'capacity', }); expect(methods.markRunResumeClaimed).not.toHaveBeenCalled(); + expect(methods.releaseLeaseByHolder).toHaveBeenCalledWith('s1', 'resume:resume-token'); }); it('normalizes an exhausted slot-collision retry to public capacity conflict', async () => { @@ -981,6 +1005,33 @@ describe('scheduled resume capacity', () => { expect(methods.markRunResumeClaimed).not.toHaveBeenCalled(); }); + it('refuses an edit that wins during async policy checks at the final schedule CAS', async () => { + const { service, methods } = makeResumeService({ takenSlots: [], unslotted: 0 }); + methods.acquireResumeLease.mockResolvedValue(null as never); + + await expect( + service.claimScheduleResume('s1', '2026-08-17T12:00:00.000Z', { + expectedConfigRevision: 3, + automatic: true, + }), + ).resolves.toEqual({ conflict: 'inactive' }); + expect(methods.acquireResumeLease).toHaveBeenCalledWith('s1', 3, true, 60_000); + expect(methods.markRunResumeClaimed).not.toHaveBeenCalled(); + }); + + it('releases the old holder when the post-approval schedule handoff loses', async () => { + const { service, methods } = makeResumeService({ takenSlots: [], unslotted: 0 }); + methods.consumeResumeLease.mockResolvedValue(false); + + await expect( + service.finalizeScheduleResumeClaim('s1', 'resume-token', 'resume:resume-token', { + expectedConfigRevision: 3, + automatic: true, + }), + ).resolves.toBe(false); + expect(methods.releaseLeaseByHolder).toHaveBeenCalledWith('s1', 'resume:resume-token'); + }); + it('refuses a disabled automatic resume but preserves explicit Run Now semantics', async () => { const { service, methods } = makeResumeService({ takenSlots: [], unslotted: 0 }); methods.getScheduleById.mockResolvedValue({ @@ -1001,7 +1052,11 @@ describe('scheduled resume capacity', () => { expectedConfigRevision: 3, automatic: false, }), - ).resolves.toEqual({ capacitySlot: 0 }); + ).resolves.toEqual({ + capacitySlot: 0, + claimToken: 'resume-token', + leaseBy: 'resume:resume-token', + }); }); it('refuses a resume when the owner lost schedule access before the capacity handoff', async () => { diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index a7ae234fc0..4993b96797 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -50,6 +50,9 @@ const QUIESCE_SETTLE_ERRORS: Partial> = error: 'Run ended in error', }; +/** Short schedule-document fence spanning capacity admission -> approval CAS. */ +const SCHEDULE_RESUME_LEASE_MS = 60_000; + /** * Whether this process may arm the scheduler at all. * @@ -93,7 +96,7 @@ export interface RecordScheduleOutcomeInput { } export type ScheduleResumeClaimResult = - | { capacitySlot: number } + | { capacitySlot: number; claimToken: string; leaseBy: string } | { conflict: 'capacity' | 'overlap' | 'not-paused' | 'inactive' }; /** @@ -172,6 +175,16 @@ export interface SchedulesService { scheduledFor: string | Date, capacitySlot: number, ) => Promise; + /** Atomically validates the live schedule generation after the approval CAS and + * releases its short-lived document fence. This is the resume linearization point. */ + finalizeScheduleResumeClaim: ( + scheduleId: string, + claimToken: string, + leaseBy: string, + options?: { expectedConfigRevision?: number; automatic?: boolean }, + ) => Promise; + /** Releases only the schedule-document fence after an unconsumed/ambiguous approval CAS. */ + releaseScheduleResumeFence: (scheduleId: string, leaseBy: string) => Promise; /** * Whether a schedule is still live (exists and not soft-deleted). The loopback * chat controller calls this right after creating the generation job to re-fence @@ -850,32 +863,57 @@ export function createSchedulesService( if (!ownerLimits.enabled || globallyDisabled || !hasAccess) { return { conflict: 'inactive' }; } - const allocation = await engineDeps.withGlobalCapacitySlot( - Math.min(ownerLimits.fireConcurrency, deploymentLimits.fireConcurrency), - async (capacitySlot) => { - const attempt = await methods.markRunResumeClaimed( - scheduleId, - new Date(scheduledFor), - capacitySlot, - ); - return 'conflict' in attempt && attempt.conflict === 'slot-taken' - ? 'slot-taken' - : { claimed: attempt }; - }, + // FINAL schedule-side admission fence. Everything above is asynchronous and an + // owner edit/disable can land while it runs. Claim the schedule document under + // the expected config generation now, carry this lease through the approval CAS, + // and atomically consume it before provider execution begins. + const resumeLease = await methods.acquireResumeLease( + scheduleId, + options?.expectedConfigRevision, + options?.automatic !== false, + SCHEDULE_RESUME_LEASE_MS, ); - if (allocation === 'capacity') { - return { conflict: 'capacity' }; + if (resumeLease?.claimToken == null || resumeLease.leaseBy == null) { + return { conflict: 'inactive' }; } - const claimed = allocation.claimed; - if ('conflict' in claimed) { - // withCapacitySlot consumes this internal collision sentinel by retrying the - // next free slot. Normalize defensively at the public boundary as well so a - // custom/test allocator can never leak an implementation-only conflict. + let retainResumeLease = false; + try { + const allocation = await engineDeps.withGlobalCapacitySlot( + Math.min(ownerLimits.fireConcurrency, deploymentLimits.fireConcurrency), + async (capacitySlot) => { + const attempt = await methods.markRunResumeClaimed( + scheduleId, + new Date(scheduledFor), + capacitySlot, + ); + return 'conflict' in attempt && attempt.conflict === 'slot-taken' + ? 'slot-taken' + : { claimed: attempt }; + }, + ); + if (allocation === 'capacity') { + return { conflict: 'capacity' }; + } + const claimed = allocation.claimed; + if ('conflict' in claimed) { + // withCapacitySlot consumes this internal collision sentinel by retrying the + // next free slot. Normalize defensively at the public boundary as well so a + // custom/test allocator can never leak an implementation-only conflict. + return { + conflict: claimed.conflict === 'slot-taken' ? 'capacity' : claimed.conflict, + }; + } + retainResumeLease = true; return { - conflict: claimed.conflict === 'slot-taken' ? 'capacity' : claimed.conflict, + capacitySlot: claimed.capacitySlot, + claimToken: resumeLease.claimToken, + leaseBy: resumeLease.leaseBy, }; + } finally { + if (!retainResumeLease) { + await methods.releaseLeaseByHolder(scheduleId, resumeLease.leaseBy); + } } - return { capacitySlot: claimed.capacitySlot }; }); } @@ -887,6 +925,32 @@ export function createSchedulesService( return methods.releaseRunResumeClaim(scheduleId, new Date(scheduledFor), capacitySlot); } + async function finalizeScheduleResumeClaim( + scheduleId: string, + claimToken: string, + leaseBy: string, + options?: { expectedConfigRevision?: number; automatic?: boolean }, + ): Promise { + const consumed = await methods.consumeResumeLease( + scheduleId, + claimToken, + leaseBy, + options?.automatic !== false, + options?.expectedConfigRevision, + ); + if (!consumed) { + // An owner edit rotates claimToken but deliberately leaves the old holder's + // lease in place. Clear only that holder so the edited schedule is immediately + // usable; a takeover changed leaseBy and is therefore untouched. + await methods.releaseLeaseByHolder(scheduleId, leaseBy); + } + return consumed; + } + + function releaseScheduleResumeFence(scheduleId: string, leaseBy: string): Promise { + return methods.releaseLeaseByHolder(scheduleId, leaseBy); + } + /** Aborts an active run's loopback job (identity-guarded). Returns whether the * abort was delivered (false when the job wasn't reachable — e.g. a peer worker's * private store, or a transient error). */ @@ -1319,6 +1383,8 @@ export function createSchedulesService( recordScheduleOutcome, claimScheduleResume, releaseScheduleResumeClaim, + finalizeScheduleResumeClaim, + releaseScheduleResumeFence, isScheduleLive, deleteScheduleForOwner, quiesceUserSchedules, diff --git a/packages/data-schemas/src/methods/schedule.methods.spec.ts b/packages/data-schemas/src/methods/schedule.methods.spec.ts index ba2bd21e76..d1bdeb45ab 100644 --- a/packages/data-schemas/src/methods/schedule.methods.spec.ts +++ b/packages/data-schemas/src/methods/schedule.methods.spec.ts @@ -1124,6 +1124,37 @@ describe('acquireManualRunLease / releaseLease', () => { }); }); +describe('scheduled resume lease fencing', () => { + it('atomically consumes the same live config generation and releases its lease', async () => { + const schedule = await methods.createSchedule(scheduleData()); + const fence = await methods.acquireResumeLease(schedule.id, 0, true, 60_000); + + expect(fence?.claimToken).toBeTruthy(); + expect(fence?.leaseBy).toMatch(/^resume:/); + await expect( + methods.consumeResumeLease(schedule.id, fence!.claimToken!, fence!.leaseBy!, true, 0), + ).resolves.toBe(true); + const consumed = await getSchedule(schedule.id); + expect(consumed.leaseUntil).toBeUndefined(); + expect(consumed.leaseBy).toBeUndefined(); + }); + + it('rejects an owner edit that wins before the final resume handoff', async () => { + const schedule = await methods.createSchedule(scheduleData({ name: 'before' })); + const fence = await methods.acquireResumeLease(schedule.id, 0, true, 60_000); + await methods.updateScheduleById(schedule.id, schedule.user, { name: 'after' }); + + await expect( + methods.consumeResumeLease(schedule.id, fence!.claimToken!, fence!.leaseBy!, true, 0), + ).resolves.toBe(false); + // Owner edits intentionally leave the in-flight holder fields intact. The stale + // resume can release only its unique holder without touching any newer takeover. + expect((await getSchedule(schedule.id)).leaseBy).toBe(fence!.leaseBy); + await methods.releaseLeaseByHolder(schedule.id, fence!.leaseBy!); + expect((await getSchedule(schedule.id)).leaseBy).toBeUndefined(); + }); +}); + describe('claim-token fencing (stale worker cannot mutate an edited/deleted schedule)', () => { it('an owner edit rotates the token so a stale disable/advance no-ops and revalidate fails', async () => { const schedule = await methods.createSchedule(scheduleData()); diff --git a/packages/data-schemas/src/methods/schedule.ts b/packages/data-schemas/src/methods/schedule.ts index f0c25e1fe5..14c08a6067 100644 --- a/packages/data-schemas/src/methods/schedule.ts +++ b/packages/data-schemas/src/methods/schedule.ts @@ -168,7 +168,20 @@ export type ScheduleMethods = { userId: string | Types.ObjectId, leaseMs: number, ) => Promise; - releaseLease: (id: string, expectedClaimToken?: string) => Promise; + acquireResumeLease: ( + id: string, + expectedConfigRevision: number | undefined, + requireEnabled: boolean, + leaseMs: number, + ) => Promise; + consumeResumeLease: ( + id: string, + expectedClaimToken: string, + expectedLeaseBy: string, + requireEnabled: boolean, + expectedConfigRevision?: number, + ) => Promise; + releaseLease: (id: string, expectedClaimToken?: string) => Promise; releaseLeaseByHolder: (id: string, leaseBy: string) => Promise; revalidateClaim: (id: string, claimToken: string, requireEnabled?: boolean) => Promise; advanceSchedule: ( @@ -176,7 +189,7 @@ export type ScheduleMethods = { nextRunAt: Date | null, expectedNextRunAt?: Date | null, expectedClaimToken?: string, - ) => Promise; + ) => Promise; disableSchedule: ( id: string, reason: ScheduleDisabledReason, @@ -540,17 +553,90 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche .lean(); } + /** + * Takes a short-lived schedule-document fence for a paused run's resume. The + * approval and its ScheduleRun row live in different stores, so this claim is + * carried through the approval CAS and atomically consumed immediately after it. + * An owner edit rotates claimToken; a policy disable clears the lease/enabled bit. + * Either therefore makes consumeResumeLease fail before provider execution starts. + */ + async function acquireResumeLease( + id: string, + expectedConfigRevision: number | undefined, + requireEnabled: boolean, + leaseMs: number, + ): Promise { + const claimToken = randomUUID(); + const leaseBy = `resume:${claimToken}`; + const now = new Date(); + const takeoverCutoff = new Date(now.getTime() - LEASE_SKEW_MARGIN_MS); + return Schedule() + .findOneAndUpdate( + { + id, + deleting: { $ne: true }, + ...(requireEnabled ? { enabled: true } : {}), + ...(expectedConfigRevision !== undefined + ? { configRevision: expectedConfigRevision } + : {}), + $or: [ + { leaseUntil: { $exists: false } }, + { leaseUntil: null }, + { leaseUntil: { $lt: takeoverCutoff } }, + ], + }, + { + $set: { + leaseUntil: new Date(now.getTime() + leaseMs), + leaseBy, + claimToken, + }, + }, + { new: true }, + ) + .lean(); + } + + /** + * Atomically linearizes a scheduled resume against owner edits, deletion, policy + * disable, lease expiry, and takeover, then releases the short-lived fence. A + * successful match means the approval was consumed under the same live schedule + * generation that acquired the fence; a miss must stop before provider execution. + */ + async function consumeResumeLease( + id: string, + expectedClaimToken: string, + expectedLeaseBy: string, + requireEnabled: boolean, + expectedConfigRevision?: number, + ): Promise { + const result = await Schedule().updateOne( + { + id, + claimToken: expectedClaimToken, + leaseBy: expectedLeaseBy, + deleting: { $ne: true }, + leaseUntil: { $gt: new Date() }, + ...(requireEnabled ? { enabled: true } : {}), + ...(expectedConfigRevision !== undefined ? { configRevision: expectedConfigRevision } : {}), + }, + { $unset: { leaseUntil: 1, leaseBy: 1 } }, + ); + return (result.matchedCount ?? 0) > 0; + } + /** * Releases a lease WITHOUT advancing nextRunAt (manual runs never reschedule). * Fenced on the claim token when provided so a stale worker cannot strip a lease * a different claimer now holds. */ - async function releaseLease(id: string, expectedClaimToken?: string): Promise { + async function releaseLease(id: string, expectedClaimToken?: string): Promise { const filter: Record = { id }; if (expectedClaimToken !== undefined) { filter.claimToken = expectedClaimToken; } - await Schedule().updateOne(filter, { $unset: { leaseUntil: 1, leaseBy: 1 } }); + const result = await Schedule().updateOne(filter, { $unset: { leaseUntil: 1, leaseBy: 1 } }); + return (result.matchedCount ?? 0) > 0; } /** @@ -607,7 +693,7 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche nextRunAt: Date | null, expectedNextRunAt?: Date | null, expectedClaimToken?: string, - ): Promise { + ): Promise { const filter: Record = { id }; if (expectedNextRunAt !== undefined) { filter.nextRunAt = expectedNextRunAt; @@ -615,10 +701,11 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche if (expectedClaimToken !== undefined) { filter.claimToken = expectedClaimToken; } - await Schedule().updateOne(filter, { + const result = await Schedule().updateOne(filter, { $set: { ...(nextRunAt ? { nextRunAt } : {}) }, $unset: { leaseUntil: 1, leaseBy: 1, ...(nextRunAt ? {} : { nextRunAt: 1 }) }, }); + return (result.matchedCount ?? 0) > 0; } /** @@ -1745,6 +1832,8 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche countSchedulesByUser, claimDueSchedule, acquireManualRunLease, + acquireResumeLease, + consumeResumeLease, releaseLease, releaseLeaseByHolder, revalidateClaim,