From b34bbd83eff6595686303a9917ad4c2ff3712fe5 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 3 Aug 2026 10:28:02 -0400 Subject: [PATCH] =?UTF-8?q?fix:=20lease-token=20lifecycle=20fences=20?= =?UTF-8?q?=E2=80=94=20same-generation=20isolation,=20admission=20fence,?= =?UTF-8?q?=20undelivered-Stop=20retention,=20legacy=20abort=20enrollment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fifth external review round; four P1s, handled as the requested consolidated lifecycle-fence pass. - Lease tokens. Marker fields were (streamId, createdAt), shared by every contender on the same generation — completion and Stop, or two racing Stops — so a losing contender's cleanup erased the winner's still-live marker. Every registrant now carries a unique lease token in the field and may only ever clear its own lease; unqualified legacy clears cannot touch qualified entries. - Admission fence. Authentication can pass before the deletion barrier goes up, and the durable createJob is several async steps later — a deletion quiesce in that window saw neither an active job nor a marker and could cascade before the admitted request created its job. The controller now registers an admission lease and THEN rereads the deletion barrier: the ordering guarantees either this request observes the barrier (403, lease released, slot/claim cleanup) or the quiesce observes the lease and defers. Held until createJob is durable; released on every refusal and initialization-error path. Fail closed when the lease itself cannot be registered (503 retryable). - Undelivered-Stop retention. abortJob released its lease in a finally even when delivery AND publication had provably failed — the job reads terminal (invisible to active-set scans), a user Stop writes no durable abort fence, and the remote owner keeps generating and will persist its abort-catch writes whenever the signal finally lands. The lease is now retained in exactly that case, and each resignal attempt heartbeats a fresh lease so the fence outlives the TTL for as long as delivery is still being driven. The abort-winning turn's own loser-side pending saves are additionally fenced in the controller catch (best-effort — those writes are already in flight). - Legacy abort enrollment. abortMiddleware (assistants abort route fallback for non-assistants endpoints) awaited abortJob and then spent usage and saved the stopped response AFTER the abort's lease was released. Both writes now run inside `beforePublish`, between the abort CAS and publication, covered by the same lease as every other abort. Barrier tests, each verified to fail without its fix: same-generation lease isolation (store), racing two-Stop loser cleanup (manager, stale-read forced CAS race), undelivered-Stop lease retention, resignal heartbeat, admission refusal with lease-before-reread ordering plus release-on-durable-create, and legacy-abort persistence inside beforePublish. --- .../agents/__tests__/jobReplacement.spec.js | 1 + .../request.partialDisconnect.spec.js | 1 + .../__tests__/request.resumeMetadata.spec.js | 97 ++++++++++++- .../agents/__tests__/resume.spec.js | 5 + api/server/controllers/agents/request.js | 79 +++++++++++ api/server/controllers/agents/resume.js | 8 ++ .../__tests__/abortMiddleware.spec.js | 107 +++++++++++++++ api/server/middleware/abortMiddleware.js | 127 ++++++++++-------- .../api/src/stream/GenerationJobManager.ts | 25 +++- .../__tests__/abortSignalHonesty.spec.ts | 95 ++++++++++++- .../__tests__/userFinalizations.spec.ts | 12 ++ .../implementations/InMemoryJobStore.ts | 9 +- .../stream/implementations/RedisJobStore.ts | 6 +- .../api/src/stream/interfaces/IJobStore.ts | 28 +++- 14 files changed, 531 insertions(+), 69 deletions(-) create mode 100644 api/server/middleware/__tests__/abortMiddleware.spec.js diff --git a/api/server/controllers/agents/__tests__/jobReplacement.spec.js b/api/server/controllers/agents/__tests__/jobReplacement.spec.js index cc000f6bf1..a9de913fc8 100644 --- a/api/server/controllers/agents/__tests__/jobReplacement.spec.js +++ b/api/server/controllers/agents/__tests__/jobReplacement.spec.js @@ -59,6 +59,7 @@ jest.mock('@librechat/api', () => ({ })); jest.mock('~/models', () => ({ + isUserDeleting: jest.fn(async () => false), saveMessage: (...args) => mockSaveMessage(...args), })); diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 9cca7e0f2c..d8d5a6b0ab 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -127,6 +127,7 @@ jest.mock('~/cache', () => ({ })); jest.mock('~/models', () => ({ + isUserDeleting: jest.fn(async () => false), saveMessage: (...args) => mockSaveMessage(...args), getMessages: (...args) => mockGetMessages(...args), getConvo: (...args) => mockGetConvo(...args), diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index a379036de6..810ccc3ae0 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -228,6 +228,7 @@ jest.mock('~/server/services/Schedules', () => ({ })); jest.mock('~/models', () => ({ + isUserDeleting: jest.fn(async () => false), saveMessage: (...args) => mockSaveMessage(...args), getMessages: (...args) => mockGetMessages(...args), getConvo: (...args) => mockGetConvo(...args), @@ -3050,18 +3051,112 @@ describe('ResumableAgentController resume metadata', () => { await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); await nextTick(); - // Generation-qualified: a predecessor's late clear must not drop this marker. + // Generation- and lease-qualified: a predecessor's (or a losing same-generation + // contender's) late clear must not drop this marker. The FIRST registration is + // the admission lease covering the auth-to-createJob window. + expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( + 'user-123', + 'conversation-123', + undefined, + undefined, + expect.stringMatching(/^admission-/), + ); expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( 'user-123', 'conversation-123', undefined, 1000, + expect.any(String), ); expect( mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], ).toBeLessThan(mockGenerationJobManager.claimTerminalJob.mock.invocationCallOrder[0]); }); + it('refuses admission when the deletion barrier is up, lease registered FIRST', async () => { + // The order is the fence: lease durable BEFORE the barrier reread, so either this + // request sees the barrier (and refuses) or the deletion quiesce sees the lease + // (and defers). Without it, a request authenticated before the barrier could + // create its job AFTER the cascade swept. + const models = require('~/models'); + models.isUserDeleting.mockResolvedValueOnce(true); + const initializeClient = jest.fn(); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Racing a deletion.', + messageId: 'user-msg', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(res.status).toHaveBeenCalledWith(403); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + expect(initializeClient).not.toHaveBeenCalled(); + const admissionRegister = mockGenerationJobManager.registerUserFinalization.mock.calls.find( + (call) => String(call[4]).startsWith('admission-'), + ); + expect(admissionRegister).toBeDefined(); + // Lease before barrier read. + expect( + mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], + ).toBeLessThan(models.isUserDeleting.mock.invocationCallOrder[0]); + // The refusal releases its own lease. + expect(mockGenerationJobManager.clearUserFinalization).toHaveBeenCalledWith( + 'user-123', + 'conversation-123', + undefined, + undefined, + admissionRegister[4], + ); + }); + + it('releases the admission lease once createJob is durable', async () => { + mockGenerationJobManager.getJob.mockResolvedValue({ createdAt: 1000, status: 'running' }); + const client = { + options: {}, + sendMessage: jest.fn().mockResolvedValue({ + userMessage: { + messageId: 'user-msg', + parentMessageId: 'parent-msg', + conversationId: 'conversation-123', + text: 'Follow up.', + }, + responseMessage: { messageId: 'resp-msg', conversationId: 'conversation-123' }, + }), + }; + const initializeClient = jest.fn().mockResolvedValue({ client }); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Follow up.', + messageId: 'user-msg', + parentMessageId: 'parent-msg', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + + await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); + await nextTick(); + + const admissionRegister = mockGenerationJobManager.registerUserFinalization.mock.calls.find( + (call) => String(call[4]).startsWith('admission-'), + ); + expect(admissionRegister).toBeDefined(); + const admissionClear = mockGenerationJobManager.clearUserFinalization.mock.calls.find( + (call) => call[4] === admissionRegister[4], + ); + // Held exactly until the job is durable (active-set visible), then released. + expect(admissionClear).toBeDefined(); + }); + it('fails CLOSED when the finalization marker cannot be registered: no terminal CAS at all', async () => { // A transient marker-store failure must not reopen the deletion race the marker // fences: neither the complete-claim NOR the error-path completeJob may run diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index 2d10f36568..759696f1fa 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -2196,12 +2196,16 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { // first-turn title) are still ahead — without the marker a deletion quiesce in // that window sees neither an active job nor a fence. // Generation-qualified: a predecessor's late clear must not drop this marker. + // Lease-qualified: a losing same-generation contender (a racing Stop) must + // only ever release its own lease, never this one. expect(mockGenerationJobManager.registerUserFinalization).toHaveBeenCalledWith( USER_ID, CONVO_ID, TENANT_ID, 1000, + expect.any(String), ); + const resumeLeaseId = mockGenerationJobManager.registerUserFinalization.mock.calls[0][4]; expect( mockGenerationJobManager.registerUserFinalization.mock.invocationCallOrder[0], ).toBeLessThan(mockGenerationJobManager.claimTerminalJob.mock.invocationCallOrder[0]); @@ -2210,6 +2214,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { CONVO_ID, TENANT_ID, 1000, + resumeLeaseId, ); // Cleared only after the persistence phase (the finish of the terminal claim), // never between the CAS and the save. diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index a6ab43208c..70ebe2b349 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -915,6 +915,71 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } startupTelemetry?.mark('request_admitted'); + // ADMISSION FENCE against account deletion. Authentication can pass BEFORE the + // deletion barrier goes up, and the durable createJob below is still several async + // steps away — in that window a deletion quiesce sees neither an active job nor a + // marker and can cascade while this admitted request goes on to create one. The + // ORDER is the fence: the lease is durable BEFORE the barrier reread, so either we + // observe their barrier (and refuse) or their quiesce observes our lease (and + // defers). Held until createJob is durable — from then on the job itself is + // deletion-visible. FAIL CLOSED on an unregistrable lease: nothing has been + // admitted into generation yet, so refusing is free and plainly retryable. + const admissionLeaseId = `admission-${crypto.randomUUID()}`; + let admissionLeaseHeld = false; + const releaseAdmissionLease = () => { + if (!admissionLeaseHeld) { + return; + } + admissionLeaseHeld = false; + void GenerationJobManager.clearUserFinalization( + userId, + streamId, + req.user?.tenantId, + undefined, + admissionLeaseId, + ).catch(() => undefined); + }; + const refuseAdmission = async (status, body) => { + releaseAdmissionLease(); + if (!exemptFromConcurrency) { + await decrementPendingRequest(userId).catch(() => undefined); + } + if (ownedIdempotencyClaim) { + await GenerationJobManager.releaseGeneration( + userId, + clientRequestId, + streamId, + ownedIdempotencyClaim, + ).catch(() => {}); + } + startupTelemetry?.end('rejected'); + return sendGenerationJson(res, status, body, generationProtocolVersion); + }; + admissionLeaseHeld = await GenerationJobManager.registerUserFinalization( + userId, + streamId, + req.user?.tenantId, + undefined, + admissionLeaseId, + ) + .then(() => true) + .catch((err) => { + logger.warn( + `[ResumableAgentController] Admission lease registration failed: ${streamId}`, + err, + ); + return false; + }); + if (!admissionLeaseHeld) { + res.set('Retry-After', '2'); + return refuseAdmission(503, { + error: 'Could not fence the request against account deletion. Please retry.', + }); + } + if (await isUserDeleting(userId).catch(() => true)) { + return refuseAdmission(403, { error: 'Account deletion is in progress' }); + } + let client = null; let jobCreatedAt; let providerExecutionId; @@ -1370,6 +1435,9 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit * an active job nor a fence and cascades mid-title — and cleared once the title * work settles (or immediately when no post-terminal title runs). */ let userFinalizationRegistered = false; + /** Unique per contender: completion and Stop race the same generation, and a + * losing contender's cleanup must only ever release its own lease. */ + const finalizationLeaseId = crypto.randomUUID(); const registerUserFinalizationFence = async () => { if (userFinalizationRegistered) { return true; @@ -1383,6 +1451,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit streamId, req.user?.tenantId, jobCreatedAt, + finalizationLeaseId, ) .then(() => true) .catch((err) => { @@ -1405,6 +1474,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit streamId, req.user?.tenantId, jobCreatedAt, + finalizationLeaseId, ).catch(() => undefined); }; @@ -2302,6 +2372,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } } + // A stop that won the terminal CAS leaves THIS controller's pending saves + // (background user-message/convo writes) racing deletion with no cover of + // their own — the Stop side's lease fences its `beforePublish` writes, not + // ours. Best-effort: the writes are already in flight either way. + if (wasAborted) { + await registerUserFinalizationFence(); + } let errorScheduleOutcomeRecorded = true; let errorScheduleOutcome = null; if (scheduleId) { @@ -2496,6 +2573,8 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit }); } catch (error) { logger.error('[ResumableAgentController] Initialization error:', error); + // A failed initialization creates no durable job for the lease to hand off to. + releaseAdmissionLease(); const initializationFailure = getInitializationFailure(error); try { if (!res.headersSent) { diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 80b59e7cb7..61dfb491e6 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -449,12 +449,16 @@ async function finalizeResumedTurn({ // fenced settlement either registers its marker or leaves the job paused/active — // deletion-visible either way. Cleared in the persistence try's finally. let finalizationRegistered = false; + /** Unique per contender: this resume and a racing Stop contend on the same + * generation, and a losing contender's cleanup must only release its own lease. */ + const finalizationLeaseId = crypto.randomUUID(); for (let attempt = 1; attempt <= 2 && !finalizationRegistered; attempt++) { finalizationRegistered = await GenerationJobManager.registerUserFinalization( userId, streamId, req.user?.tenantId, job.createdAt, + finalizationLeaseId, ) .then(() => true) .catch((err) => { @@ -474,6 +478,7 @@ async function finalizeResumedTurn({ streamId, req.user?.tenantId, job.createdAt, + finalizationLeaseId, ).catch(() => undefined); }; @@ -1532,11 +1537,13 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) // the transition and leaves the job visible to deletion (the stale reaper // recovers it). Cleared immediately after — the error path owes no // message-recreating writes past its CAS. + const errorLeaseId = crypto.randomUUID(); const errorFenced = await GenerationJobManager.registerUserFinalization( userId, streamId, req.user?.tenantId, job.createdAt, + errorLeaseId, ) .then(() => true) .catch(() => false); @@ -1567,6 +1574,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) streamId, req.user?.tenantId, job.createdAt, + errorLeaseId, ).catch(() => undefined); } } diff --git a/api/server/middleware/__tests__/abortMiddleware.spec.js b/api/server/middleware/__tests__/abortMiddleware.spec.js new file mode 100644 index 0000000000..4555dc07f4 --- /dev/null +++ b/api/server/middleware/__tests__/abortMiddleware.spec.js @@ -0,0 +1,107 @@ +const mockAbortJob = jest.fn(); +const mockSaveMessage = jest.fn(); +const mockSpendTokens = jest.fn(); +const mockGetConvo = jest.fn(); +const mockRecordCollectedUsage = jest.fn(); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { error: jest.fn(), warn: jest.fn(), info: jest.fn(), debug: jest.fn() }, +})); + +jest.mock('@librechat/api', () => ({ + isEnabled: jest.fn(() => false), + sendEvent: jest.fn(), + countTokens: jest.fn(async () => 7), + GenerationJobManager: { + abortJob: (...args) => mockAbortJob(...args), + }, + recordCollectedUsage: (...args) => mockRecordCollectedUsage(...args), + sanitizeMessageForTransmit: (message) => message, + buildAbortedResponseMetadata: jest.fn(() => null), +})); + +jest.mock('~/app/clients/prompts', () => ({ + truncateText: (text) => text, + smartTruncateText: (text) => text, +})); +jest.mock('~/cache/clearPendingReq', () => jest.fn()); +jest.mock('~/server/middleware/error', () => ({ sendError: jest.fn() })); +jest.mock('../abortRun', () => ({ abortRun: jest.fn() })); +jest.mock('~/models', () => ({ + saveMessage: (...args) => mockSaveMessage(...args), + spendTokens: (...args) => mockSpendTokens(...args), + getConvo: (...args) => mockGetConvo(...args), +})); + +const { handleAbort } = require('../abortMiddleware'); + +function createRes() { + const res = { headersSent: false }; + res.status = jest.fn().mockReturnValue(res); + res.send = jest.fn().mockReturnValue(res); + res.setHeader = jest.fn().mockReturnValue(res); + return res; +} + +describe('legacy abortMessage persistence enrollment', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetConvo.mockResolvedValue({ title: 'A chat' }); + mockSaveMessage.mockResolvedValue({ messageId: 'resp-1' }); + mockSpendTokens.mockResolvedValue(undefined); + }); + + it('runs its usage spend and response save INSIDE beforePublish, before abortJob resolves', async () => { + // The old shape awaited abortJob and persisted afterwards — after abortJob had + // already released its finalization lease, so an account-deletion cascade could + // race the spend/save. Enrolled in beforePublish, the writes stay covered. + const persistedDuringAbort = { spend: false, save: false }; + mockAbortJob.mockImplementation(async (streamId, options) => { + await options.beforePublish({ + jobData: { + responseMessageId: 'resp-1', + userMessage: { messageId: 'user-1' }, + conversationId: 'conv-1', + sender: 'AI', + promptTokens: 3, + }, + content: [{ type: 'text', text: 'partial' }], + text: 'partial', + collectedUsage: [], + }); + persistedDuringAbort.spend = mockSpendTokens.mock.calls.length > 0; + persistedDuringAbort.save = mockSaveMessage.mock.calls.length > 0; + return { + success: true, + jobData: { + responseMessageId: 'resp-1', + userMessage: { messageId: 'user-1' }, + conversationId: 'conv-1', + }, + content: [], + text: 'partial', + collectedUsage: [], + finalEvent: null, + }; + }); + const req = { user: { id: 'user-1', email: 'a@b.com' }, body: { abortKey: 'conv-1:x' } }; + const res = createRes(); + + await handleAbort()(req, res); + + expect(persistedDuringAbort).toEqual({ spend: true, save: true }); + expect(res.send).toHaveBeenCalled(); + }); + + it('answers 204 without persisting when the abort found nothing', async () => { + mockAbortJob.mockResolvedValue({ success: false, jobData: null, content: [], text: '' }); + const req = { user: { id: 'user-1', email: 'a@b.com' }, body: { abortKey: 'conv-1:x' } }; + const res = createRes(); + + await handleAbort()(req, res); + + expect(res.status).toHaveBeenCalledWith(204); + expect(mockSaveMessage).not.toHaveBeenCalled(); + expect(mockSpendTokens).not.toHaveBeenCalled(); + }); +}); diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js index feed002b0e..5366c5cfd2 100644 --- a/api/server/middleware/abortMiddleware.js +++ b/api/server/middleware/abortMiddleware.js @@ -109,8 +109,71 @@ async function abortMessage(req, res) { const conversationId = abortKey?.split(':')?.[0] ?? req.user.id; const userId = req.user.id; - // Use GenerationJobManager to abort the job (streamId === conversationId) - const abortResult = await GenerationJobManager.abortJob(conversationId); + let responseMessage = null; + // Use GenerationJobManager to abort the job (streamId === conversationId). + // The usage spend and response save run INSIDE `beforePublish` — between the + // abort CAS and publication — so they stay covered by abortJob's finalization + // lease. Running them after abortJob returned (as this middleware used to) left + // them racing an account-deletion cascade with the lease already released. + const abortResult = await GenerationJobManager.abortJob(conversationId, { + beforePublish: async (pendingAbortResult) => { + const { jobData, content, text, collectedUsage } = pendingAbortResult; + const completionTokens = await countTokens(text); + const promptTokens = jobData?.promptTokens ?? 0; + + responseMessage = { + messageId: jobData?.responseMessageId, + parentMessageId: jobData?.userMessage?.messageId, + conversationId: jobData?.conversationId, + content, + text, + sender: jobData?.sender ?? 'AI', + finish_reason: 'incomplete', + endpoint: jobData?.endpoint, + iconURL: jobData?.iconURL, + model: jobData?.model, + unfinished: false, + error: false, + isCreatedByUser: false, + tokenCount: completionTokens, + }; + + /** Persist the usage/cost rollup + context breakdown for the stopped response + * so its branch/total cost and granular rows survive a reload, matching the + * normal completion path. */ + const abortMetadata = buildAbortedResponseMetadata(jobData); + if (abortMetadata) { + responseMessage.metadata = abortMetadata; + } + + // Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo) + if (collectedUsage && collectedUsage.length > 0) { + await spendCollectedUsage({ + userId, + conversationId: jobData?.conversationId, + collectedUsage, + fallbackModel: jobData?.model, + messageId: jobData?.responseMessageId, + }); + } else { + // Fallback: no collected usage, use text-based token counting for primary model only + await db.spendTokens( + { ...responseMessage, context: 'incomplete', user: userId }, + { promptTokens, completionTokens }, + ); + } + + await db.saveMessage( + { + userId: req?.user?.id, + isTemporary: req?.body?.isTemporary, + interfaceConfig: req?.config?.interfaceConfig, + }, + { ...responseMessage, user: userId }, + { context: 'api/server/middleware/abortMiddleware.js' }, + ); + }, + }); if (!abortResult.success) { if (!res.headersSent) { @@ -118,63 +181,13 @@ async function abortMessage(req, res) { } return; } - - const { jobData, content, text, collectedUsage } = abortResult; - - const completionTokens = await countTokens(text); - const promptTokens = jobData?.promptTokens ?? 0; - - const responseMessage = { - messageId: jobData?.responseMessageId, - parentMessageId: jobData?.userMessage?.messageId, - conversationId: jobData?.conversationId, - content, - text, - sender: jobData?.sender ?? 'AI', - finish_reason: 'incomplete', - endpoint: jobData?.endpoint, - iconURL: jobData?.iconURL, - model: jobData?.model, - unfinished: false, - error: false, - isCreatedByUser: false, - tokenCount: completionTokens, - }; - - /** Persist the usage/cost rollup + context breakdown for the stopped response - * so its branch/total cost and granular rows survive a reload, matching the - * normal completion path. */ - const abortMetadata = buildAbortedResponseMetadata(jobData); - if (abortMetadata) { - responseMessage.metadata = abortMetadata; + if (abortResult.persistenceFailed) { + // Same observable behavior as the old inline throw: handleAbort logs it and no + // normal final envelope is sent for a stop whose persistence failed. + throw new Error(`Failed to persist aborted response for ${conversationId}`); } - // Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo) - if (collectedUsage && collectedUsage.length > 0) { - await spendCollectedUsage({ - userId, - conversationId: jobData?.conversationId, - collectedUsage, - fallbackModel: jobData?.model, - messageId: jobData?.responseMessageId, - }); - } else { - // Fallback: no collected usage, use text-based token counting for primary model only - await db.spendTokens( - { ...responseMessage, context: 'incomplete', user: userId }, - { promptTokens, completionTokens }, - ); - } - - await db.saveMessage( - { - userId: req?.user?.id, - isTemporary: req?.body?.isTemporary, - interfaceConfig: req?.config?.interfaceConfig, - }, - { ...responseMessage, user: userId }, - { context: 'api/server/middleware/abortMiddleware.js' }, - ); + const { jobData } = abortResult; // Get conversation for title const conversation = await db.getConvo(userId, conversationId); diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index c05d57cab9..7e7a3a98f7 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -3771,6 +3771,21 @@ class GenerationJobManagerClass { logger.error(`[GenerationJobManager] Failed to republish abort for ${streamId}:`, err); } } + if (!delivered && !published) { + // HEARTBEAT the fence: the winning abort's lease is TTL-bounded and never + // renewed on its own, so each retry re-registers a fresh lease while the stop + // remains undelivered — the deletion quiesce keeps deferring for as long as + // someone is still trying to deliver it. + await this.jobStore + .registerUserFinalization?.( + jobData.userId, + streamId, + jobData.tenantId, + jobData.createdAt, + randomUUID(), + ) + .catch(() => undefined); + } return { delivered, published }; } @@ -3935,12 +3950,14 @@ class GenerationJobManagerClass { * has been aborted yet, so refusing merely defers, and the Stop route/quiesce * both treat the failure as retryable. Generation-qualified so a late clear * from a predecessor can never drop a replacement's marker. */ + const abortLeaseId = randomUUID(); try { await this.jobStore.registerUserFinalization?.( jobData.userId, streamId, jobData.tenantId, jobData.createdAt, + abortLeaseId, ); } catch (fenceError) { logger.error( @@ -3959,7 +3976,13 @@ class GenerationJobManagerClass { } const clearAbortFinalization = () => { void this.jobStore - .clearUserFinalization?.(jobData.userId, streamId, jobData.tenantId, jobData.createdAt) + .clearUserFinalization?.( + jobData.userId, + streamId, + jobData.tenantId, + jobData.createdAt, + abortLeaseId, + ) .catch(() => undefined); }; diff --git a/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts b/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts index 3b49d11659..d9f9837373 100644 --- a/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts +++ b/packages/api/src/stream/__tests__/abortSignalHonesty.spec.ts @@ -134,7 +134,13 @@ describe('abort signal honesty across replicas', () => { const result = await owner.abortJob('conv-6'); expect(result.success).toBe(true); - expect(register).toHaveBeenCalledWith('user-1', 'conv-6', undefined, createdAt); + expect(register).toHaveBeenCalledWith( + 'user-1', + 'conv-6', + undefined, + createdAt, + expect.any(String), + ); // BEFORE the CAS: registering after it leaves the window uncovered. expect(register.mock.invocationCallOrder[0]).toBeLessThan( transition.mock.invocationCallOrder[0], @@ -160,6 +166,93 @@ describe('abort signal honesty across replicas', () => { await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(0); }); + it('RETAINS the abort lease while the stop is undelivered and unpublished', async () => { + const { owner, peer, peerTransport, jobStore } = await makeManagers(); + await owner.createJob('conv-8', 'user-1', 'conv-8'); + // Peer-owned generation; the acknowledged publication provably fails. + Object.assign(peerTransport, { + emitAbort: jest.fn((): void => undefined), + emitAbortConfirmed: jest.fn(async () => false), + }); + + const result = await peer.abortJob('conv-8'); + + expect(result.success).toBe(true); + expect(result.signalDelivered).toBe(false); + expect(result.signalPublished).toBe(false); + // The job is terminal (invisible to active-set scans) and a user Stop writes no + // durable abort fence — this lease is the ONLY thing deferring a deletion + // quiesce while the remote owner keeps generating. + await new Promise((resolve) => setImmediate(resolve)); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBe(1); + }); + + it('renews the fence on each undelivered resignal attempt', async () => { + const { owner, peer, peerTransport, jobStore } = await makeManagers(); + await owner.createJob('conv-9', 'user-1', 'conv-9'); + Object.assign(peerTransport, { + emitAbort: jest.fn((): void => undefined), + emitAbortConfirmed: jest.fn(async () => false), + }); + await peer.abortJob('conv-9'); + const register = jest.spyOn(jobStore, 'registerUserFinalization'); + + const resignal = await peer.resignalAbort('conv-9'); + + expect(resignal).toEqual({ delivered: false, published: false }); + // The 5-minute lease is never self-renewing: each retry heartbeats a fresh one + // so the quiesce keeps deferring for as long as delivery is still being driven. + expect(register).toHaveBeenCalledWith( + 'user-1', + 'conv-9', + undefined, + expect.any(Number), + expect.any(String), + ); + }); + + it('a losing second Stop releases only its own lease, never the winning one', async () => { + const { owner, peer, jobStore } = await makeManagers(); + await owner.createJob('conv-10', 'user-1', 'conv-10'); + // The winning abort parks inside its beforePublish persistence. + let releaseWinner!: () => void; + const winnerPersisting = new Promise((resolve) => { + releaseWinner = resolve; + }); + let winnerEntered!: () => void; + const winnerEnteredPersistence = new Promise((resolve) => { + winnerEntered = resolve; + }); + const winner = owner.abortJob('conv-10', { + beforePublish: async () => { + winnerEntered(); + await winnerPersisting; + }, + }); + await winnerEnteredPersistence; + + // A second Stop RACES the winner: its pre-abort read still sees `running` (the + // winner's CAS lands in between), so it gets past the terminal pre-check, + // registers its own lease, and only then loses the CAS. Pre-lease keying, its + // registration overwrote the winner's field and its cleanup then deleted it. + const realGetJob = jobStore.getJob.bind(jobStore); + const staleRunning = { ...(await realGetJob('conv-10'))!, status: 'running' as const }; + const getJobSpy = jest + .spyOn(jobStore, 'getJob') + .mockImplementationOnce(async () => staleRunning); + const loser = await peer.abortJob('conv-10'); + getJobSpy.mockRestore(); + expect(loser.success).toBe(false); + + // The winner's persistence is still live: its lease must have survived the + // loser's cleanup. Keyed without lease tokens, this read 0. + await new Promise((resolve) => setImmediate(resolve)); + await expect(jobStore.countUserFinalizations('user-1')).resolves.toBeGreaterThanOrEqual(1); + + releaseWinner(); + await winner; + }); + it('reports republication as succeeded when the owner acknowledges', async () => { const { owner, peer, peerTransport, jobStore } = await makeManagers(); await owner.createJob('conv-4', 'user-1', 'conv-4'); diff --git a/packages/api/src/stream/__tests__/userFinalizations.spec.ts b/packages/api/src/stream/__tests__/userFinalizations.spec.ts index 29bf6d07db..5e9a005733 100644 --- a/packages/api/src/stream/__tests__/userFinalizations.spec.ts +++ b/packages/api/src/stream/__tests__/userFinalizations.spec.ts @@ -66,6 +66,18 @@ describe('InMemoryJobStore user finalizations', () => { await expect(store.countUserFinalizations('user-1')).resolves.toBe(0); }); + it('isolates same-generation contenders by lease: a losing Stop cannot drop completion', async () => { + // Completion and Stop (or two Stops) contend on the SAME (stream, generation). + // Keyed without a lease token they shared one entry, and the losing contender's + // cleanup erased the winner's still-live marker. + await store.registerUserFinalization('user-1', 'conv-1', undefined, 1000, 'completion-lease'); + await store.registerUserFinalization('user-1', 'conv-1', undefined, 1000, 'stop-lease'); + + await store.clearUserFinalization('user-1', 'conv-1', undefined, 1000, 'stop-lease'); + + await expect(store.countUserFinalizations('user-1')).resolves.toBe(1); + }); + it('an unqualified legacy clear does not drop generation-qualified markers', async () => { await store.registerUserFinalization('user-1', 'conv-1', undefined, 1000); diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index 26e771ee5c..b2634cf7e0 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -1348,6 +1348,7 @@ export class InMemoryJobStore implements IJobStoreV2 { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise { const userKey = tenantId ? `${tenantId}:${userId}` : userId; let entries = this.userFinalizations.get(userKey); @@ -1355,7 +1356,10 @@ export class InMemoryJobStore implements IJobStoreV2 { entries = new Map(); this.userFinalizations.set(userKey, entries); } - entries.set(finalizationField(streamId, generationId), Date.now() + USER_FINALIZATION_TTL_MS); + entries.set( + finalizationField(streamId, generationId, leaseId), + Date.now() + USER_FINALIZATION_TTL_MS, + ); } async clearUserFinalization( @@ -1363,13 +1367,14 @@ export class InMemoryJobStore implements IJobStoreV2 { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise { const userKey = tenantId ? `${tenantId}:${userId}` : userId; const entries = this.userFinalizations.get(userKey); if (!entries) { return; } - entries.delete(finalizationField(streamId, generationId)); + entries.delete(finalizationField(streamId, generationId, leaseId)); if (entries.size === 0) { this.userFinalizations.delete(userKey); } diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index 7d920fee8d..6b85f333d4 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -3056,6 +3056,7 @@ export class RedisJobStore implements IJobStoreV2 { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise { const key = KEYS.userFinalizations(userId, tenantId); const expiresAt = Date.now() + USER_FINALIZATION_TTL_MS; @@ -3068,7 +3069,7 @@ export class RedisJobStore implements IJobStoreV2 { REGISTER_FINALIZATION_LUA, 1, key, - finalizationField(streamId, generationId), + finalizationField(streamId, generationId, leaseId), String(expiresAt), String(Math.ceil(USER_FINALIZATION_TTL_MS / 1000)), ); @@ -3079,10 +3080,11 @@ export class RedisJobStore implements IJobStoreV2 { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise { await this.redis.hdel( KEYS.userFinalizations(userId, tenantId), - finalizationField(streamId, generationId), + finalizationField(streamId, generationId, leaseId), ); } diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index e122b2f9b1..0657a72cb6 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -323,11 +323,25 @@ export type JobMetadataPatch = Partial< * stores so an expired approval is kept out of active-job listings (the client * stops polling; cleanup/expiry finalizes it). */ -/** Generation-qualified finalization-marker field, shared by the stores so a late - * clear from generation A can never drop the marker generation B registered on the - * same stream. Unqualified (legacy) callers fall back to the bare streamId. */ -export function finalizationField(streamId: string, generationId?: number): string { - return generationId != null ? `${streamId}:${generationId}` : streamId; +/** Lease-qualified finalization-marker field, shared by the stores. Generation + * qualification stops a late clear from generation A dropping generation B's marker + * on the same stream; the LEASE token goes further — completion and Stop (or two + * Stops) contend on the SAME generation, and a losing contender's cleanup must only + * ever release its own lease, never the winner's still-live one. Unqualified + * (legacy) callers fall back to the bare streamId. */ +export function finalizationField( + streamId: string, + generationId?: number, + leaseId?: string, +): string { + let field = streamId; + if (generationId != null) { + field += `:${generationId}`; + } + if (leaseId != null) { + field += `:${leaseId}`; + } + return field; } export function isPendingActionExpired(job: Pick): boolean { @@ -801,6 +815,7 @@ export interface IJobStore { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise; /** Clear a finalization registered by {@link registerUserFinalization}. Must be @@ -812,6 +827,7 @@ export interface IJobStore { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise; /** Count live (unexpired) finalizations for a user. */ @@ -871,12 +887,14 @@ export interface IJobStoreV2 extends IJobStore { streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise; clearUserFinalization( userId: string, streamId: string, tenantId?: string, generationId?: number, + leaseId?: string, ): Promise; countUserFinalizations(userId: string, tenantId?: string): Promise;