From c7f3f8983b93f9982156822ea7358bc3363853e4 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 22 Jul 2026 23:51:28 -0400 Subject: [PATCH] feat: scheduled-run lifecycle epoch, durable resume lease, DB-enforced capacity Foundation for the new merge blockers. Collapses B1/B3/B8/B10 (and B7's per-run half) into one coherent ScheduleRun state machine rather than four overlapping patches, since they all mutate the same row. B1 - pause epoch + CAS: ScheduleRun.resumeSeq is a monotonic per-occurrence segment counter, incremented by every resume reservation. recordRunOutcome's pause branch now CASes on expectResumeSeq (both the guard read and the row flip), so a stale requires_action callback from a superseded segment can no longer demote a run that a resume already promoted. resume.js threads its segment's resumeSeq into the re-pause. B3 - durable resuming state: promoteRunToStarted (whose lossy 'missing' -> 'ok' admitted both racers) is replaced by acquireResumeLease/markResumeClaimed/commitResumeLease/ releaseResumeLease. The lease is holder-fenced and deadline-bounded, and records whether the approval was already consumed: a PRE-claim crash is adoptable (rolled back to requires_action, slot freed, approval stays actionable), a POST-claim crash must roll forward. resume.js drives all four transitions. B8 - atomic global capacity: a unique partial index on {capacitySlot} where status:'started' makes fireConcurrency a DB-enforced bound instead of a read-then-compare count. Both admissions (fire insert, resume promotion) claim a slot in the SAME write via withCapacitySlot, so two admissions of DIFFERENT schedules can no longer both pass a cap-1 check. Capacity is now refused BEFORE inserting, so the old reserve-then-rollback path is gone. B10 - settlement: requestRunAbort records abortRequestedAt WITHOUT freeing the slot; the slot is released only by a terminal outcome, i.e. when the generation owner confirms the run actually stopped. Aborted runs therefore keep counting capacity until settled. B7 (per-run half) - ScheduleRun.configRevision captures the schedule's revision at claim time; applyTerminalBookkeeping filters on it so a run started under older config cannot apply counters or auto-disable a schedule the owner has since edited. Schedule. configRevision is bumped ONLY by updateScheduleById (owner edits), atomically with the claimToken rotation. B2 (field only) - User.deletionRequestedAt added as the durable deletion-barrier field; the barrier logic itself is not wired yet. Also fixes a latent discriminator bug: isActiveRunConflict matched any duplicate-key lacking scheduledFor, so the new {capacitySlot} index would have been misread as a per-schedule overlap. It now matches scheduleId exactly. Absent fields disable each fence, so existing rows/schedules keep current behavior and no migration is required. --- api/server/controllers/agents/resume.js | 70 ++++- api/server/services/Schedules/index.js | 3 + packages/api/src/schedules/capacity.ts | 53 ++++ packages/api/src/schedules/fire.spec.ts | 98 ++++-- packages/api/src/schedules/fire.ts | 59 ++-- packages/api/src/schedules/service.ts | 153 ++++++++-- packages/api/src/schedules/types.ts | 12 + packages/data-schemas/src/methods/schedule.ts | 287 ++++++++++++++++-- packages/data-schemas/src/schema/schedule.ts | 11 + .../data-schemas/src/schema/scheduleRun.ts | 50 +++ packages/data-schemas/src/schema/user.ts | 9 + packages/data-schemas/src/types/schedule.ts | 15 + packages/data-schemas/src/types/user.ts | 2 + 13 files changed, 716 insertions(+), 106 deletions(-) create mode 100644 packages/api/src/schedules/capacity.ts diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 0e45428f73..afcb48e05d 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -23,7 +23,13 @@ const { getMCPRequestContext, cleanupMCPRequestContextForReq, } = require('~/server/services/MCPRequestContext'); -const { recordScheduleOutcome, reserveScheduledResume } = require('~/server/services/Schedules'); +const { + recordScheduleOutcome, + reserveScheduledResume, + markScheduledResumeClaimed, + commitScheduledResume, + releaseScheduledResume, +} = require('~/server/services/Schedules'); const { saveMessage, getConvo, getMessages } = require('~/models'); /** @@ -544,6 +550,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) // approval drives whatever is `started`, and an unclaimed promotion self-heals when // the reconciler surfaces the still-paused job (so a losing racer never flips the // winner's active row). + let resumeLease = null; if (job.metadata?.scheduleId) { let reservation; try { @@ -559,10 +566,10 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) logger.error('[ResumeAgentController] Scheduled resume reservation failed', err); return res.status(500).json({ error: 'Failed to resume' }); } - if (reservation !== 'ok') { + if (reservation.outcome !== 'ok') { await decrementPendingRequest(userId); logger.debug( - `[ResumeAgentController] Deferring scheduled resume (${reservation}): ${streamId}`, + `[ResumeAgentController] Deferring scheduled resume (${reservation.outcome}): ${streamId}`, ); const deferrals = { gone: { status: 410, error: 'This schedule no longer exists' }, @@ -571,10 +578,22 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) status: 409, error: 'Another run for this schedule is in progress; try again shortly', }, + held: { + status: 409, + error: 'This scheduled run is already resuming; try again shortly', + }, }; - const { status, error } = deferrals[reservation] ?? deferrals.overlap; + const { status, error } = deferrals[reservation.outcome] ?? deferrals.overlap; return res.status(status).json({ error }); } + // Own the durable RESUMING lease: every later transition is fenced on this holder, + // and resumeSeq is the CAS token any pause written by THIS segment must carry. + resumeLease = { + scheduleId: job.metadata.scheduleId, + scheduledFor: job.metadata.scheduledFor, + holder: reservation.holder, + resumeSeq: reservation.resumeSeq, + }; } // Atomically claim the resume. The single winner drives the run; a racing second @@ -587,13 +606,42 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) claimed = await GenerationJobManager.approvals.resolve(streamId, pendingAction.actionId); } catch (err) { await decrementPendingRequest(userId); + // Pre-claim failure: roll the reservation back so the approval stays actionable + // and the capacity slot isn't held by a resume that never started. + if (resumeLease) { + await releaseScheduledResume( + resumeLease.scheduleId, + resumeLease.scheduledFor, + resumeLease.holder, + ).catch(() => undefined); + } logger.error('[ResumeAgentController] Failed to claim resume', err); return res.status(500).json({ error: 'Failed to resume' }); } if (!claimed) { await decrementPendingRequest(userId); + // Still pre-claim: roll the reservation back so the approval stays actionable and + // the capacity slot is freed rather than held by a resume that never ran. + if (resumeLease) { + await releaseScheduledResume( + resumeLease.scheduleId, + resumeLease.scheduledFor, + resumeLease.holder, + ).catch(() => undefined); + } return res.status(409).json({ error: 'This action was already resolved or has expired' }); } + // The approval is consumed: from here a crash must roll FORWARD (the reconciler must + // finish or terminalize this run), so the lease is no longer adoptable by a retry. + if (resumeLease) { + await markScheduledResumeClaimed( + resumeLease.scheduleId, + resumeLease.scheduledFor, + resumeLease.holder, + ).catch((err) => + logger.error('[ResumeAgentController] Failed to mark resume lease claimed', err), + ); + } // Seed the run-scoped MCP request-context store BEFORE the ACK: once `res.json` // finishes the response, a later `getMCPRequestContext(req, res)` (from tool loading) @@ -691,6 +739,17 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) // the paused job's createdAt — used by the re-pause CAS pre-check + checkpoint prune to // avoid acting on a job a newer request has since replaced. client.jobCreatedAt = job.createdAt; + // Reconstruction succeeded, so the RESUMING phase is over: drop the lease and let + // the row be an ordinary `started` run again. A crash before this point leaves the + // lease for the reconciler (post-claim => roll forward), never a silently stranded + // `running` job the reconciler would skip. + if (resumeLease) { + await commitScheduledResume( + resumeLease.scheduleId, + resumeLease.scheduledFor, + resumeLease.holder, + ).catch((err) => logger.error('[ResumeAgentController] Failed to commit resume lease', err)); + } client.responseMessageId = job.metadata.responseMessageId; client.parentMessageId = job.metadata.userMessage?.messageId ?? Constants.NO_PARENT; // Read the pre-pause content BEFORE swapping the store's content reference: the @@ -742,6 +801,9 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) scheduledFor: job.metadata.scheduledFor, status: 'requires_action', conversationId, + // Epoch CAS: if a newer resume already bumped resumeSeq, this pause belongs to + // a superseded segment and must not demote the run that resume promoted. + expectResumeSeq: resumeLease?.resumeSeq, }); } return; diff --git a/api/server/services/Schedules/index.js b/api/server/services/Schedules/index.js index 8dfaea8fc3..f327445991 100644 --- a/api/server/services/Schedules/index.js +++ b/api/server/services/Schedules/index.js @@ -26,6 +26,9 @@ module.exports = { recordScheduleOutcome: service.recordScheduleOutcome, isScheduleLive: service.isScheduleLive, reserveScheduledResume: service.reserveScheduledResume, + markScheduledResumeClaimed: service.markScheduledResumeClaimed, + commitScheduledResume: service.commitScheduledResume, + releaseScheduledResume: service.releaseScheduledResume, deleteScheduleForOwner: service.deleteScheduleForOwner, quiesceUserSchedules: service.quiesceUserSchedules, initializeScheduleEngine: service.initializeScheduleEngine, diff --git a/packages/api/src/schedules/capacity.ts b/packages/api/src/schedules/capacity.ts new file mode 100644 index 0000000000..79770b5e5d --- /dev/null +++ b/packages/api/src/schedules/capacity.ts @@ -0,0 +1,53 @@ +/** Occupancy of the global scheduled-run capacity slots. `unslotted` counts legacy + * `started` rows written before slots existed; they shrink the effective cap so the + * bound stays conservative during rollout instead of transiently overshooting. */ +export interface CapacityOccupancy { + takenSlots: number[]; + unslotted: number; +} + +/** A claim either succeeded (carrying the caller's own result) or lost the slot race. */ +export type SlotClaimResult = { claimed: T } | 'slot-taken'; + +/** + * Allocates the lowest free global capacity slot and hands it to `claim`, retrying the + * next free slot when the DB rejects a collision on the unique partial index. + * + * This replaces "count active runs, compare to the cap, then insert": the count is a + * read-then-write race, so two admissions of DIFFERENT schedules could both observe + * cap-1 and both proceed. Here the slot itself is the contended resource and the + * unique index is the arbiter, so the cap is enforced by the database. + * + * Bounded by cap+1 attempts: every collision advances to a strictly higher free slot. + */ +export async function withCapacitySlot( + cap: number, + readOccupancy: () => Promise, + claim: (slot: number) => Promise>, +): Promise<{ claimed: T } | 'capacity'> { + if (cap <= 0) { + return 'capacity'; + } + for (let attempt = 0; attempt <= cap; attempt++) { + const { takenSlots, unslotted } = await readOccupancy(); + if (takenSlots.length + unslotted >= cap) { + return 'capacity'; + } + const taken = new Set(takenSlots); + let slot = -1; + for (let candidate = 0; candidate < cap; candidate++) { + if (!taken.has(candidate)) { + slot = candidate; + break; + } + } + if (slot < 0) { + return 'capacity'; + } + const result = await claim(slot); + if (result !== 'slot-taken') { + return result; + } + } + return 'capacity'; +} diff --git a/packages/api/src/schedules/fire.spec.ts b/packages/api/src/schedules/fire.spec.ts index fadde38e81..f36cdd89a0 100644 --- a/packages/api/src/schedules/fire.spec.ts +++ b/packages/api/src/schedules/fire.spec.ts @@ -1,5 +1,6 @@ import type { ScheduleEngineDeps, ScheduleLimits, ScheduleUserContext } from './types'; import type { FireableSchedule } from './types'; +import { withCapacitySlot } from './capacity'; import { fireSchedule } from './fire'; const OWNER: ScheduleUserContext = { id: 'user-1', tenantId: 't1', role: 'USER' }; @@ -34,7 +35,10 @@ function makeSchedule(overrides: Partial = {}): FireableSchedu /** In-memory run store exercising the real insert/count/delete/idempotency interplay. */ function makeMethods() { - const runs = new Map(); + const runs = new Map< + string, + { status: string; conversationId?: string; capacitySlot?: number } + >(); const calls = { advance: 0, releaseLease: 0, @@ -73,21 +77,54 @@ function makeMethods() { // Mirrors the partial-unique-index semantics: same-occurrence row => 'duplicate'; // any OTHER started run for the schedule => 'overlap'; else reserve the slot. reserveStartedRun: jest.fn( - async (data: { scheduleId: string; scheduledFor: Date; conversationId?: string }) => { + async (data: { + scheduleId: string; + scheduledFor: Date; + conversationId?: string; + capacitySlot?: number; + }) => { const k = key(data.scheduleId, data.scheduledFor); if (runs.has(k)) { return { conflict: 'duplicate' as const }; } + // Mirrors the unique {capacitySlot} partial index (status:'started'). + if ( + data.capacitySlot != null && + [...runs.values()].some( + (r) => r.status === 'started' && r.capacitySlot === data.capacitySlot, + ) + ) { + return { conflict: 'slot-taken' as const }; + } const overlap = [...runs.entries()].some( ([rk, r]) => rk.startsWith(`${data.scheduleId}:`) && r.status === 'started', ); if (overlap) { return { conflict: 'overlap' as const }; } - runs.set(k, { status: 'started', conversationId: data.conversationId }); + runs.set(k, { + status: 'started', + conversationId: data.conversationId, + capacitySlot: data.capacitySlot, + }); return { run: { scheduleId: data.scheduleId, scheduledFor: data.scheduledFor } }; }, ), + getCapacityOccupancy: jest.fn(async () => { + const takenSlots: number[] = []; + let unslotted = 0; + for (const r of runs.values()) { + if (r.status !== 'started') { + continue; + } + if (typeof r.capacitySlot === 'number') { + takenSlots.push(r.capacitySlot); + } else { + unslotted += 1; + } + } + return { takenSlots, unslotted }; + }), revalidateClaim: jest.fn(async () => true), holdsLease: jest.fn(async () => true), scheduleExists: jest.fn(async () => true), @@ -134,6 +171,12 @@ function makeDeps( clearReconciledJob: async () => undefined, isJobStoreShared: () => true, countActiveRunsGlobal: async () => methods.countActiveRuns(), + withGlobalCapacitySlot: (cap: number, claim: (slot: number) => Promise) => + withCapacitySlot( + cap, + () => methods.getCapacityOccupancy(), + claim as Parameters[2], + ), ...over, } as ScheduleEngineDeps; } @@ -248,30 +291,43 @@ describe('fireSchedule', () => { expect(global.fetch).not.toHaveBeenCalled(); }); - it('rolls back the reservation when over the global capacity cap', async () => { + it('refuses the fire at the global capacity cap WITHOUT inserting a run', async () => { const { methods, runs } = makeMethods(); - // 5 already active → this insert makes 6 > cap(5), so it must roll back. + // All 5 slots taken → the allocator finds no free slot and never inserts. for (let i = 0; i < 5; i++) { - runs.set(`other-${i}:x`, { status: 'started' }); + runs.set(`other-${i}:x`, { status: 'started', capacitySlot: i }); } mockFetch(async () => okResponse()); const result = await fireSchedule(makeDeps(methods), makeSchedule(), LIMITS, dueAt()); expect(result.skipped).toBe('capacity'); expect(global.fetch).not.toHaveBeenCalled(); - // The reservation was rolled back — only the 5 pre-existing remain. + // Slot-based capacity is decided BEFORE the write, so there is nothing to roll back. expect([...runs.values()].filter((r) => r.status === 'started')).toHaveLength(5); - expect(methods.deleteScheduleRun).toHaveBeenCalledTimes(1); + expect(methods.reserveStartedRun).not.toHaveBeenCalled(); + expect(methods.deleteScheduleRun).not.toHaveBeenCalled(); }); - it('capacity rollback re-fires cleanly next tick, exactly once', async () => { + it('claims a free slot and never exceeds the cap when slots collide', async () => { + const { methods, runs } = makeMethods(); + // Slots 0 and 2 are taken; the allocator must land the fire on slot 1. + runs.set('other-a:x', { status: 'started', capacitySlot: 0 }); + runs.set('other-b:x', { status: 'started', capacitySlot: 2 }); + mockFetch(async () => okResponse()); + const result = await fireSchedule(makeDeps(methods), makeSchedule(), LIMITS, dueAt()); + expect(result.fired).toBe(true); + const own = [...runs.entries()].find(([k]) => k.startsWith('sched-1:')); + expect(own?.[1].capacitySlot).toBe(1); + }); + + it('re-fires cleanly next tick once capacity frees, exactly once', async () => { const { methods, runs } = makeMethods(); for (let i = 0; i < 5; i++) { - runs.set(`other-${i}:x`, { status: 'started' }); + runs.set(`other-${i}:x`, { status: 'started', capacitySlot: i }); } mockFetch(async () => okResponse()); const schedule = makeSchedule(); const when = dueAt(); - // Tick 1: at capacity → rolled back. + // Tick 1: every slot taken → refused before any insert. const first = await fireSchedule(makeDeps(methods), schedule, LIMITS, when); expect(first.skipped).toBe('capacity'); // Capacity frees up before the next tick. @@ -279,7 +335,7 @@ describe('fireSchedule', () => { // Tick 2: same occurrence re-claimed → now fires, exactly one live run. const second = await fireSchedule(makeDeps(methods), schedule, LIMITS, when); expect(second.fired).toBe(true); - expect(methods.reserveStartedRun).toHaveBeenCalledTimes(2); // reserve, rollback, reserve + expect(methods.reserveStartedRun).toHaveBeenCalledTimes(1); // only the successful tick inserts expect( [...runs.entries()].filter(([k, r]) => k.startsWith('sched-1:') && r.status === 'started'), ).toHaveLength(1); @@ -288,15 +344,14 @@ describe('fireSchedule', () => { it('preserves the reserved run for reconcile when the lease was taken over', async () => { const { methods, runs } = makeMethods(); - // At capacity, so this fire will roll back its reservation. - for (let i = 0; i < 5; i++) { - runs.set(`other-${i}:x`, { status: 'started' }); - } + // An owner edit/takeover superseded this fire after it reserved its run, which is + // the path that now drives rollbackReservation (capacity no longer inserts at all). + (methods.revalidateClaim as jest.Mock).mockResolvedValue(false); // Simulate a lease takeover: this worker no longer holds the claim. (methods.holdsLease as jest.Mock).mockResolvedValue(false); mockFetch(async () => okResponse()); const result = await fireSchedule(makeDeps(methods), makeSchedule(), LIMITS, dueAt()); - expect(result.skipped).toBe('capacity'); + expect(result.skipped).toBe('superseded'); // The reserved row must NOT be deleted (another worker owns the occurrence now); // it's left for the reconciler so the occurrence stays reconcilable. expect(methods.deleteScheduleRun).not.toHaveBeenCalled(); @@ -305,17 +360,14 @@ describe('fireSchedule', () => { it('deletes the reserved run when the schedule was hard-deleted mid-fire', async () => { const { methods, runs } = makeMethods(); - // At capacity, so this fire rolls back its reservation. - for (let i = 0; i < 5; i++) { - runs.set(`other-${i}:x`, { status: 'started' }); - } // Account deletion hard-deleted the schedule after this fire reserved its run: - // the lease is not held AND the schedule no longer exists. + // revalidation fails, the lease is not held AND the schedule no longer exists. + (methods.revalidateClaim as jest.Mock).mockResolvedValue(false); (methods.holdsLease as jest.Mock).mockResolvedValue(false); (methods.scheduleExists as jest.Mock).mockResolvedValue(false); mockFetch(async () => okResponse()); const result = await fireSchedule(makeDeps(methods), makeSchedule(), LIMITS, dueAt()); - expect(result.skipped).toBe('capacity'); + expect(result.skipped).toBe('superseded'); // The orphaned reservation (no schedule left to own it) is deleted, not leaked. expect(methods.deleteScheduleRun).toHaveBeenCalledWith('sched-1', expect.any(Date), 'started'); expect([...runs.entries()].some(([k]) => k.startsWith('sched-1:'))).toBe(false); diff --git a/packages/api/src/schedules/fire.ts b/packages/api/src/schedules/fire.ts index 11e50ed641..3814b7db00 100644 --- a/packages/api/src/schedules/fire.ts +++ b/packages/api/src/schedules/fire.ts @@ -331,11 +331,40 @@ export async function fireSchedule( // rejects a second `started` run for the schedule, so a concurrent occurrence // surfaces as 'overlap' with no read-then-insert race. const conversationId = randomUUID(); - const reservation = await methods.reserveStartedRun({ - ...baseRun, - conversationId, - firedAt: new Date(), - }); + // The GLOBAL fireConcurrency cap is enforced by claiming a unique capacity slot in + // the SAME insert that reserves the run, so it is decided by the DB rather than by + // a count read before the write. The allocator advances to the next free slot when + // another admission wins one, and reports 'capacity' only when genuinely saturated. + // Occupancy is read system-scoped so the cap stays global across tenants. + const allocation = await deps.withGlobalCapacitySlot( + ownerLimits.fireConcurrency, + async (capacitySlot) => { + const attempt = await methods.reserveStartedRun({ + ...baseRun, + conversationId, + firedAt: new Date(), + capacitySlot, + ...(typeof schedule.configRevision === 'number' + ? { configRevision: schedule.configRevision } + : {}), + }); + return 'conflict' in attempt && attempt.conflict === 'slot-taken' + ? 'slot-taken' + : { claimed: attempt }; + }, + ); + if (allocation === 'capacity') { + // Automatic claims keep the claim's lease as a backoff so the nextRunAt-sorted + // claimer doesn't immediately re-pick this row and starve others; nextRunAt is + // untouched, so the occurrence retries once the lease expires. A manual run-now + // MUST release its lease, or repeated Run-now clicks hit a misleading "already + // in progress" 409 for the full manual-lease TTL even after capacity frees. + if (options?.manual) { + await methods.releaseLease(schedule.id, claimToken); + } + return { fired: false, skipped: 'capacity' as const }; + } + const reservation = allocation.claimed; if ('conflict' in reservation) { if (reservation.conflict === 'overlap') { // Another occurrence of this schedule is already active. Record the skip @@ -348,26 +377,6 @@ export async function fireSchedule( return { fired: false, skipped: 'duplicate' as const }; } - // Reserve-then-verify GLOBAL capacity: the reservation above is atomic per - // schedule, but the cross-schedule fireConcurrency cap is a count. The count is - // GLOBAL (system tenant) — under the owner's tenant context it would only see - // this tenant's runs and multiple tenants could collectively exceed the cap. - // Compare against the OWNER-resolved fireConcurrency (matching run-now's - // request-scoped check). Roll back (status-fenced) if over. - const active = await deps.countActiveRunsGlobal(); - if (active > ownerLimits.fireConcurrency) { - await rollbackReservation(); - // Automatic claims keep the claim's lease as a backoff so the nextRunAt-sorted - // claimer doesn't immediately re-pick this row and starve others; nextRunAt is - // untouched, so the occurrence retries once the lease expires. A manual run-now - // MUST release its lease, or repeated Run-now clicks hit a misleading "already - // in progress" 409 for the full manual-lease TTL even after capacity frees. - if (options?.manual) { - await methods.releaseLease(schedule.id, claimToken); - } - return { fired: false, skipped: 'capacity' as const }; - } - // Last check before the point of no return: re-verify this fire still holds an // authoritative claim (same claim token, lease unexpired, not deleting; and for // an automatic fire, still enabled). An owner delete/edit or a lease-expiry diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index 5ce197f4aa..e3700cacf2 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -1,3 +1,4 @@ +import { randomUUID } from 'node:crypto'; import { logger, runAsSystem, tenantStorage } from '@librechat/data-schemas'; import { getRefillEligibilityDate, Permissions, PermissionTypes } from 'librechat-data-provider'; import type { ScheduleMethods, AppConfig, IBalance } from '@librechat/data-schemas'; @@ -23,6 +24,7 @@ import { getAppConfigOptionsFromUser } from '../app/service'; import { DEFAULT_SCHEDULE_LIMITS } from './types'; import { getBalanceConfig } from '../app/config'; import { startScheduleEngine } from './engine'; +import { withCapacitySlot } from './capacity'; /** Recordable terminal/paused run outcome, as accepted by `recordRunOutcome`. */ type ScheduleRunOutcomeStatus = Parameters[0]['status']; @@ -33,7 +35,18 @@ type ScheduleRunOutcomeStatus = Parameters[ * is active; `capacity` = the global cap is saturated; each defers with the approval * left unconsumed. `ok` = reserved (or already active), let the claim proceed. */ -export type ResumeCheck = 'ok' | 'gone' | 'overlap' | 'capacity'; +export type ResumeCheck = 'ok' | 'gone' | 'overlap' | 'capacity' | 'held'; + +/** + * Outcome of the durable resume reservation. `ok` carries the lease identity the + * caller must fence every later transition on: `holder` for commit/release and + * `resumeSeq` as the CAS token for any pause written by this segment. + */ +export type ResumeReservation = + | { outcome: 'ok'; holder: string; resumeSeq: number } + /** Not a scheduled resume (ordinary interactive HITL): there is no lease to drive. */ + | { outcome: 'not-scheduled' } + | { outcome: Exclude }; /** Whether a persisted job still carries a given scheduled occurrence's identity. */ function jobMatchesIdentity(job: SerializableJobData, identity: JobIdentity): boolean { @@ -49,6 +62,8 @@ export interface RecordScheduleOutcomeInput { status: ScheduleRunOutcomeStatus; conversationId?: string; error?: string; + /** Pause CAS token: the resumeSeq the writing segment observed. */ + expectResumeSeq?: number; } /** @@ -105,16 +120,37 @@ export interface SchedulesService { */ isScheduleLive: (scheduleId: string) => Promise; /** - * Atomically reserves the active slot for a HITL resume BEFORE the approval claim - * (so a deferral leaves the approval claimable): promotes the paused run to - * `started` (the single-active partial index makes per-schedule overlap atomic) - * and reserve-then-verifies the global fireConcurrency cap, rolling back its OWN - * promotion on overshoot. 'gone' = the schedule was deleted; 'overlap' = another - * occurrence is active; 'capacity' = global cap saturated. The caller must NOT - * release on a lost approval claim — the claim winner drives whatever is `started`; - * an unclaimed promotion self-heals when the reconciler surfaces the pause. + * Acquires the durable RESUMING lease for a HITL resume, BEFORE the approval claim + * (so a deferral leaves the approval claimable). Per-schedule overlap is enforced by + * the single-active partial index and the GLOBAL fireConcurrency cap by the unique + * capacity-slot index, so two resumes of different schedules can never both admit. + * 'gone' = deleted/terminal; 'overlap' = another occurrence is active; 'capacity' = + * cap saturated; 'held' = another live resume attempt owns this occurrence. On 'ok' + * the caller MUST drive the lease: markResumeClaimed after consuming the approval, + * then commitResumeLease on success or releaseResumeLease while still pre-claim. */ - reserveScheduledResume: (scheduleId: string, scheduledFor: string | Date) => Promise; + reserveScheduledResume: ( + scheduleId: string, + scheduledFor: string | Date, + ) => Promise; + /** Marks the resume lease as having consumed the approval (crash must roll forward). */ + markScheduledResumeClaimed: ( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ) => Promise; + /** Commits a resume once its generation is reconstructed and running. */ + commitScheduledResume: ( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ) => Promise; + /** Rolls a still-unclaimed resume back to `requires_action`, freeing its slot. */ + releaseScheduledResume: ( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ) => Promise; /** Soft-deletes an owner's schedule: stop claims, abort active runs, drain, erase. */ deleteScheduleForOwner: (scheduleId: string, userId: string) => Promise; /** Quiesces all of a user's schedules ahead of account deletion (stop + abort). */ @@ -160,6 +196,9 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer } const MANUAL_RUN_LEASE_MS = 5 * 60 * 1000; + // Lease held by an in-flight HITL resume. Long enough to cover reconstruction, short + // enough that a crashed PRE-CLAIM resume is reclaimable without a human. + const RESUME_LEASE_MS = 2 * 60 * 1000; // Bounded wait for aborted scheduled runs to settle during account-deletion quiesce, // before the message/conversation cascade runs. Long enough to cover a generation that // already returned from the model finishing its persistence; capped so account deletion @@ -330,6 +369,11 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer // Counted in system scope so the cap is GLOBAL — a per-owner (tenant-scoped) // count would let multiple tenants collectively exceed fireConcurrency. countActiveRunsGlobal: () => runAsSystem(() => methods.countActiveRuns()), + // Occupancy is read in SYSTEM scope so the cap is global across tenants (the + // owner's tenant context would only see its own runs); the claim itself stays in + // the caller's context so the inserted row keeps correct tenant ownership. + withGlobalCapacitySlot: (cap, claim) => + withCapacitySlot(cap, () => runAsSystem(() => methods.getCapacityOccupancy()), claim), }; let engine: ReturnType | undefined; @@ -424,6 +468,7 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer status, conversationId, error, + expectResumeSeq, }: RecordScheduleOutcomeInput): Promise { if (!scheduleId || !scheduledFor) { return true; @@ -435,6 +480,7 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer const schedule = await methods.getScheduleById(scheduleId); const owner = schedule ? await engineDeps.getUserContext(schedule.user) : null; const limits = await getLimits(owner ?? undefined); + const run = await methods.getRun(scheduleId, new Date(scheduledFor)); await methods.recordRunOutcome({ scheduleId, scheduledFor: new Date(scheduledFor), @@ -442,6 +488,10 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer conversationId, error, autoDisableAfterFailures: limits.autoDisableAfterFailures, + expectResumeSeq, + // Fence terminal bookkeeping/auto-disable to the config this run started + // under, so an owner edit or re-enable since then is never acted on. + expectConfigRevision: run?.configRevision, }); return true; } catch (err) { @@ -479,46 +529,82 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer async function reserveScheduledResume( scheduleId: string, scheduledFor: string | Date, - ): Promise { + ): Promise { if (!scheduleId || !scheduledFor) { - return 'ok'; + return { outcome: 'not-scheduled' }; } // getScheduleById hides deleted/soft-deleted schedules, so a null here means the // owner already deleted the schedule — its paused run must not be resumable even // if the delete's best-effort abort raced. Reject before touching the run. const schedule = await methods.getScheduleById(scheduleId); if (schedule == null) { - return 'gone'; + return { outcome: 'gone' }; } const when = new Date(scheduledFor); // Overlap = a DIFFERENT occurrence is currently `started`. Exclude this run's own // row: if its pause bookkeeping failed transiently the row can still be `started`, // and treating that as an overlap would wrongly reject resuming THIS same occurrence. if (await methods.hasOtherActiveRun(scheduleId, when)) { - return 'overlap'; + return { outcome: 'overlap' }; } - // Read-only capacity gate BEFORE promoting, so we never mutate a row a concurrent - // same-pause resume may already be driving (no rollback path exists). Discount - // this occurrence's OWN `started` row when present (a transient pause-bookkeeping - // failure): resuming it adds no new active run, so the global count already - // includes it and must not block the resume. const owner = await engineDeps.getUserContext(schedule.user); const limits = await getLimits(owner ?? undefined); - const selfActive = await methods.isOccurrenceStarted(scheduleId, when); - if (!selfActive && (await engineDeps.countActiveRunsGlobal()) >= limits.fireConcurrency) { - return 'capacity'; + const holder = randomUUID(); + // ATOMIC global capacity: the slot (not a count) is the contended resource, so two + // resumes of DIFFERENT schedules can no longer both pass a cap-1 check — the unique + // partial index rejects the loser and the allocator advances it to the next free + // slot, or reports 'capacity' when genuinely saturated. + const reserved = await runAsSystem(() => + withCapacitySlot( + limits.fireConcurrency, + () => methods.getCapacityOccupancy(), + async (slot) => { + const lease = await methods.acquireResumeLease({ + scheduleId, + scheduledFor: when, + holder, + ttlMs: RESUME_LEASE_MS, + capacitySlot: slot, + }); + return lease.outcome === 'slot-taken' ? 'slot-taken' : { claimed: lease }; + }, + ), + ); + if (reserved === 'capacity') { + return { outcome: 'capacity' }; } - // Reserve the single active slot. If a different occurrence won the slot since - // the read-only check above, promoteRunToStarted returns 'overlap' — defer the - // resume (approval stays claimable) rather than running a second concurrent - // occurrence with the paused row still `requires_action` (which overlap/capacity - // accounting would miss). 'missing' means a concurrent same-pause resume already - // promoted it — proceed. Never rolled back. - const promoted = await methods.promoteRunToStarted(scheduleId, when); - if (promoted === 'overlap') { - return 'overlap'; + const lease = reserved.claimed; + if (lease.outcome === 'acquired') { + return { outcome: 'ok', holder: lease.holder, resumeSeq: lease.resumeSeq }; } - return 'ok'; + // 'held' = another live resume attempt owns this occurrence (previously collapsed + // into a lossy 'missing' -> 'ok', which admitted both racers). 'overlap'/'gone' + // defer with the approval left unconsumed. + return { outcome: lease.outcome === 'held' ? 'held' : lease.outcome }; + } + + async function markScheduledResumeClaimed( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ): Promise { + return methods.markResumeClaimed(scheduleId, new Date(scheduledFor), holder, RESUME_LEASE_MS); + } + + async function commitScheduledResume( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ): Promise { + return methods.commitResumeLease(scheduleId, new Date(scheduledFor), holder); + } + + async function releaseScheduledResume( + scheduleId: string, + scheduledFor: string | Date, + holder: string, + ): Promise { + return methods.releaseResumeLease(scheduleId, new Date(scheduledFor), holder); } /** Aborts an active run's loopback job (identity-guarded). Returns whether the @@ -645,6 +731,9 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer recordScheduleOutcome, isScheduleLive, reserveScheduledResume, + markScheduledResumeClaimed, + commitScheduledResume, + releaseScheduledResume, deleteScheduleForOwner, quiesceUserSchedules, initializeScheduleEngine, diff --git a/packages/api/src/schedules/types.ts b/packages/api/src/schedules/types.ts index 4e861e04ba..7f0821f02f 100644 --- a/packages/api/src/schedules/types.ts +++ b/packages/api/src/schedules/types.ts @@ -1,5 +1,6 @@ import type { ScheduleMethods, ISchedule } from '@librechat/data-schemas'; import type { Types } from 'mongoose'; +import type { SlotClaimResult } from './capacity'; export interface ScheduleLimits { /** Feature-level switch: when false the engine claims/fires nothing. */ @@ -96,6 +97,17 @@ export interface ScheduleEngineDeps { clearReconciledJob: (conversationId: string, identity: JobIdentity) => Promise; /** Global in-flight scheduled-run count (system tenant scope) for the fire cap. */ countActiveRunsGlobal: () => Promise; + /** + * Runs `claim` against the lowest free GLOBAL capacity slot, retrying the next slot + * when the unique partial index rejects a collision. Enforces fireConcurrency in the + * database instead of via a read-then-compare count, so concurrent admissions of + * different schedules cannot both pass a cap-1 check. Occupancy is read in system + * scope so the cap stays global across tenants. + */ + withGlobalCapacitySlot: ( + cap: number, + claim: (slot: number) => Promise>, + ) => Promise<{ claimed: T } | 'capacity'>; } /** The immutable scheduled identity of a generation job, for reconcile/abort fencing. */ diff --git a/packages/data-schemas/src/methods/schedule.ts b/packages/data-schemas/src/methods/schedule.ts index 34657783c9..616aa82051 100644 --- a/packages/data-schemas/src/methods/schedule.ts +++ b/packages/data-schemas/src/methods/schedule.ts @@ -30,14 +30,25 @@ function isOccurrenceDuplicate(error: unknown): boolean { return err?.code === DUPLICATE_KEY && err.keyPattern != null && 'scheduledFor' in err.keyPattern; } -/** A duplicate-key error whose conflict is the single-active-run partial index. */ +/** A duplicate-key error whose conflict is the single-active-run partial index + * ({scheduleId} where status:'started'). Matched EXACTLY on scheduleId so the + * global {capacitySlot} index below is never misread as a per-schedule overlap. */ function isActiveRunConflict(error: unknown): boolean { const err = error as DuplicateKeyError; return ( - err?.code === DUPLICATE_KEY && err.keyPattern != null && !('scheduledFor' in err.keyPattern) + err?.code === DUPLICATE_KEY && + err.keyPattern != null && + 'scheduleId' in err.keyPattern && + !('scheduledFor' in err.keyPattern) ); } +/** A duplicate-key error whose conflict is the GLOBAL {capacitySlot} cap index. */ +function isCapacitySlotConflict(error: unknown): boolean { + const err = error as DuplicateKeyError; + return err?.code === DUPLICATE_KEY && err.keyPattern != null && 'capacitySlot' in err.keyPattern; +} + /** A duplicate-key error whose conflict is the per-user {user, slot} cap index. */ function isSlotConflict(error: unknown): boolean { const err = error as DuplicateKeyError; @@ -57,6 +68,12 @@ export interface RecordRunOutcomeParams { error?: string; durationMs?: number; autoDisableAfterFailures: number; + /** Pause CAS token: the resumeSeq the writing segment observed. A pause whose epoch + * no longer matches is dropped, so it cannot demote an already-resumed run. */ + expectResumeSeq?: number; + /** The configRevision this run started under. Terminal bookkeeping / auto-disable is + * skipped when the owner has since edited the schedule (revision moved on). */ + expectConfigRevision?: number; } /** Result of claiming/leasing a schedule: the snapshot plus the fencing token to carry. */ @@ -66,7 +83,32 @@ export interface ScheduleClaim { } /** Outcome of reserving the single-active-run slot for a fired occurrence. */ -export type StartedRunReservation = { run: IScheduleRun } | { conflict: 'duplicate' | 'overlap' }; +export type StartedRunReservation = + | { run: IScheduleRun } + | { conflict: 'duplicate' | 'overlap' | 'slot-taken' }; + +export interface AcquireResumeLeaseParams { + scheduleId: string; + scheduledFor: Date; + /** Unique identity of this resume attempt; every later transition is fenced on it. */ + holder: string; + ttlMs: number; + /** Global capacity slot to claim if the row does not already hold one. */ + capacitySlot: number; +} + +/** + * `acquired` = this attempt owns the occurrence; `held` = another live attempt owns it; + * `overlap` = a DIFFERENT occurrence holds the schedule's active slot; `slot-taken` = + * the global capacity slot collided (allocator should try the next one); `gone` = the + * occurrence no longer exists or is already terminal. + */ +export type ResumeLeaseResult = + | { outcome: 'acquired'; holder: string; resumeSeq: number; capacitySlot?: number } + | { outcome: 'held' } + | { outcome: 'overlap' } + | { outcome: 'slot-taken' } + | { outcome: 'gone' }; /** Outcome of promoting a paused occurrence back into the single active slot. */ export type PromoteRunResult = 'promoted' | 'overlap' | 'missing'; @@ -114,7 +156,18 @@ export type ScheduleMethods = { ) => Promise; insertScheduleRun: (data: Partial) => Promise; reserveStartedRun: (data: Partial) => Promise; - promoteRunToStarted: (scheduleId: string, scheduledFor: Date) => Promise; + getCapacityOccupancy: () => Promise<{ takenSlots: number[]; unslotted: number }>; + acquireResumeLease: (params: AcquireResumeLeaseParams) => Promise; + markResumeClaimed: ( + scheduleId: string, + scheduledFor: Date, + holder: string, + ttlMs: number, + ) => Promise; + commitResumeLease: (scheduleId: string, scheduledFor: Date, holder: string) => Promise; + releaseResumeLease: (scheduleId: string, scheduledFor: Date, holder: string) => Promise; + requestRunAbort: (scheduleId: string, scheduledFor: Date) => Promise; + getRun: (scheduleId: string, scheduledFor: Date) => Promise; setRunFireDetails: ( scheduleId: string, scheduledFor: Date, @@ -235,7 +288,17 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche return Schedule() .findOneAndUpdate( { id, user: userId, deleting: { $ne: true } }, - { $set: { ...update, claimToken: randomUUID() }, ...(unset ? { $unset: unset } : {}) }, + { + $set: { ...update, claimToken: randomUUID() }, + // The ONLY writer of configRevision: an owner edit moves the config + // generation forward atomically with the claim-token rotation, so a run + // that started under the old config can detect it and skip bookkeeping. + // Worker/policy writes (claim, lease, advance, disable, bookkeeping) never + // bump it, and deletion deliberately does not either — a draining run must + // still be able to record its outcome before erasure. + $inc: { configRevision: 1 }, + ...(unset ? { $unset: unset } : {}), + }, { new: true }, ) .lean(); @@ -540,6 +603,12 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche if (isOccurrenceDuplicate(error)) { return { conflict: 'duplicate' }; } + // Checked BEFORE overlap: the global cap index and the per-schedule active + // index are different failures and drive different caller behavior (retry the + // next slot vs skip the occurrence). + if (isCapacitySlotConflict(error)) { + return { conflict: 'slot-taken' }; + } if (isActiveRunConflict(error)) { return { conflict: 'overlap' }; } @@ -547,31 +616,171 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche } } + /** Capacity-slot occupancy for the allocator: which slots are held by `started` + * runs, plus how many legacy rows hold no slot (they shrink the effective cap so + * the bound stays conservative during rollout rather than transiently overshooting). */ + async function getCapacityOccupancy(): Promise<{ takenSlots: number[]; unslotted: number }> { + const rows = await ScheduleRun() + .find({ status: 'started' }) + .select('capacitySlot') + .lean>(); + const takenSlots: number[] = []; + let unslotted = 0; + for (const row of rows) { + if (typeof row.capacitySlot === 'number') { + takenSlots.push(row.capacitySlot); + } else { + unslotted += 1; + } + } + return { takenSlots, unslotted }; + } + /** - * Promotes a paused occurrence back into the single active slot on HITL resume. - * The partial unique index makes overlap atomic: if a newer occurrence is - * already `started`, the update raises a duplicate key and returns 'overlap' - * rather than creating a second concurrent active run. 'missing' means the row - * is no longer `requires_action` (already terminalized/reconciled). + * Acquires the durable RESUMING lease for a HITL resume, atomically: promotes the + * occurrence into the single active slot, claims a global capacity slot, mints the + * next `resumeSeq` (the CAS token that fences stale pause writes), and stamps the + * holder so only this attempt may drive the row. + * + * Adoptable states (the $or): a paused row; a row already `started` with NO holder + * (its pause bookkeeping never landed); or a row whose PRE-CLAIM lease expired (a + * crashed resume that never consumed the approval). A post-claim lease is never + * stolen — that attempt must roll forward, not be re-run. */ - async function promoteRunToStarted( - scheduleId: string, - scheduledFor: Date, - ): Promise { + async function acquireResumeLease(params: AcquireResumeLeaseParams): Promise { + const { scheduleId, scheduledFor, holder, ttlMs, capacitySlot } = params; try { - const result = await ScheduleRun().updateOne( - { scheduleId, scheduledFor, status: 'requires_action' }, - { $set: { status: 'started' } }, - ); - return (result.matchedCount ?? 0) > 0 ? 'promoted' : 'missing'; + const row = await ScheduleRun() + .findOneAndUpdate( + { + scheduleId, + scheduledFor, + $or: [ + { status: 'requires_action' }, + { status: 'started', resumeHolder: { $exists: false } }, + { + status: 'started', + resumeClaimedAt: { $exists: false }, + $expr: { $lt: [{ $ifNull: ['$resumeExpiresAt', new Date(0)] }, '$$NOW'] }, + }, + ], + }, + [ + { + $set: { + status: 'started', + resumeSeq: { $add: [{ $ifNull: ['$resumeSeq', 0] }, 1] }, + resumeHolder: holder, + resumeExpiresAt: { $add: ['$$NOW', ttlMs] }, + capacitySlot, + }, + }, + { $unset: ['resumeClaimedAt'] }, + ], + { new: true }, + ) + .lean(); + if (row == null) { + // Distinguish "another attempt holds it" from "the occurrence is gone/terminal" + // instead of collapsing both to a lossy 'missing'. + const current = await ScheduleRun() + .findOne({ scheduleId, scheduledFor }) + .select('status resumeHolder') + .lean<{ status?: ScheduleRunStatus; resumeHolder?: string }>(); + if (current == null) { + return { outcome: 'gone' }; + } + if (current.resumeHolder != null && current.status === 'started') { + return { outcome: 'held' }; + } + return { outcome: 'gone' }; + } + return { + outcome: 'acquired', + holder, + resumeSeq: row.resumeSeq ?? 0, + capacitySlot: row.capacitySlot, + }; } catch (error) { + if (isCapacitySlotConflict(error)) { + return { outcome: 'slot-taken' }; + } if (isActiveRunConflict(error)) { - return 'overlap'; + return { outcome: 'overlap' }; } throw error; } } + /** Marks the resume lease as having consumed the approval. After this point a crash + * must roll FORWARD (the approval is spent) and the lease is no longer adoptable. */ + async function markResumeClaimed( + scheduleId: string, + scheduledFor: Date, + holder: string, + ttlMs: number, + ): Promise { + const result = await ScheduleRun().updateOne( + { scheduleId, scheduledFor, resumeHolder: holder }, + [{ $set: { resumeClaimedAt: '$$NOW', resumeExpiresAt: { $add: ['$$NOW', ttlMs] } } }], + ); + return (result.matchedCount ?? 0) > 0; + } + + /** Commits the resume: the generation is reconstructed and running, so the row is a + * normal `started` run again and the lease fields are dropped. Holder-fenced. */ + async function commitResumeLease( + scheduleId: string, + scheduledFor: Date, + holder: string, + ): Promise { + const result = await ScheduleRun().updateOne( + { scheduleId, scheduledFor, resumeHolder: holder }, + { $unset: { resumeHolder: 1, resumeExpiresAt: 1, resumeClaimedAt: 1 } }, + ); + return (result.matchedCount ?? 0) > 0; + } + + /** Rolls a PRE-CLAIM resume back to `requires_action`, freeing the capacity slot so + * the approval stays actionable. Refuses once the approval was claimed (post-claim + * must terminalize instead) and is holder-fenced so a stale attempt can't demote. */ + async function releaseResumeLease( + scheduleId: string, + scheduledFor: Date, + holder: string, + ): Promise { + const result = await ScheduleRun().updateOne( + { + scheduleId, + scheduledFor, + resumeHolder: holder, + resumeClaimedAt: { $exists: false }, + }, + { + $set: { status: 'requires_action' }, + $unset: { resumeHolder: 1, resumeExpiresAt: 1, capacitySlot: 1 }, + }, + ); + return (result.matchedCount ?? 0) > 0; + } + + /** Records that an abort was requested WITHOUT freeing the capacity slot: the run + * keeps counting against fireConcurrency until its generation owner confirms + * settlement by writing a terminal outcome. */ + async function requestRunAbort(scheduleId: string, scheduledFor: Date): Promise { + const result = await ScheduleRun().updateOne( + { scheduleId, scheduledFor, status: { $in: ACTIVE_RUN_STATUSES } }, + [{ $set: { abortRequestedAt: { $ifNull: ['$abortRequestedAt', '$$NOW'] } } }], + ); + return (result.matchedCount ?? 0) > 0; + } + + /** The occurrence's run row, or null. Used to read the revision/epoch a run + * started under before writing its outcome. */ + async function getRun(scheduleId: string, scheduledFor: Date): Promise { + return ScheduleRun().findOne({ scheduleId, scheduledFor }).lean(); + } + async function hasActiveRun(scheduleId: string): Promise { const row = await ScheduleRun().findOne({ scheduleId, status: 'started' }).select('_id').lean(); return row != null; @@ -606,8 +815,14 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche // land atomically WITH the count goes in this one update so a crash can't leave // it half-applied: the balance-skip streak resets on ANY non-balance outcome, // and a success clears the failure streak inline (never a lost follow-up). + // CONFIG-REVISION FENCE: a run that started under an older owner config must not + // apply counters (or walk toward auto-disable) against a schedule the owner has + // since edited or re-enabled. Absent on either side disables the fence, so + // pre-existing rows/schedules keep today's behavior instead of wedging. + const revisionFilter = + params.expectConfigRevision != null ? { configRevision: params.expectConfigRevision } : {}; await Schedule().updateOne( - { id: params.scheduleId, countedFor: { $ne: params.scheduledFor } }, + { id: params.scheduleId, countedFor: { $ne: params.scheduledFor }, ...revisionFilter }, { $set: { lastRun, @@ -659,11 +874,19 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche // instead hide the pause from the card until some later terminal outcome. Guard on // a matching active run first so a spoofed scheduleId can't write a card. Keyed on // existence, not modification, so a retried pause still re-affirms the card. + // EPOCH CAS: a pause may only land on the segment that produced it. A resume + // bumps resumeSeq, so a stale requires_action callback from the PREVIOUS segment + // no longer matches and cannot demote the run the resume just promoted. + const epochFilter = + params.expectResumeSeq != null + ? { resumeSeq: params.expectResumeSeq } + : ({} as Record); const activeRun = await ScheduleRun() .findOne({ scheduleId: params.scheduleId, scheduledFor: params.scheduledFor, status: { $in: ['started', 'requires_action'] }, + ...epochFilter, }) .select('_id') .lean(); @@ -683,12 +906,16 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche scheduleId: params.scheduleId, scheduledFor: params.scheduledFor, status: { $in: ['started', 'requires_action'] }, + ...epochFilter, }, { $set: { status: 'requires_action', ...(params.conversationId ? { conversationId: params.conversationId } : {}), }, + // Leaving `started` frees the global capacity slot; the resume claims a + // fresh one from the allocator rather than re-adopting a possibly-taken slot. + $unset: { capacitySlot: 1 }, }, ); return; @@ -710,6 +937,16 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche ...(params.error ? { error: params.error } : {}), ...(params.durationMs != null ? { durationMs: params.durationMs } : {}), }, + // SETTLEMENT: a terminal outcome is the generation owner confirming the run + // actually stopped, so this is the ONLY place the global capacity slot is + // released. An abort request alone does not free it (see requestRunAbort). + // Any in-flight resume lease ends with the run. + $unset: { + capacitySlot: 1, + resumeHolder: 1, + resumeExpiresAt: 1, + resumeClaimedAt: 1, + }, }, ); // No-match guard: never touch schedule bookkeeping without a matching run @@ -1007,7 +1244,13 @@ export function createScheduleMethods(mongoose: typeof import('mongoose')): Sche disableSchedule, insertScheduleRun, reserveStartedRun, - promoteRunToStarted, + getCapacityOccupancy, + acquireResumeLease, + markResumeClaimed, + commitResumeLease, + releaseResumeLease, + requestRunAbort, + getRun, setRunFireDetails, hasActiveRun, countActiveRuns, diff --git a/packages/data-schemas/src/schema/schedule.ts b/packages/data-schemas/src/schema/schedule.ts index 7442a34a63..cdf6196283 100644 --- a/packages/data-schemas/src/schema/schedule.ts +++ b/packages/data-schemas/src/schema/schedule.ts @@ -96,6 +96,17 @@ const scheduleSchema: Schema = new Schema( claimToken: { type: String, }, + /** + * Owner-config generation. Bumped ONLY by an owner edit (updateScheduleById), + * atomically with the claimToken rotation. Distinct from claimToken, which also + * rotates on every lease acquisition. A run captures this at claim time so its + * terminal bookkeeping / auto-disable cannot act on config it never ran under. + */ + configRevision: { + type: Number, + default: 0, + min: 0, + }, /** * Soft-delete marker. A delete disables + marks the schedule `deleting` (so * it is hidden from the owner and never re-claimed) and aborts in-flight diff --git a/packages/data-schemas/src/schema/scheduleRun.ts b/packages/data-schemas/src/schema/scheduleRun.ts index d45bd8c44b..26348fe221 100644 --- a/packages/data-schemas/src/schema/scheduleRun.ts +++ b/packages/data-schemas/src/schema/scheduleRun.ts @@ -59,6 +59,44 @@ const scheduleRunSchema: Schema = new Schema( bookkept: { type: Boolean, }, + /** Monotonic per-occurrence segment counter. Incremented by every resume + * reservation; a pause write must CAS on the epoch its segment observed, so a + * stale `requires_action` callback can never demote an already-resumed run. */ + resumeSeq: { + type: Number, + min: 0, + }, + /** Identity of the in-flight resume attempt. Presence means RESUMING: the run is + * `started` but its generation has not been reconstructed yet. */ + resumeHolder: { + type: String, + }, + /** Deadline for the current resume phase, so a crashed resume is reclaimable. */ + resumeExpiresAt: { + type: Date, + }, + /** Set once the approval claim succeeded; discriminates pre-claim (rollback-safe) + * from post-claim (must roll forward) recovery. */ + resumeClaimedAt: { + type: Date, + }, + /** Global concurrency slot held while `started`. The unique partial index below + * turns fireConcurrency into a DB-enforced bound instead of a racy count. */ + capacitySlot: { + type: Number, + min: 0, + }, + /** When an abort was requested. The run keeps holding its capacity slot until the + * generation owner confirms settlement, so capacity is never freed early. */ + abortRequestedAt: { + type: Date, + }, + /** The schedule's configRevision at claim time. Fences terminal bookkeeping and + * auto-disable from owner edits/re-enables that landed after this run started. */ + configRevision: { + type: Number, + min: 0, + }, }, { timestamps: true, @@ -74,6 +112,18 @@ scheduleRunSchema.index( { scheduleId: 1 }, { unique: true, partialFilterExpression: { status: 'started' } }, ); +// GLOBAL fireConcurrency, enforced by the DB rather than a read-then-compare count. +// Every transition into `started` (fire insert, resume promotion) claims a slot in +// [0, fireConcurrency) in the SAME write; a duplicate slot is rejected atomically, so +// two concurrent admissions of DIFFERENT schedules can never both pass a cap-1 check. +// Partial + $exists so legacy slotless rows (written before this field) never collide. +scheduleRunSchema.index( + { capacitySlot: 1 }, + { + unique: true, + partialFilterExpression: { status: 'started', capacitySlot: { $exists: true } }, + }, +); scheduleRunSchema.index({ scheduleId: 1, firedAt: -1 }); // Reconciliation sweeps by status; keeps `started` (capacity) fetch cheap and // prevents long-lived `requires_action` rows from starving the scan. diff --git a/packages/data-schemas/src/schema/user.ts b/packages/data-schemas/src/schema/user.ts index 706c73403b..b992eecf3c 100644 --- a/packages/data-schemas/src/schema/user.ts +++ b/packages/data-schemas/src/schema/user.ts @@ -157,6 +157,15 @@ const userSchema: Schema = new Schema( of: Boolean, default: () => new Map(), }, + /** + * Durable account-deletion barrier. Set BEFORE the deletion cascade quiesces + * anything, so every later scheduling admission (create/update/run-now, engine + * claim, loopback fire) can refuse this user. Absent means live, so existing + * documents need no migration. + */ + deletionRequestedAt: { + type: Date, + }, /** Field for external source identification (for consistency with TPrincipal schema) */ idOnTheSource: { type: String, diff --git a/packages/data-schemas/src/types/schedule.ts b/packages/data-schemas/src/types/schedule.ts index 94b4701d54..28f2203d47 100644 --- a/packages/data-schemas/src/types/schedule.ts +++ b/packages/data-schemas/src/types/schedule.ts @@ -25,6 +25,8 @@ export interface ISchedule { leaseUntil?: Date; leaseBy?: string; claimToken?: string; + /** Owner-config generation; bumped only by an owner edit. */ + configRevision?: number; deleting?: boolean; slot?: number; lastRun?: { @@ -57,6 +59,19 @@ export interface IScheduleRun { droppedFileIds?: string[]; durationMs?: number; bookkept?: boolean; + /** Monotonic segment counter; the CAS token a pause write must match. */ + resumeSeq?: number; + /** Identity of the in-flight resume attempt (presence == RESUMING). */ + resumeHolder?: string; + resumeExpiresAt?: Date; + /** Set once the approval claim succeeded (pre-claim vs post-claim recovery). */ + resumeClaimedAt?: Date; + /** Global concurrency slot held while `started`. */ + capacitySlot?: number; + /** When an abort was requested; capacity is held until settlement is confirmed. */ + abortRequestedAt?: Date; + /** The schedule's configRevision at claim time. */ + configRevision?: number; createdAt?: Date; updatedAt?: Date; } diff --git a/packages/data-schemas/src/types/user.ts b/packages/data-schemas/src/types/user.ts index 5b302ad275..9ac943e5fa 100644 --- a/packages/data-schemas/src/types/user.ts +++ b/packages/data-schemas/src/types/user.ts @@ -57,6 +57,8 @@ export interface IUser extends Document { skillStates?: Record; createdAt?: Date; updatedAt?: Date; + /** Set when account deletion begins; durably blocks all new scheduling for this user. */ + deletionRequestedAt?: Date; /** Field for external source identification (for consistency with TPrincipal schema) */ idOnTheSource?: string; tenantId?: string;