diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index 8024b8d408..addbdadc60 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -13,10 +13,19 @@ const mockGenerationJobManager = { completeJob: jest.fn(), getResumeState: jest.fn(), updateMetadata: jest.fn(), + claimGeneration: jest.fn(), + releaseGeneration: jest.fn(), + hasJob: jest.fn(), }; const mockCheckAndIncrementPendingRequest = jest.fn(); const mockDecrementPendingRequest = jest.fn(); +const mockGetViolationInfo = jest.fn(() => ({ + type: 'concurrent', + limit: 2, + pendingRequests: 3, + score: 1, +})); const mockFilterPersistableAbortContent = jest.fn((content) => content.filter((part) => part?.type !== 'tool_call'), ); @@ -82,7 +91,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ sendEvent: jest.fn(), - getViolationInfo: jest.fn(), + getViolationInfo: (...args) => mockGetViolationInfo(...args), buildMessageFiles: jest.fn(() => []), resolveTitleTiming: jest.fn(() => 'immediate'), GenerationJobManager: mockGenerationJobManager, @@ -186,6 +195,10 @@ describe('ResumableAgentController resume metadata', () => { mockGenerationJobManager.getResumeState.mockResolvedValue(null); mockGenerationJobManager.updateMetadata.mockResolvedValue(undefined); mockGenerationJobManager.emitError.mockResolvedValue(undefined); + mockGenerationJobManager.completeJob.mockResolvedValue(undefined); + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + mockGenerationJobManager.releaseGeneration.mockResolvedValue(undefined); + mockGenerationJobManager.hasJob.mockResolvedValue(true); mockSaveMessage.mockResolvedValue({}); }); @@ -618,4 +631,269 @@ describe('ResumableAgentController resume metadata', () => { expect.any(Object), ); }); + + it('dedups a retried start-generation request to the original stream', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ + claimed: false, + existing: { streamId: 'orig-stream', conversationId: 'orig-convo' }, + }); + mockGenerationJobManager.hasJob.mockResolvedValue(true); + const initializeClient = jest.fn(); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Retried after a lost response.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(res.json).toHaveBeenCalledWith({ + streamId: 'orig-stream', + conversationId: 'orig-convo', + status: 'resumed', + }); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); + expect(initializeClient).not.toHaveBeenCalled(); + }); + + it('resumes when the job is missing but the claim is old (original completed and was cleaned up)', async () => { + // An old claim with no job means the original already ran and was cleaned up; the deduped + // response must attach (client 404 handler refetches) rather than loop on readiness. + mockGenerationJobManager.claimGeneration.mockResolvedValue({ + claimed: false, + existing: { + streamId: 'orig-stream', + conversationId: 'orig-convo', + claimedAt: Date.now() - 60000, + }, + }); + mockGenerationJobManager.hasJob.mockResolvedValue(false); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Retry after a fast, already-cleaned-up generation.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.json).toHaveBeenCalledWith({ + streamId: 'orig-stream', + conversationId: 'orig-convo', + status: 'resumed', + }); + expect(res.status).not.toHaveBeenCalledWith(503); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + }); + + it('returns 503 SERVER_NOT_READY when a fresh claim still has no job (winner is between claim and createJob)', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ + claimed: false, + existing: { + streamId: 'orig-stream', + conversationId: 'orig-convo', + claimedAt: Date.now(), + }, + }); + mockGenerationJobManager.hasJob.mockResolvedValue(false); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Concurrent duplicate before the winner wrote its job.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.set).toHaveBeenCalledWith('Retry-After', '1'); + expect(res.status).toHaveBeenCalledWith(503); + expect(res.json).toHaveBeenCalledWith(expect.objectContaining({ code: 'SERVER_NOT_READY' })); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + }); + + it('never starts a second generation when the job lookup fails for a confirmed duplicate', async () => { + // A store hiccup while checking an existing claim must not fail open into createJob. + mockGenerationJobManager.claimGeneration.mockResolvedValue({ + claimed: false, + existing: { streamId: 'orig-stream', conversationId: 'orig-convo', claimedAt: Date.now() }, + }); + mockGenerationJobManager.hasJob.mockRejectedValue(new Error('redis down')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Duplicate during a Redis hiccup.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).toHaveBeenCalledWith(503); + expect(res.json).toHaveBeenCalledWith(expect.objectContaining({ code: 'SERVER_NOT_READY' })); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + }); + + it('finalizes the failed job before releasing the idempotency claim', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + const initializeClient = jest.fn().mockRejectedValue(new Error('init boom after res.json')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Start fails after the initial JSON.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + 'conversation-123', + expect.any(String), + ); + expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + // completeJob must finalize the failed job BEFORE the claim is released, or a racing + // retry could win the key, createJob the same streamId, and be aborted by this completeJob. + expect(mockGenerationJobManager.completeJob.mock.invocationCallOrder[0]).toBeLessThan( + mockGenerationJobManager.releaseGeneration.mock.invocationCallOrder[0], + ); + }); + + it('still releases the claim and pending slot when completeJob fails during init-error cleanup', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + mockGenerationJobManager.completeJob.mockRejectedValue(new Error('store hiccup')); + const initializeClient = jest.fn().mockRejectedValue(new Error('init boom after res.json')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Start fails while the store is degraded.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = createResumableResponse(); + + await AgentController(req, res, jest.fn(), initializeClient, null); + + // A completeJob rejection must not wedge the retry behind the claim or leak the slot. + expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + expect(mockDecrementPendingRequest).toHaveBeenCalledWith('user-123'); + }); + + it('proceeds to create the job when it wins the idempotency claim', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Fresh submission.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { + headersSent: true, + json: jest.fn(() => { + res.headersSent = true; + }), + status: jest.fn(() => res), + set: jest.fn(), + }; + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(mockCheckAndIncrementPendingRequest).toHaveBeenCalledWith('user-123'); + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + 'conversation-123', + 'user-123', + 'conversation-123', + ); + }); + + it('releases the idempotency claim on a 429 only when it won the claim', async () => { + mockGenerationJobManager.claimGeneration.mockResolvedValue({ claimed: true }); + mockCheckAndIncrementPendingRequest.mockResolvedValue({ + allowed: false, + pendingRequests: 3, + limit: 2, + }); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Over the limit.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).toHaveBeenCalledWith(429); + expect(mockGenerationJobManager.releaseGeneration).toHaveBeenCalledWith('user-123', 'req-abc'); + }); + + it('does not release a claim it never won when a fail-open duplicate hits the limiter', async () => { + mockGenerationJobManager.claimGeneration.mockRejectedValue(new Error('redis down')); + mockCheckAndIncrementPendingRequest.mockResolvedValue({ + allowed: false, + pendingRequests: 3, + limit: 2, + }); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Duplicate while the original runs.', + messageId: 'user-msg', + clientRequestId: 'req-abc', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res), set: jest.fn() }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).toHaveBeenCalledWith(429); + expect(mockGenerationJobManager.releaseGeneration).not.toHaveBeenCalled(); + }); }); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 5b63c06ef8..49d9329f0c 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -176,6 +176,30 @@ async function finishResumableRequest(req, userId) { } } +const JOB_RECORD_WAIT_ATTEMPTS = 5; +const JOB_RECORD_WAIT_DELAY_MS = 60; + +// A winner writes its job record within a few ms of claiming; if a losing duplicate still +// sees no job within this window of the claim, the winner is still starting (retry rather +// than hand back a stream that would 404). Past it, a missing job means the original +// already completed and was cleaned up (attach and let the client refetch). +const IDEMPOTENCY_STARTUP_GRACE_MS = 5000; + +/** + * Poll briefly for a job record to appear. A deduped retry that loses the idempotency + * claim must not be handed the winner's stream until its job exists, or the client's + * subscribe 404s terminally. The winner writes the record a few ms after claiming. + */ +async function waitForJobRecord(streamId) { + for (let attempt = 0; attempt < JOB_RECORD_WAIT_ATTEMPTS; attempt++) { + if (await GenerationJobManager.hasJob(streamId)) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, JOB_RECORD_WAIT_DELAY_MS)); + } + return GenerationJobManager.hasJob(streamId); +} + function rejectPreliminaryParentMessageId(res) { return res.status(409).json({ error: @@ -219,13 +243,6 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit * Resolved from the agent's actual endpoint once the client is initialized. */ let titleTiming = 'immediate'; - const { allowed, pendingRequests, limit } = await checkAndIncrementPendingRequest(userId); - if (!allowed) { - const violationInfo = getViolationInfo(pendingRequests, limit); - await logViolation(req, res, ViolationTypes.CONCURRENT, violationInfo, violationInfo.score); - return res.status(429).json(violationInfo); - } - // Generate conversationId upfront if not provided - streamId === conversationId always // Treat "new" as a placeholder that needs a real UUID (frontend may send "new" for new convos) const isNewConvo = !reqConversationId || reqConversationId === 'new'; @@ -233,6 +250,96 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit const streamId = conversationId; req.body.conversationId = conversationId; + // Idempotency: a lost/reset start-generation response makes the client re-POST the + // identical payload, which would otherwise start a second fully-billed generation. + // Claim the submission's clientRequestId before creating the job so a retry attaches + // to the original stream instead of spawning a duplicate. Runs before the concurrency + // check so a deduped retry is never counted against the limiter. Fail-open on errors. + const clientRequestId = req.body?.clientRequestId; + let ownsIdempotencyClaim = false; + if (clientRequestId) { + let claim = null; + try { + claim = await GenerationJobManager.claimGeneration( + userId, + clientRequestId, + streamId, + conversationId, + ); + } catch (err) { + // The claim itself could not be determined (store unavailable): fail open and proceed + // as a fresh request rather than blocking the send. This is the ONLY fail-open path — + // once a duplicate is confirmed below, an error must never fall through to a second + // billed generation. + logger.error( + '[ResumableAgentController] Idempotency claim failed; proceeding without dedup', + err, + ); + } + + if (claim?.claimed) { + ownsIdempotencyClaim = true; + } else if (claim?.existing) { + // A duplicate is confirmed. Attach to the original stream — and never fall through to + // a second generation, even if the job lookup hiccups. + const existingStreamId = claim.existing.streamId; + let jobExists = false; + try { + // Wait briefly for the winner to write the job record (it does so a few ms after + // claiming) so a still-live stream isn't handed back before its job exists. + jobExists = await waitForJobRecord(existingStreamId); + } catch (err) { + // Store hiccup while checking the job: ask the client to retry rather than starting + // a second generation for a request we know is a duplicate. + logger.error( + '[ResumableAgentController] Job lookup failed for an existing claim; asking the client to retry', + err, + ); + res.set('Retry-After', '1'); + return res.status(503).json({ + code: 'SERVER_NOT_READY', + error: 'Generation is still starting. Please retry shortly.', + }); + } + const claimAgeMs = Date.now() - (claim.existing.claimedAt ?? 0); + if (!jobExists && claimAgeMs < IDEMPOTENCY_STARTUP_GRACE_MS) { + // The winner claimed but has not written the job yet (still between claim and + // createJob). Handing back the stream now would 404 and tear down the client while + // the winner goes on to generate and bill with no UI attached — ask the client to + // retry via the readiness path instead. + res.set('Retry-After', '1'); + return res.status(503).json({ + code: 'SERVER_NOT_READY', + error: 'Generation is still starting. Please retry shortly.', + }); + } + // Job exists (live), or the grace elapsed with none (the original already completed + // and was cleaned up, or the winner died): attach. A then-missing job recovers via + // the client's subscribe 404 handler (refetch persisted messages) rather than an + // indefinite readiness loop. + logger.debug('[ResumableAgentController] Deduped retried start-generation request', { + userId, + clientRequestId, + streamId: existingStreamId, + }); + return res.json({ + streamId: existingStreamId, + conversationId: claim.existing.conversationId, + status: 'resumed', + }); + } + } + + const { allowed, pendingRequests, limit } = await checkAndIncrementPendingRequest(userId); + if (!allowed) { + if (ownsIdempotencyClaim) { + await GenerationJobManager.releaseGeneration(userId, clientRequestId).catch(() => {}); + } + const violationInfo = getViolationInfo(pendingRequests, limit); + await logViolation(req, res, ViolationTypes.CONCURRENT, violationInfo, violationInfo.score); + return res.status(429).json(violationInfo); + } + let client = null; try { @@ -941,7 +1048,22 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // JSON already sent, emit error to stream so client can receive it await GenerationJobManager.emitError(streamId, error.message || 'Failed to start generation'); } - GenerationJobManager.completeJob(streamId, error.message); + // Finalize THIS failed job before releasing the idempotency claim. Releasing first would + // let the client's retry win the same key and createJob() the same streamId while we are + // still here — and completeJob() is not guarded by the original createdAt, so it would + // abort/error that replacement. A completeJob() rejection (store hiccup) must NOT skip the + // release + pending-request decrement below, or the retry stays wedged behind the claim + // and the concurrency slot leaks — so swallow its error. (A failed completeJob did not + // finalize anything, so releasing afterward can't let it abort a later replacement.) + await GenerationJobManager.completeJob(streamId, error.message).catch((completeErr) => { + logger.warn( + '[ResumableAgentController] completeJob failed during init-error cleanup', + completeErr, + ); + }); + if (ownsIdempotencyClaim) { + await GenerationJobManager.releaseGeneration(userId, clientRequestId).catch(() => {}); + } await finishResumableRequest(req, userId); if (client) { disposeClient(client); diff --git a/client/src/hooks/Chat/useChatFunctions.ts b/client/src/hooks/Chat/useChatFunctions.ts index 16fc0b9572..8419f5d767 100644 --- a/client/src/hooks/Chat/useChatFunctions.ts +++ b/client/src/hooks/Chat/useChatFunctions.ts @@ -412,6 +412,10 @@ export default function useChatFunctions({ // construct the query message // this is not a real messageId, it is used as placeholder before real messageId returned const intermediateId = overrideUserMessageId ?? v4(); + /** Stable idempotency key for this submission: fresh per `ask()` (so regenerate differs) + * but reused across the client's start-generation network retries, letting the server + * dedup a retried request instead of starting a second billed generation. */ + const clientRequestId = v4(); if (parentMessageId == null) { parentMessageId = getAppendParentMessageId({ latestMessage, currentMessages }); } @@ -665,6 +669,7 @@ export default function useChatFunctions({ editedContent, addedConvo, manualSkills: manualSkills.length > 0 ? manualSkills : undefined, + clientRequestId, }; if (isRegenerate) { diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index 129bccd626..39c422172e 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -430,6 +430,50 @@ describe('useResumableSSE', () => { unmount(); }); + it('reconciles conversations via refetch instead of removing them on a resume 404', async () => { + mockFindAll.mockReturnValue([{ queryKey: [QueryKeys.allConversations] }]); + // A deduped start returns status: 'resumed', so the client subscribes with resume=true. + (request.post as jest.Mock).mockResolvedValue({ streamId: 'stream-123', status: 'resumed' }); + const submission = buildSubmission({ + conversation: {}, + userMessage: { + messageId: 'msg-1', + conversationId: null, + text: 'Hello', + isCreatedByUser: true, + sender: 'User', + parentMessageId: Constants.NO_PARENT, + }, + initialResponse: { + messageId: 'msg-1_', + conversationId: null, + text: '', + isCreatedByUser: false, + sender: 'Assistant', + }, + }); + const chatHelpers = buildChatHelpers(); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + + await act(async () => { + await Promise.resolve(); + }); + + const sse = getLastSSE(); + await act(async () => { + sse._emit('error', { responseCode: 404 }); + }); + + // Reconcile against the server (refetch) rather than dropping a possibly-persisted + // conversation. The handler is a mutually-exclusive isResume ? invalidate : remove, so + // asserting the invalidate proves the immediate removal did not run. + expect(mockInvalidateQueries).toHaveBeenCalledWith({ + queryKey: [QueryKeys.allConversations], + }); + unmount(); + }); + it('closes the SSE connection on 404', async () => { const { sse, unmount } = await render404Scenario(); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index aadb6c9474..11f15b7206 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -90,6 +90,12 @@ const getStartGenerationStreamId = (data: unknown): string | null => { return typeof streamId === 'string' && streamId.length > 0 ? streamId : null; }; +/** The server returns `status: 'resumed'` when a duplicate start request was deduped to an + * already-running stream — the client must subscribe with resume=true to replay its state + * (prior content and any pending-action) rather than only receiving live events. */ +const isResumedStartResponse = (data: unknown): boolean => + data != null && typeof data === 'object' && (data as { status?: unknown }).status === 'resumed'; + const parseSSEErrorData = (body: string): unknown | null => { const blocks = body.split(/\r?\n\r?\n/); for (const block of blocks) { @@ -1226,7 +1232,17 @@ export default function useResumableSSE( !createdStreamIdsRef.current.has(currentStreamId) && optimisticStreamIdsRef.current.has(currentStreamId) ) { - removeConvoFromAllQueries(queryClient, currentStreamId); + if (isResume) { + // A resumed subscribe attaches to an already-adopted stream (e.g. a deduped + // start request). A 404 means the job is gone — but the conversation may be + // persisted (the original completed and was cleaned up) or may never have + // existed (the winner died before persisting). Don't guess: reconcile against + // the server so a real conversation stays and a phantom is dropped. + queryClient.invalidateQueries({ queryKey: [QueryKeys.allConversations] }); + } else { + // Fresh optimistic stream that never started: prune immediately. + removeConvoFromAllQueries(queryClient, currentStreamId); + } } setIsSubmitting(false); setShowStopButton(false); @@ -1523,7 +1539,10 @@ export default function useResumableSSE( * Readiness retries honor Retry-After until cleanup or the readiness window expires. */ const startGeneration = useCallback( - async (currentSubmission: TSubmission, signal?: AbortSignal): Promise => { + async ( + currentSubmission: TSubmission, + signal?: AbortSignal, + ): Promise<{ streamId: string; resumed: boolean } | null> => { const payloadData = createPayload(currentSubmission); let { payload } = payloadData; payload = removeNullishValues(payload) as TPayload; @@ -1548,8 +1567,9 @@ export default function useResumableSSE( } const streamId = getStartGenerationStreamId(data); if (streamId) { - logger.log('ResumableSSE', 'Generation started:', { streamId }); - return streamId; + const resumed = isResumedStartResponse(data); + logger.log('ResumableSSE', 'Generation started:', { streamId, resumed }); + return { streamId, resumed }; } lastError = { response: { data } }; @@ -1666,11 +1686,12 @@ export default function useResumableSSE( } else { // New generation: start and then subscribe logger.log('ResumableSSE', 'Starting NEW generation'); - const newStreamId = await startGeneration(submission, signal); + const startResult = await startGeneration(submission, signal); if (signal.aborted) { return; } - if (newStreamId) { + if (startResult) { + const { streamId: newStreamId, resumed } = startResult; setStreamId(newStreamId); // Optimistically add to active jobs addActiveJob(newStreamId); @@ -1687,7 +1708,9 @@ export default function useResumableSSE( } const streamSubmission = addOptimisticConversation(newStreamId, submission); submissionRef.current = streamSubmission; - subscribeToStream(newStreamId, streamSubmission); + // A deduped retry (status: 'resumed') attaches to an already-running stream, so + // subscribe with resume=true to replay its state instead of only live events. + subscribeToStream(newStreamId, streamSubmission, resumed); } else { logger.error('ResumableSSE', 'Failed to get streamId from startGeneration'); } diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 46f819effd..53d0c1d20e 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -21,6 +21,7 @@ import type { UsageMetadata, AbortResult, IJobStore, + IdempotencyClaimResult, } from './interfaces/IJobStore'; import type { SteerOwner, SteerContentView } from './SteeringLifecycle'; import type { GenerationJobStore } from '~/app/metrics'; @@ -48,6 +49,10 @@ const APPROVAL_EXPIRED_ERROR = 'Approval expired before a decision was made'; /** Error surfaced to any client still attached when a stale/hung job is reaped. */ const REAPED_JOB_ERROR = 'Generation timed out'; + +/** Lifetime of a start-generation idempotency claim (matches the running-job TTL: 20 min), + * so a late retry still dedups for the whole generation window. */ +const IDEMPOTENCY_TTL_SECONDS = 1200; const OAUTH_TOOL_CALL_PREFIX = `oauth${Constants.mcp_delimiter}`; function getToolCallName(toolCall: unknown): unknown { @@ -690,6 +695,33 @@ class GenerationJobManagerClass { return this.jobStore.hasJob(streamId); } + /** + * Atomically claim a start-generation request for `(userId, clientRequestId)`. + * The first caller wins (`claimed: true`) and should create the job; a retried + * request for the same submission loses and receives the original stream so it + * can attach to it instead of starting a second billed generation. + */ + async claimGeneration( + userId: string, + clientRequestId: string, + streamId: string, + conversationId: string, + ): Promise { + return this.jobStore.claimIdempotencyKey( + `${userId}:${clientRequestId}`, + { streamId, conversationId, claimedAt: Date.now() }, + IDEMPOTENCY_TTL_SECONDS, + ); + } + + /** + * Release a start-generation claim so the submission can be retried (e.g. the + * start failed before generation began). + */ + async releaseGeneration(userId: string, clientRequestId: string): Promise { + await this.jobStore.releaseIdempotencyKey(`${userId}:${clientRequestId}`); + } + /** * Get job status. */ diff --git a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts index 99a43142f2..13a0aa98c6 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -2274,4 +2274,103 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); }); + + describe('Idempotency claims (#14339 duplicate-billing guard)', () => { + test('grants the first claim and returns the original stream to a duplicate', async () => { + if (!ioredisClient) { + return; + } + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const key = `user-1:req-${Date.now()}`; + const first = await store.claimIdempotencyKey( + key, + { streamId: 's1', conversationId: 'c1' }, + 1200, + ); + expect(first).toEqual({ claimed: true }); + + const second = await store.claimIdempotencyKey( + key, + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(second).toEqual({ + claimed: false, + existing: { streamId: 's1', conversationId: 'c1' }, + }); + + await store.destroy(); + }); + + test('sets a bounded TTL on the claim', async () => { + if (!ioredisClient) { + return; + } + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const key = `user-1:req-ttl-${Date.now()}`; + await store.claimIdempotencyKey(key, { streamId: 's1', conversationId: 'c1' }, 1200); + + const pttl = await ioredisClient.pttl(`stream:idem:{${key}}`); + expect(pttl).toBeGreaterThan(0); + expect(pttl).toBeLessThanOrEqual(1200 * 1000); + + await store.destroy(); + }); + + test('releaseIdempotencyKey frees the claim', async () => { + if (!ioredisClient) { + return; + } + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const key = `user-1:req-rel-${Date.now()}`; + await store.claimIdempotencyKey(key, { streamId: 's1', conversationId: 'c1' }, 1200); + await store.releaseIdempotencyKey(key); + + const reclaimed = await store.claimIdempotencyKey( + key, + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(reclaimed).toEqual({ claimed: true }); + + await store.destroy(); + }); + + test('two concurrent claims for one key elect exactly one winner', async () => { + if (!ioredisClient) { + return; + } + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const key = `user-1:req-race-${Date.now()}`; + const [a, b] = await Promise.all([ + store.claimIdempotencyKey(key, { streamId: 'sa', conversationId: 'ca' }, 1200), + store.claimIdempotencyKey(key, { streamId: 'sb', conversationId: 'cb' }, 1200), + ]); + + const winners = [a, b].filter((r) => r.claimed); + const losers = [a, b].filter((r) => !r.claimed); + expect(winners).toHaveLength(1); + expect(losers).toHaveLength(1); + // The loser attaches to whichever stream the winner registered. + expect(losers[0].existing).toEqual( + winners[0] === a + ? { streamId: 'sa', conversationId: 'ca' } + : { streamId: 'sb', conversationId: 'cb' }, + ); + + await store.destroy(); + }); + }); }); diff --git a/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts new file mode 100644 index 0000000000..e33b74b57d --- /dev/null +++ b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts @@ -0,0 +1,156 @@ +import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; +import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; +import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; + +jest.spyOn(console, 'log').mockImplementation(); + +/** + * Start-generation idempotency: a retried start request for the SAME submission must + * attach to the original stream instead of spawning a second billed generation, while a + * distinct submission (including a regenerate) must NOT be deduped. See issue #14339. + */ +describe('InMemoryJobStore.claimIdempotencyKey', () => { + let store: InMemoryJobStore; + + beforeEach(() => { + store = new InMemoryJobStore({ ttlAfterComplete: 0 }); + }); + + it('grants the first claim and returns the original stream to a duplicate', async () => { + const first = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's1', conversationId: 'c1' }, + 1200, + ); + expect(first).toEqual({ claimed: true }); + + // A retry that computed a different streamId still gets the ORIGINAL stream back. + const second = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(second).toEqual({ claimed: false, existing: { streamId: 's1', conversationId: 'c1' } }); + }); + + it('lets a released key be claimed again', async () => { + await store.claimIdempotencyKey('user:req', { streamId: 's1', conversationId: 'c1' }, 1200); + await store.releaseIdempotencyKey('user:req'); + + const reclaimed = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(reclaimed).toEqual({ claimed: true }); + }); + + it('clears claims on destroy so a reused store does not falsely dedup', async () => { + await store.claimIdempotencyKey('user:req', { streamId: 's1', conversationId: 'c1' }, 1200); + await store.destroy(); + const reclaimed = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(reclaimed).toEqual({ claimed: true }); + }); + + it('treats distinct keys independently', async () => { + const a = await store.claimIdempotencyKey( + 'user:reqA', + { streamId: 's1', conversationId: 'c1' }, + 1200, + ); + const b = await store.claimIdempotencyKey( + 'user:reqB', + { streamId: 's2', conversationId: 'c2' }, + 1200, + ); + expect(a).toEqual({ claimed: true }); + expect(b).toEqual({ claimed: true }); + }); + + it('lets the key be reclaimed after its TTL elapses', async () => { + jest.useFakeTimers(); + try { + jest.setSystemTime(new Date('2026-07-20T00:00:00Z')); + const first = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's1', conversationId: 'c1' }, + 1, + ); + expect(first).toEqual({ claimed: true }); + + // Still held one moment before expiry. + jest.setSystemTime(new Date('2026-07-20T00:00:00.999Z')); + const held = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's2', conversationId: 'c2' }, + 1, + ); + expect(held.claimed).toBe(false); + + // Expired: the next caller wins. + jest.setSystemTime(new Date('2026-07-20T00:00:02Z')); + const expired = await store.claimIdempotencyKey( + 'user:req', + { streamId: 's3', conversationId: 'c3' }, + 1, + ); + expect(expired).toEqual({ claimed: true }); + } finally { + jest.useRealTimers(); + } + }); +}); + +describe('GenerationJobManager start-generation claim', () => { + let manager: GenerationJobManagerClass; + + beforeEach(() => { + manager = new GenerationJobManagerClass(); + manager.configure({ + jobStore: new InMemoryJobStore({ ttlAfterComplete: 0 }), + eventTransport: new InMemoryEventTransport(), + isRedis: false, + }); + manager.initialize(); + }); + + afterEach(async () => { + await manager.destroy(); + }); + + it('dedups a retry of the same submission to the original stream', async () => { + const first = await manager.claimGeneration('user-1', 'req-1', 'stream-a', 'convo-a'); + expect(first).toEqual({ claimed: true }); + + const retry = await manager.claimGeneration('user-1', 'req-1', 'stream-b', 'convo-b'); + expect(retry.claimed).toBe(false); + expect(retry.existing).toEqual( + expect.objectContaining({ streamId: 'stream-a', conversationId: 'convo-a' }), + ); + expect(typeof retry.existing?.claimedAt).toBe('number'); + }); + + it('does NOT dedup a distinct submission (e.g. regenerate reuses the user message id)', async () => { + await manager.claimGeneration('user-1', 'req-1', 'stream-a', 'convo-a'); + // A regenerate is a fresh ask() → fresh clientRequestId, so it must start its own generation. + const regenerate = await manager.claimGeneration('user-1', 'req-2', 'stream-a', 'convo-a'); + expect(regenerate).toEqual({ claimed: true }); + }); + + it('scopes claims per user', async () => { + await manager.claimGeneration('user-1', 'req-1', 'stream-a', 'convo-a'); + const otherUser = await manager.claimGeneration('user-2', 'req-1', 'stream-z', 'convo-z'); + expect(otherUser).toEqual({ claimed: true }); + }); + + it('allows a fresh claim after release', async () => { + await manager.claimGeneration('user-1', 'req-1', 'stream-a', 'convo-a'); + await manager.releaseGeneration('user-1', 'req-1'); + const again = await manager.claimGeneration('user-1', 'req-1', 'stream-c', 'convo-c'); + expect(again).toEqual({ claimed: true }); + }); +}); diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index a56e4c6de3..1fae66e510 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -8,6 +8,8 @@ import type { IJobStore, JobStatus, JobStatusTransition, + IdempotencyClaimValue, + IdempotencyClaimResult, } from '~/stream/interfaces/IJobStore'; import { STEER_ENQUEUE_NOT_RUNNING, @@ -64,6 +66,12 @@ export class InMemoryJobStore implements IJobStore { * default completeJob path deletes the job record immediately). */ private parkedSteers = new Map(); + /** Maps idempotency key -> claimed stream + expiry, deduping retried start requests. */ + private idempotencyClaims = new Map< + string, + { value: IdempotencyClaimValue; expiresAt: number } + >(); + /** Time to keep completed jobs before cleanup (0 = immediate) */ private ttlAfterComplete = 0; @@ -191,6 +199,24 @@ export class InMemoryJobStore implements IJobStore { return true; } + async claimIdempotencyKey( + key: string, + value: IdempotencyClaimValue, + ttlSeconds: number, + ): Promise { + const now = Date.now(); + const existing = this.idempotencyClaims.get(key); + if (existing && existing.expiresAt > now) { + return { claimed: false, existing: existing.value }; + } + this.idempotencyClaims.set(key, { value, expiresAt: now + ttlSeconds * 1000 }); + return { claimed: true }; + } + + async releaseIdempotencyKey(key: string): Promise { + this.idempotencyClaims.delete(key); + } + async deleteJob(streamId: string): Promise { this.jobs.delete(streamId); this.contentState.delete(streamId); @@ -237,6 +263,14 @@ export class InMemoryJobStore implements IJobStore { } } + // Idempotency keys are unique per submission, so expired claims are never + // overwritten — prune them here to keep the map bounded. + for (const [key, claim] of this.idempotencyClaims) { + if (claim.expiresAt <= now) { + this.idempotencyClaims.delete(key); + } + } + for (const [streamId, job] of this.jobs) { const isFinished = ['complete', 'error', 'aborted'].includes(job.status); if (isFinished && job.completedAt) { @@ -366,6 +400,7 @@ export class InMemoryJobStore implements IJobStore { this.steerQueues.clear(); this.closedSteerQueues.clear(); this.parkedSteers.clear(); + this.idempotencyClaims.clear(); logger.debug('[InMemoryJobStore] Destroyed'); } diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index d8a6acaccc..fb4a5b0aab 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -10,6 +10,8 @@ import type { IJobStore, JobStatus, JobStatusTransition, + IdempotencyClaimValue, + IdempotencyClaimResult, } from '~/stream/interfaces/IJobStore'; import { STEER_ENQUEUE_NOT_RUNNING, @@ -46,6 +48,18 @@ const JOB_CAS_LUA = 'redis.call("EXPIRE", KEYS[1], ttl) ' + 'return 1'; +/** + * Atomic idempotency claim. Single-key `SET NX PX`: returns nil when this caller + * won the claim, or the already-stored stream JSON when a prior request holds it. + * Touches ONLY KEYS[1], so it is atomic on single-node and Redis Cluster. + * + * KEYS: [idempotency] + * ARGV: [valueJson, ttlMs] + */ +const IDEMPOTENCY_CLAIM_LUA = + 'if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", tonumber(ARGV[2])) then return false end ' + + 'return redis.call("GET", KEYS[1])'; + /** * Atomic job (re)creation for the two same-slot keys: reset the steer queue * and write the job hash in ONE script. A `/chat/steer` request can then @@ -280,6 +294,8 @@ const KEYS = { /** User's active jobs set, tenant-qualified when tenantId is available */ userJobs: (userId: string, tenantId?: string) => tenantId ? `stream:user:{${tenantId}:${userId}}:jobs` : `stream:user:{${userId}}:jobs`, + /** Idempotency claim for a start-generation request: stream:idem:{userId:clientRequestId} */ + idempotency: (key: string) => `stream:idem:{${key}}`, }; /** @@ -691,6 +707,33 @@ export class RedisJobStore implements IJobStore { return true; } + async claimIdempotencyKey( + key: string, + value: IdempotencyClaimValue, + ttlSeconds: number, + ): Promise { + const result = await this.redis.eval( + IDEMPOTENCY_CLAIM_LUA, + 1, + KEYS.idempotency(key), + JSON.stringify(value), + String(ttlSeconds * 1000), + ); + if (result == null) { + return { claimed: true }; + } + try { + return { claimed: false, existing: JSON.parse(result as string) as IdempotencyClaimValue }; + } catch { + // Unreachable in practice (we wrote the JSON); proceed rather than dedup to a broken target. + return { claimed: false }; + } + } + + async releaseIdempotencyKey(key: string): Promise { + await this.redis.del(KEYS.idempotency(key)); + } + private async updateExistingJobHash(key: string, fields: string[]): Promise { const updated = await this.redis.eval( 'if redis.call("EXISTS", KEYS[1]) == 1 then redis.call("HSET", KEYS[1], unpack(ARGV)) return 1 else return 0 end', diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 1e4de99cc4..1f48f6b847 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -191,6 +191,24 @@ export interface JobStatusTransition { expectActionId?: string; } +/** Value stored under an idempotency claim: the stream a retried request should attach to. */ +export interface IdempotencyClaimValue { + streamId: string; + conversationId: string; + /** Epoch ms the claim was written — lets a losing duplicate tell a winner that is still + * starting (recent, no job yet → retry) from one that already finished and was cleaned + * up (old, no job → attach and let the client refetch). */ + claimedAt?: number; +} + +/** Result of an atomic {@link IJobStore.claimIdempotencyKey} attempt. */ +export interface IdempotencyClaimResult { + /** True when this caller won the claim and should create the job. */ + claimed: boolean; + /** When `claimed` is false, the stream the original request is already driving. */ + existing?: IdempotencyClaimValue; +} + /** * Usage metadata for token spending across different LLM providers. * @@ -363,6 +381,32 @@ export interface IJobStore { */ transitionStatus(streamId: string, args: JobStatusTransition): Promise; + /** + * Atomically claim an idempotency key so a retried start-generation request + * attaches to the original stream instead of starting a second billed + * generation. The first caller gets `{ claimed: true }` and should create the + * job; a later caller for the same key gets `{ claimed: false, existing }` + * carrying the stream the original request is already driving. + * + * Atomicity: single-key `SET NX` on Redis (one hash slot, cluster-safe) / + * check-and-set on the single-threaded in-memory store. + * + * @param key - Caller-scoped key, e.g. `${userId}:${clientRequestId}`. + * @param value - The stream a duplicate request should attach to. + * @param ttlSeconds - Claim lifetime; outlive the generation so a late retry still dedups. + */ + claimIdempotencyKey( + key: string, + value: IdempotencyClaimValue, + ttlSeconds: number, + ): Promise; + + /** + * Release a previously-claimed idempotency key so the submission can be retried + * (e.g. the start failed before generation began). No-op if the key is absent. + */ + releaseIdempotencyKey(key: string): Promise; + /** Delete a job */ deleteJob(streamId: string): Promise; diff --git a/packages/data-provider/src/createPayload.ts b/packages/data-provider/src/createPayload.ts index da20e7d1ac..c31e22aef3 100644 --- a/packages/data-provider/src/createPayload.ts +++ b/packages/data-provider/src/createPayload.ts @@ -24,6 +24,7 @@ export default function createPayload(submission: t.TSubmission) { ephemeralAgent, endpointOption, manualSkills, + clientRequestId, } = submission; const { conversationId } = s.tConvoUpdateSchema.parse(conversation); const { endpoint: _e, endpointType } = endpointOption as { @@ -52,6 +53,7 @@ export default function createPayload(submission: t.TSubmission) { ephemeralAgent: s.isAssistantsEndpoint(endpoint) ? undefined : ephemeralAgent, manualSkills: s.isAssistantsEndpoint(endpoint) ? undefined : manualSkills, timezone: getUserTimezone(), + clientRequestId, }; return { server, payload }; diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 9c9b8cb71e..e89ebaf51e 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -139,6 +139,13 @@ export type TPayload = Partial & manualSkills?: string[]; /** Browser IANA timezone (e.g. `America/New_York`) used to resolve local-time prompt variables server-side. */ timezone?: string; + /** + * Stable per-submission idempotency key (uuid) generated once per `ask()`. Identical + * across the client's start-generation network retries, unique per user action (including + * regenerate). The server dedups retried start requests on it so a lost/reset response + * cannot trigger a second billed generation. + */ + clientRequestId?: string; }; export type TEditedContent = @@ -172,6 +179,8 @@ export type TSubmission = { addedConvo?: TConversation; /** Skills the user invoked via the `$` popover for this submission. */ manualSkills?: string[]; + /** Stable per-submission idempotency key (uuid) forwarded to the server to dedup retried start-generation requests. */ + clientRequestId?: string; }; export type EventSubmission = Omit & { initialResponse: TMessage };