diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index 838009e178..d05334fcb1 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -25,6 +25,10 @@ const mockGenerationJobManager = { const mockSaveMessage = jest.fn(); +const mockRecordScheduleOutcome = jest.fn(); +const mockBeginScheduledStop = jest.fn(); +const mockAcknowledgeScheduledStopPersistence = jest.fn(); + jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), logger: mockLogger, @@ -45,6 +49,13 @@ jest.mock('~/models', () => ({ saveMessage: (...args) => mockSaveMessage(...args), })); +jest.mock('~/server/services/Schedules', () => ({ + recordScheduleOutcome: (...args) => mockRecordScheduleOutcome(...args), + beginScheduledStop: (...args) => mockBeginScheduledStop(...args), + acknowledgeScheduledStopPersistence: (...args) => + mockAcknowledgeScheduledStopPersistence(...args), +})); + jest.mock('~/server/middleware', () => ({ uaParser: (req, res, next) => next(), checkBan: (req, res, next) => next(), @@ -86,6 +97,12 @@ describe('Agent Abort Endpoint', () => { mockGenerationJobManager.getActiveJobIdsForUser.mockReset(); mockSaveMessage.mockReset(); mockSaveMessage.mockImplementation(async (_context, message) => message); + mockRecordScheduleOutcome.mockReset(); + mockRecordScheduleOutcome.mockResolvedValue(true); + mockBeginScheduledStop.mockReset(); + mockBeginScheduledStop.mockResolvedValue(true); + mockAcknowledgeScheduledStopPersistence.mockReset(); + mockAcknowledgeScheduledStopPersistence.mockResolvedValue(undefined); }); describe('POST /chat/abort', () => { @@ -846,5 +863,110 @@ describe('Agent Abort Endpoint', () => { }); }); }); + + describe('Scheduled Stop persistence protocol', () => { + const scheduledJob = { + status: 'running', + createdAt: 111, + metadata: { + userId: 'test-user-123', + generationProtocolVersion: 2, + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + conversationId: 'conv-1', + }, + }; + + it('stamps the Stop before signalling abort, then acknowledges after persistence', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(scheduledJob); + mockGenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: { conversationId: 'conv-1' }, + content: [], + text: '', + }); + + const response = await request(app) + .post('/api/agents/chat/abort') + .set('X-LibreChat-Generation-Protocol', '2') + .send({ conversationId: 'conv-1', generationProtocolVersion: 2 }); + + expect(response.status).toBe(200); + expect(mockBeginScheduledStop).toHaveBeenCalledWith({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + }); + // Stamp BEFORE the abort signal; acknowledgement AFTER it (persistence done). + expect(mockBeginScheduledStop.mock.invocationCallOrder[0]).toBeLessThan( + mockGenerationJobManager.abortJob.mock.invocationCallOrder[0], + ); + expect(mockAcknowledgeScheduledStopPersistence).toHaveBeenCalledWith({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + }); + expect(mockGenerationJobManager.abortJob.mock.invocationCallOrder[0]).toBeLessThan( + mockAcknowledgeScheduledStopPersistence.mock.invocationCallOrder[0], + ); + }); + + it('returns STOP_IN_PROGRESS and never signals a second abort when a Stop is already live', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(scheduledJob); + mockBeginScheduledStop.mockResolvedValue('in_progress'); + + const response = await request(app) + .post('/api/agents/chat/abort') + .set('X-LibreChat-Generation-Protocol', '2') + .send({ conversationId: 'conv-1', generationProtocolVersion: 2 }); + + expect(response.status).toBe(409); + expect(response.body.code).toBe('STOP_IN_PROGRESS'); + expect(mockGenerationJobManager.abortJob).not.toHaveBeenCalled(); + expect(mockAcknowledgeScheduledStopPersistence).not.toHaveBeenCalled(); + }); + + it('does NOT acknowledge when the abort persistence failed (run stays preserved)', async () => { + mockGenerationJobManager.getJob.mockResolvedValue(scheduledJob); + mockGenerationJobManager.abortJob.mockResolvedValue({ + success: true, + persistenceFailed: true, + jobData: { conversationId: 'conv-1' }, + content: [], + text: '', + }); + + const response = await request(app) + .post('/api/agents/chat/abort') + .set('X-LibreChat-Generation-Protocol', '2') + .send({ conversationId: 'conv-1', generationProtocolVersion: 2 }); + + expect(response.status).toBe(200); + expect(response.body.persistenceFailed).toBe(true); + expect(mockBeginScheduledStop).toHaveBeenCalled(); + expect(mockAcknowledgeScheduledStopPersistence).not.toHaveBeenCalled(); + }); + + it('leaves non-scheduled aborts untouched by the Stop protocol', async () => { + mockGenerationJobManager.getJob.mockResolvedValue({ + status: 'running', + createdAt: 1, + metadata: { userId: 'test-user-123', generationProtocolVersion: 2 }, + }); + mockGenerationJobManager.abortJob.mockResolvedValue({ + success: true, + jobData: { conversationId: 'conv-x' }, + content: [], + text: '', + }); + + const response = await request(app) + .post('/api/agents/chat/abort') + .set('X-LibreChat-Generation-Protocol', '2') + .send({ conversationId: 'conv-x', generationProtocolVersion: 2 }); + + expect(response.status).toBe(200); + expect(mockBeginScheduledStop).not.toHaveBeenCalled(); + expect(mockAcknowledgeScheduledStopPersistence).not.toHaveBeenCalled(); + }); + }); }); }); diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index e6ba230f94..baf380c667 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -38,7 +38,11 @@ const { negotiateExistingGenerationProtocol, } = require('~/server/controllers/agents/protocol'); const { saveMessage } = require('~/models'); -const { recordScheduleOutcome } = require('~/server/services/Schedules'); +const { + recordScheduleOutcome, + beginScheduledStop, + acknowledgeScheduledStopPersistence, +} = require('~/server/services/Schedules'); const responses = require('./responses'); const openai = require('./openai'); const { v1 } = require('./v1'); @@ -673,6 +677,27 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { throwOnError: true, }) : undefined; + // Stamp a scheduled run's Stop BEFORE signalling the abort. `abortJob` flips the job + // to `aborted` immediately, then runs the partial-message/checkpoint persistence in + // `beforePublish`. Without this stamp the owner settlement, reconciliation, and + // schedule/account deletion could observe `aborted` and terminalize/erase the run + // mid-write. The stamp is serialized: a fresh Stop already owning it means another + // request is persisting, so we must not signal a second abort. + const stopScheduleId = job.metadata?.scheduleId; + const stopScheduledFor = job.metadata?.scheduledFor; + const isScheduledStop = stopScheduleId != null && stopScheduledFor != null; + let scheduledStopStamped = false; + if (isScheduledStop) { + const stopStamp = await beginScheduledStop({ + scheduleId: stopScheduleId, + scheduledFor: stopScheduledFor, + }); + if (stopStamp === 'in_progress') { + res.set('Retry-After', '1'); + return res.status(409).json({ code: 'STOP_IN_PROGRESS', generationProtocolVersion }); + } + scheduledStopStamped = stopStamp === true; + } const abortResult = await GenerationJobManager.abortJob(jobStreamId, { expectedCreatedAt: job.createdAt, transformAbortContent: (content, abortJobData) => { @@ -811,6 +836,18 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { } }, }); + // The abort did not land (replaced/still-active/already-settled), so no persistence + // is in flight: release the Stop barrier we armed rather than deferring this + // occurrence's settlement for the full stale-owner window. A retry or a replacement + // generation re-stamps its own; the acknowledgement is fenced to the occurrence, not + // a generation, so it cannot settle a successor through its predecessor. + if (scheduledStopStamped && !abortResult.success) { + await acknowledgeScheduledStopPersistence({ + scheduleId: stopScheduleId, + scheduledFor: stopScheduledFor, + }); + scheduledStopStamped = false; + } if (abortResult.failureReason === 'generation_replaced') { return res.status(409).json({ code: 'RUN_REPLACED', generationProtocolVersion }); } @@ -862,11 +899,28 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { abortResultResponseMessageId: abortResult.jobData?.responseMessageId, }); + // `beforePublish` has run: its partial-message and checkpoint writes have either + // landed or failed. Acknowledge the Stop ONLY on success — that releases the owner's + // settlement barrier so the run can terminalize. On a persistence failure we leave + // the barrier unresolved: the run stays preserved (client retries; the stale-owner + // timeout is the bounded recovery) rather than settling over an incomplete write. + if (scheduledStopStamped && !abortResult.persistenceFailed) { + await acknowledgeScheduledStopPersistence({ + scheduleId: stopScheduleId, + scheduledFor: stopScheduledFor, + }); + } + // A paused generation has no live provider owner left to report the stop. // Persist its scheduled occurrence here, after abortJob's required partial - // response/checkpoint work, while running generations continue to settle from - // their owning request/resume controller. - if (job.status === 'requires_action' && job.metadata?.scheduleId) { + // response/checkpoint work AND its acknowledgement above, while running generations + // continue to settle from their owning request/resume controller. A failed + // persistence skips settlement so the incomplete run is not terminalized. + if ( + job.status === 'requires_action' && + job.metadata?.scheduleId && + !abortResult.persistenceFailed + ) { await recordScheduleOutcome({ scheduleId: job.metadata.scheduleId, scheduledFor: job.metadata.scheduledFor, diff --git a/api/server/services/Schedules/index.js b/api/server/services/Schedules/index.js index 5724811ff5..5944ad2c1b 100644 --- a/api/server/services/Schedules/index.js +++ b/api/server/services/Schedules/index.js @@ -83,6 +83,8 @@ module.exports = { getLimits: invoke('getLimits'), fireScheduleNow: invoke('fireScheduleNow'), recordScheduleOutcome: invoke('recordScheduleOutcome'), + beginScheduledStop: invoke('beginScheduledStop'), + acknowledgeScheduledStopPersistence: invoke('acknowledgeScheduledStopPersistence'), claimScheduleResume: invoke('claimScheduleResume'), releaseScheduleResumeClaim: invoke('releaseScheduleResumeClaim'), finalizeScheduleResumeClaim: invoke('finalizeScheduleResumeClaim'), diff --git a/packages/api/src/schedules/service.spec.ts b/packages/api/src/schedules/service.spec.ts index 2ca606165b..eea132bf3b 100644 --- a/packages/api/src/schedules/service.spec.ts +++ b/packages/api/src/schedules/service.spec.ts @@ -51,6 +51,8 @@ function makeService( getActiveRunsForUser, countActiveRuns: jest.fn(async () => 0), requestRunAbort: jest.fn(async () => true), + getScheduleRunAbortState: jest.fn(async () => null), + markRunAbortPersisted: jest.fn(async () => undefined), recordRunOutcome, }; const deps = { @@ -68,6 +70,8 @@ function makeService( return createSchedulesService(deps, { drainTimeoutMs: 400, drainPollMs: 25, + stopBarrierTimeoutMs: 400, + stopBarrierPollMs: 25, }); } @@ -608,6 +612,124 @@ describe('erase-on-settle', () => { }); }); +describe('interactive Stop persistence barrier', () => { + function outcomeService() { + const service = makeService(jest.fn, [string]>().mockResolvedValue([])); + const methods = service.engineDeps.methods as unknown as { + getScheduleById: jest.Mock; + recordRunOutcome: jest.Mock; + eraseScheduleIfDrained: jest.Mock; + getScheduleRunAbortState: jest.Mock; + requestRunAbort: jest.Mock; + markRunAbortPersisted: jest.Mock; + }; + methods.getScheduleById = jest.fn(async () => null); + methods.recordRunOutcome = jest.fn(async () => undefined); + methods.eraseScheduleIfDrained = jest.fn(async () => true); + return { service, methods }; + } + + it('beginScheduledStop stamps a serialized stop; acknowledge marks persistence', async () => { + const { service, methods } = outcomeService(); + methods.requestRunAbort = jest.fn(async () => 'in_progress' as const); + methods.markRunAbortPersisted = jest.fn(async () => undefined); + + const stamp = await service.beginScheduledStop({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + }); + expect(stamp).toBe('in_progress'); + expect(methods.requestRunAbort).toHaveBeenCalledWith('s1', expect.any(Date), 'stop'); + + await service.acknowledgeScheduledStopPersistence({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + }); + expect(methods.markRunAbortPersisted).toHaveBeenCalledWith('s1', expect.any(Date)); + }); + + it('owner terminal settlement waits until the Stop persistence is acknowledged', async () => { + const { service, methods } = outcomeService(); + const fresh = new Date(); + let reads = 0; + // Unresolved fresh stop for the first two reads, then acknowledged. + methods.getScheduleRunAbortState = jest.fn(async () => { + reads += 1; + return { + status: 'started', + abortSource: 'stop', + abortRequestedAt: fresh, + ...(reads > 2 ? { abortPersistedAt: new Date() } : {}), + }; + }); + + await service.recordScheduleOutcome({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + status: 'success', + }); + + // Polled until the acknowledgement landed, and only THEN recorded the outcome. + expect(reads).toBeGreaterThanOrEqual(3); + expect(methods.recordRunOutcome).toHaveBeenCalled(); + expect(methods.getScheduleRunAbortState.mock.invocationCallOrder[0]).toBeLessThan( + methods.recordRunOutcome.mock.invocationCallOrder[0], + ); + }); + + it('does not wait on a non-stop abort source', async () => { + const { service, methods } = outcomeService(); + methods.getScheduleRunAbortState = jest.fn(async () => ({ + status: 'started', + abortSource: 'deletion', + abortRequestedAt: new Date(), + })); + + await service.recordScheduleOutcome({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + status: 'interrupted', + }); + + expect(methods.getScheduleRunAbortState).toHaveBeenCalledTimes(1); + expect(methods.recordRunOutcome).toHaveBeenCalled(); + }); + + it('does not wait on a STALE stop stamp (route presumed dead)', async () => { + const { service, methods } = outcomeService(); + // Older than ABORT_OWNER_PRESUMED_ALIVE_MS: the owner presumes the route dead. + const stale = new Date(Date.now() - 31 * 60_000); + methods.getScheduleRunAbortState = jest.fn(async () => ({ + status: 'started', + abortSource: 'stop', + abortRequestedAt: stale, + })); + + await service.recordScheduleOutcome({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + status: 'error', + }); + + expect(methods.getScheduleRunAbortState).toHaveBeenCalledTimes(1); + expect(methods.recordRunOutcome).toHaveBeenCalled(); + }); + + it('does not gate a pause (requires_action) on the Stop barrier', async () => { + const { service, methods } = outcomeService(); + methods.getScheduleRunAbortState = jest.fn(async () => null); + + await service.recordScheduleOutcome({ + scheduleId: 's1', + scheduledFor: '2026-01-01T00:00:00.000Z', + status: 'requires_action', + }); + + // A pause is not a settlement, so the barrier is never consulted. + expect(methods.getScheduleRunAbortState).not.toHaveBeenCalled(); + }); +}); + describe('attachment hold renewal', () => { it('renews the bounded upload hold at fire preflight, best-effort', async () => { const service = makeService(jest.fn, [string]>().mockResolvedValue([])); diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index d06f301e83..11384da6f0 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -174,6 +174,26 @@ export interface SchedulesService { limits: ScheduleLimits, ) => Promise; recordScheduleOutcome: (input: RecordScheduleOutcomeInput) => Promise; + /** + * Stamps a scheduled run's interactive Stop BEFORE the abort is signalled, so the owner + * settlement barrier, reconciliation, and schedule/account deletion hold off settling or + * erasing until {@link acknowledgeScheduledStopPersistence}. Serialized: `'in_progress'` + * means a fresh Stop already owns the stamp and the caller must not signal a second abort; + * `false` means there is no active run to stop. + */ + beginScheduledStop: (input: { + scheduleId: string; + scheduledFor: string | Date; + }) => Promise; + /** + * Releases the Stop settlement barrier once the route's partial-message/checkpoint writes + * have landed. Call ONLY after persistence succeeds — a failed persistence must leave the + * barrier unresolved so the run stays preserved and recovers via the stale-owner timeout. + */ + acknowledgeScheduledStopPersistence: (input: { + scheduleId: string; + scheduledFor: string | Date; + }) => Promise; /** Re-enters a paused occurrence into the DB-enforced global/same-schedule * capacity set before its approval job is allowed to resume. */ claimScheduleResume: ( @@ -243,6 +263,8 @@ export interface SchedulesService { export interface ScheduleServiceTimings { drainTimeoutMs?: number; drainPollMs?: number; + stopBarrierTimeoutMs?: number; + stopBarrierPollMs?: number; } /** @@ -335,6 +357,13 @@ export function createSchedulesService( // never blocks indefinitely on an unreachable peer-worker run. const QUIESCE_DRAIN_TIMEOUT_MS = timings?.drainTimeoutMs ?? 10 * 1000; const QUIESCE_DRAIN_POLL_MS = timings?.drainPollMs ?? 250; + // The owner's terminal settlement waits (briefly, bounded) for an interactive Stop to + // acknowledge its partial-message/checkpoint persistence before releasing capacity or + // erasing evidence. Short: the Stop route's persistence is a couple of writes; if it + // does not land in this window the stale-owner timeout (ABORT_OWNER_PRESUMED_ALIVE_MS) + // takes over as the bounded recovery, so the settlement never blocks indefinitely. + const STOP_BARRIER_TIMEOUT_MS = timings?.stopBarrierTimeoutMs ?? 5 * 1000; + const STOP_BARRIER_POLL_MS = timings?.stopBarrierPollMs ?? 100; /** * Whether a refill would top up this zero-credit balance record right now, @@ -726,6 +755,34 @@ export function createSchedulesService( * Mongo failure here is RETRIED (bounded) before giving up, and the failure is * surfaced to the caller (returns false) so it can keep the job when it matters. */ + /** + * OWNER-SIDE Stop barrier. An interactive Stop stamps `abortSource: 'stop'` before it + * signals the abort, and only acknowledges (`abortPersistedAt`) once its partial-message + * and checkpoint writes have landed. A terminal settlement here in between would release + * the run's capacity — and let a concurrent schedule/account deletion erase the data — + * while that write is still in flight. Wait (bounded) for the acknowledgement. A resolved + * marker, a non-stop source, or a stamp gone stale (its route presumed dead) proceeds + * immediately, so the stale-owner timeout remains the bounded recovery path. + */ + async function waitForStopPersistence(scheduleId: string, scheduledFor: Date): Promise { + const deadline = Date.now() + STOP_BARRIER_TIMEOUT_MS; + for (;;) { + const state = await methods.getScheduleRunAbortState(scheduleId, scheduledFor); + if ( + state == null || + state.abortSource !== 'stop' || + state.abortPersistedAt != null || + !hasAbortInFlight(state, Date.now()) + ) { + return; + } + if (Date.now() >= deadline) { + return; + } + await new Promise((resolve) => setTimeout(resolve, STOP_BARRIER_POLL_MS)); + } + } + async function recordScheduleOutcome({ scheduleId, scheduledFor, @@ -740,6 +797,10 @@ export function createSchedulesService( return true; } const terminal = status !== 'requires_action'; + if (terminal) { + // Honor an in-flight interactive Stop's persistence before terminalizing. + await waitForStopPersistence(scheduleId, new Date(scheduledFor)); + } if (terminal && streamId && jobCreatedAt != null) { try { await GenerationJobManager.updateMetadata( @@ -828,6 +889,32 @@ export function createSchedulesService( return false; } + async function beginScheduledStop({ + scheduleId, + scheduledFor, + }: { + scheduleId: string; + scheduledFor: string | Date; + }): Promise { + if (!scheduleId || !scheduledFor) { + return false; + } + return methods.requestRunAbort(scheduleId, new Date(scheduledFor), 'stop'); + } + + async function acknowledgeScheduledStopPersistence({ + scheduleId, + scheduledFor, + }: { + scheduleId: string; + scheduledFor: string | Date; + }): Promise { + if (!scheduleId || !scheduledFor) { + return; + } + await methods.markRunAbortPersisted(scheduleId, new Date(scheduledFor)); + } + async function isScheduleLive( scheduleId: string, expectedConfigRevision?: number, @@ -1451,6 +1538,8 @@ export function createSchedulesService( engineDeps, fireScheduleNow, recordScheduleOutcome, + beginScheduledStop, + acknowledgeScheduledStopPersistence, claimScheduleResume, releaseScheduleResumeClaim, finalizeScheduleResumeClaim,