mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
🎟️ fix: Claim Idempotency Keys to Dedup Retried Generation Requests and Prevent Double Billing (#14344)
* 🐛 fix: Dedup retried start-generation requests to prevent duplicate billing A lost or reset start-generation response makes the client re-POST the identical payload (up to 3x on network errors). The resumable-stream controller had no idempotency: createJob unconditionally overwrote the running job without aborting the prior one, so both requests ran full LLM completions and both billed while the UI showed only one (#14339). Add a stable per-submission clientRequestId (uuid, fresh per ask() so a regenerate differs, reused across the start-generation retries) and an atomic claim on the job store keyed by userId:clientRequestId. The first request wins and generates; a retried POST loses the claim and receives the original stream, which the client subscribes to and replays - no second billed generation. - IJobStore.claimIdempotencyKey/releaseIdempotencyKey (in-memory Map+TTL, Redis single-key SET NX PX + GET Lua, cluster-safe) - GenerationJobManager.claimGeneration/releaseGeneration (20m TTL) - Controller claims before the concurrency check, dedups with a resumed response, releases on start-failure/429 - clientRequestId threaded through TSubmission/TPayload/createPayload * 🐛 fix: Harden start-generation dedup (Codex review) Address three P2 findings on the idempotency path: - Resume replay: a deduped retry now subscribes with resume=true so the client replays prior content and any pending-action from the running stream instead of only live events (cross-replica / HITL correctness). startGeneration returns { streamId, resumed } and the response's status:'resumed' drives the subscribe mode. - Wait for the job record: a duplicate that loses the claim now waits briefly for the winner to create the job before returning the stream (a stream with no job 404s terminally). If the winner has not materialized, return 503 SERVER_NOT_READY so the client retries via the existing readiness path instead of attaching to a dead stream. - Release only owned claims: track whether the request actually won the claim; the 429 and init-error paths no longer release a claim owned by another in-flight generation (fail-open path could erase it and re-enable double billing). Adds controller tests covering dedup, the 503 race fallback, win-then- create, and claim-release ownership on 429 / fail-open. * 🐛 fix: Don't trap deduped retries on missing job records (Codex review) The previous round returned 503 SERVER_NOT_READY when a deduped retry's job record was absent. But a missing job usually means the original generation already completed and was cleaned up (cleanupOnComplete) — the correct recovery is to return the stream and let the client's subscribe 404 handler refetch the persisted messages. The 503 instead trapped the send in a readiness-retry loop until the client's window expired. Keep the bounded wait (it still covers the job-about-to-be-created race) but always return the resumed stream afterward; a gone/never-created job recovers via the client's existing 404 path instead of being treated as indefinitely starting. Updated the controller test accordingly. * 🐛 fix: Gate deduped resume on claim age, not just job presence (Codex review) Removing the 503 entirely (previous round) reintroduced the inverse race: if the winning request stalls between claimGeneration and createJob, a losing duplicate saw no job, returned status:'resumed' anyway, and the client subscribed to a stream that did not exist yet — the 404 handler tore the turn down while the winner went on to generate and bill with no UI attached. Distinguish the two missing-job cases by claim age (claimedAt now travels on the claim value): - fresh claim, no job yet → winner is still starting → 503 SERVER_NOT_READY so the client retries via the readiness path (bounded, not indefinite). - old claim, no job → the original already completed and was cleaned up (or the winner died) → attach; the client's 404 handler refetches. Tests cover both age branches. * 🐛 fix: Scope dedup fail-open + keep resumed convos on 404 (Codex review) - Fail-open only on claim acquisition: a store error while checking an already-confirmed existing claim no longer falls through to createJob (which would start a second billed generation during a Redis hiccup). Once claim.existing is known, a job-lookup error returns 503 retry. - Don't drop a resumed convo on 404: the optimistic-conversation cleanup in useResumableSSE now runs only for fresh (non-resume) subscribes. A deduped resume whose original completed and was cleaned up 404s, but its conversation is persisted and must stay in the sidebar. Adds a controller test for the job-lookup-error path (503, no createJob). * 🐛 fix: Reconcile resumed convos on 404 instead of guessing (Codex review) Round-4's !isResume guard fixed the completed-and-cleaned case (don't drop a persisted convo) but left the inverse: a new-conversation retry deduped to a claim whose original worker died before persisting still resumes, 404s, and — with removal skipped — leaves a phantom /c/<streamId> sidebar entry. Stop guessing keep-vs-remove on a resume 404. Reconcile against the server: invalidate the conversations list so a real (persisted) convo stays and a phantom is dropped. Fresh (non-resume) optimistic streams still prune immediately. Adds a client test for the resume path. * 🐛 fix: Finalize failed job before releasing its claim (Codex review) In the initialization-error catch, the idempotency claim was released before completeJob(streamId). A racing retry could win the released key and createJob() the same streamId while this catch was still running, and completeJob() (not guarded by the original createdAt) would then abort the replacement. Finalize the failed job first, then release the claim. Adds a controller test asserting completeJob precedes releaseGeneration. * 🐛 fix: Clear claims on destroy + survive completeJob failure (Codex review) - InMemoryJobStore.destroy() now clears the idempotencyClaims map, so a reused/reconfigured store instance doesn't dedup a fresh start against a torn-down job's stale claim. - Init-error cleanup: completeJob() is swallowed so a store-hiccup rejection can no longer skip the idempotency-key release and the pending-request decrement (which would wedge the retry behind the claim and leak the concurrency slot). A failed completeJob finalized nothing, so releasing afterward still can't abort a later replacement. Tests: claims cleared on destroy; release + pending decrement still run when completeJob rejects.
This commit is contained in:
parent
3171b86413
commit
9e245aced4
13 changed files with 908 additions and 16 deletions
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
||||
|
|
|
|||
|
|
@ -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<string | null> => {
|
||||
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');
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<IdempotencyClaimResult> {
|
||||
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<void> {
|
||||
await this.jobStore.releaseIdempotencyKey(`${userId}:${clientRequestId}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get job status.
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
156
packages/api/src/stream/__tests__/idempotencyClaim.spec.ts
Normal file
156
packages/api/src/stream/__tests__/idempotencyClaim.spec.ts
Normal file
|
|
@ -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 });
|
||||
});
|
||||
});
|
||||
|
|
@ -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<string, { payload: string; expiresAt: number }>();
|
||||
|
||||
/** 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<IdempotencyClaimResult> {
|
||||
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<void> {
|
||||
this.idempotencyClaims.delete(key);
|
||||
}
|
||||
|
||||
async deleteJob(streamId: string): Promise<void> {
|
||||
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');
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<IdempotencyClaimResult> {
|
||||
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<void> {
|
||||
await this.redis.del(KEYS.idempotency(key));
|
||||
}
|
||||
|
||||
private async updateExistingJobHash(key: string, fields: string[]): Promise<boolean> {
|
||||
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',
|
||||
|
|
|
|||
|
|
@ -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<boolean>;
|
||||
|
||||
/**
|
||||
* 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<IdempotencyClaimResult>;
|
||||
|
||||
/**
|
||||
* 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<void>;
|
||||
|
||||
/** Delete a job */
|
||||
deleteJob(streamId: string): Promise<void>;
|
||||
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
|
|
|
|||
|
|
@ -139,6 +139,13 @@ export type TPayload = Partial<TMessage> &
|
|||
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<TSubmission, 'initialResponse'> & { initialResponse: TMessage };
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue