diff --git a/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js new file mode 100644 index 0000000000..e7281ee2ed --- /dev/null +++ b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js @@ -0,0 +1,153 @@ +const express = require('express'); +const request = require('supertest'); + +const mockHasGenerationClaim = jest.fn(); +const mockIpLimiter = jest.fn((_req, res) => res.status(429).json({ limited: 'ip' })); +const mockUserLimiter = jest.fn((_req, res) => res.status(429).json({ limited: 'user' })); +const mockRetryLimiter = jest.fn((_req, _res, next) => next()); +const mockRetryProbeLimiter = jest.fn((_req, _res, next) => next()); +const mockExemptAgentTrigger = jest.fn(() => false); +const mockExemptSchedule = jest.fn(() => false); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + debug: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + info: jest.fn(), + }, +})); + +jest.mock('@librechat/api', () => ({ + isEnabled: jest.fn(() => true), + detectGenerationRetry: async (req, _res, next) => { + req._isConfirmedGenerationRetry = await mockHasGenerationClaim( + req.user?.id, + req.body?.clientRequestId, + ); + next(); + }, + isConfirmedGenerationRetry: (req) => req._isConfirmedGenerationRetry === true, + generationRetryProbeLimiter: (...args) => mockRetryProbeLimiter(...args), + generationRetryLimiter: (...args) => mockRetryLimiter(...args), + isAgentTriggerRequest: jest.fn(() => false), + captureScheduleFireContext: jest.fn(), + exemptAgentTriggerFromIpLimiter: (...args) => mockExemptAgentTrigger(...args), + exemptFromUserLimiter: (...args) => mockExemptSchedule(...args), + createMessageFilterPii: jest.fn(() => (_req, _res, next) => next()), +})); + +jest.mock('~/server/middleware', () => ({ + uaParser: (_req, _res, next) => next(), + checkBan: (_req, _res, next) => next(), + requireJwtAuth: (req, _res, next) => { + req.user = { id: 'user-1' }; + next(); + }, + moderateText: (_req, _res, next) => next(), + messageIpLimiter: (...args) => mockIpLimiter(...args), + configMiddleware: (_req, _res, next) => next(), + messageUserLimiter: (...args) => mockUserLimiter(...args), +})); + +jest.mock('~/server/routes/agents/chat', () => { + const router = require('express').Router(); + router.post('/', (_req, res) => res.status(201).json({ admitted: true })); + return router; +}); +jest.mock('~/server/routes/agents/v1', () => ({ + v1: require('express').Router(), +})); +jest.mock('~/server/routes/agents/openai', () => require('express').Router()); +jest.mock('~/server/routes/agents/responses', () => require('express').Router()); +jest.mock('~/server/controllers/agents/steer', () => { + const controller = (_req, _res, next) => next(); + controller.SteerDeliveryController = (_req, _res, next) => next(); + controller.SteerCancelController = (_req, _res, next) => next(); + controller.SteerArmController = (_req, _res, next) => next(); + return controller; +}); +jest.mock('~/models', () => ({})); +jest.mock('~/server/services/Schedules', () => ({})); + +const agentsRouter = require('../index'); +const app = express(); +app.use(express.json()); +app.use('/agents', agentsRouter); + +describe('start-generation idempotency before message limiters', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockExemptAgentTrigger.mockReturnValue(false); + mockExemptSchedule.mockReturnValue(false); + }); + + it('lets a confirmed retry reach the controller without consuming either limiter', async () => { + mockHasGenerationClaim.mockResolvedValue(true); + + const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-1' }); + + expect(response.status).toBe(201); + expect(mockRetryProbeLimiter).toHaveBeenCalledTimes(1); + expect(mockRetryLimiter).toHaveBeenCalledTimes(1); + expect(mockIpLimiter).not.toHaveBeenCalled(); + expect(mockUserLimiter).not.toHaveBeenCalled(); + }); + + it('keeps a new submission behind the configured message limiters', async () => { + mockHasGenerationClaim.mockResolvedValue(false); + + const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-2' }); + + expect(response.status).toBe(429); + expect(response.body).toEqual({ limited: 'ip' }); + expect(mockRetryProbeLimiter).toHaveBeenCalledTimes(1); + expect(mockRetryLimiter).toHaveBeenCalledTimes(1); + expect(mockIpLimiter).toHaveBeenCalledTimes(1); + expect(mockUserLimiter).not.toHaveBeenCalled(); + }); + + it('defers an excessive confirmed retry before the chat pipeline', async () => { + mockHasGenerationClaim.mockResolvedValue(true); + mockRetryLimiter.mockImplementationOnce((_req, res) => + res.status(503).json({ code: 'SERVER_NOT_READY' }), + ); + + const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-3' }); + + expect(response.status).toBe(503); + expect(response.body.code).toBe('SERVER_NOT_READY'); + expect(mockIpLimiter).not.toHaveBeenCalled(); + expect(mockUserLimiter).not.toHaveBeenCalled(); + }); + + it('bounds candidate probes before durable storage inspection', async () => { + mockRetryProbeLimiter.mockImplementationOnce((_req, res) => + res.status(503).json({ code: 'SERVER_NOT_READY' }), + ); + + const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-5' }); + + expect(response.status).toBe(503); + expect(mockHasGenerationClaim).not.toHaveBeenCalled(); + expect(mockRetryLimiter).not.toHaveBeenCalled(); + expect(mockIpLimiter).not.toHaveBeenCalled(); + expect(mockUserLimiter).not.toHaveBeenCalled(); + }); + + it.each([ + ['an agent-trigger delivery', mockExemptAgentTrigger], + ['a scheduled delivery', mockExemptSchedule], + ])('keeps %s outside the human retry bucket', async (_label, exemption) => { + mockHasGenerationClaim.mockResolvedValue(true); + exemption.mockReturnValue(true); + + const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-4' }); + + expect(response.status).toBe(201); + expect(mockRetryProbeLimiter).not.toHaveBeenCalled(); + expect(mockRetryLimiter).not.toHaveBeenCalled(); + expect(mockIpLimiter).not.toHaveBeenCalled(); + expect(mockUserLimiter).not.toHaveBeenCalled(); + }); +}); diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 460d960adc..b6c237708a 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -18,6 +18,10 @@ const { exemptAgentTriggerFromIpLimiter, captureScheduleFireContext, exemptFromUserLimiter: exemptScheduleFromUserLimiter, + detectGenerationRetry, + isConfirmedGenerationRetry, + generationRetryProbeLimiter, + generationRetryLimiter, } = require('@librechat/api'); const { createSseStreamTelemetry } = require('@librechat/api/telemetry'); const { logger } = require('@librechat/data-schemas'); @@ -1075,14 +1079,41 @@ router.post( router.use('/', v1); const chatRouter = express.Router(); +const useMessageIpLimiter = isEnabled(LIMIT_MESSAGE_IP); +const useMessageUserLimiter = isEnabled(LIMIT_MESSAGE_USER); chatRouter.use(configMiddleware); - -if (isEnabled(LIMIT_MESSAGE_IP)) { - chatRouter.use(unless(exemptAgentTriggerFromIpLimiter, messageIpLimiter)); +if (useMessageIpLimiter || useMessageUserLimiter) { + chatRouter.use( + unless( + (req) => exemptAgentTriggerFromIpLimiter(req) || exemptScheduleFromUserLimiter(req), + generationRetryProbeLimiter, + ), + ); + chatRouter.use(detectGenerationRetry); + chatRouter.use( + unless( + (req) => exemptAgentTriggerFromIpLimiter(req) || exemptScheduleFromUserLimiter(req), + generationRetryLimiter, + ), + ); } -if (isEnabled(LIMIT_MESSAGE_USER)) { - chatRouter.use(unless(exemptScheduleFromUserLimiter, messageUserLimiter)); +if (useMessageIpLimiter) { + chatRouter.use( + unless( + (req) => exemptAgentTriggerFromIpLimiter(req) || isConfirmedGenerationRetry(req), + messageIpLimiter, + ), + ); +} + +if (useMessageUserLimiter) { + chatRouter.use( + unless( + (req) => exemptScheduleFromUserLimiter(req) || isConfirmedGenerationRetry(req), + messageUserLimiter, + ), + ); } chatRouter.use('/', chat); diff --git a/package-lock.json b/package-lock.json index e3d95633c0..bdddde990b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -42783,6 +42783,7 @@ "@langchain/langgraph-checkpoint-mongodb": "^1.4.0", "cluster-key-slot": "^1.1.2", "croner": "^10.0.1", + "express-rate-limit": "^8.5.1", "helmet": "^8.3.0", "proxy-from-env": "^2.1.0", "re2js": "^2.8.6" diff --git a/packages/api/package.json b/packages/api/package.json index 63d36cbc53..0a7a950209 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -171,6 +171,7 @@ "@langchain/langgraph-checkpoint-mongodb": "^1.4.0", "cluster-key-slot": "^1.1.2", "croner": "^10.0.1", + "express-rate-limit": "^8.5.1", "helmet": "^8.3.0", "proxy-from-env": "^2.1.0", "re2js": "^2.8.6" diff --git a/packages/api/src/middleware/generationRetry.spec.ts b/packages/api/src/middleware/generationRetry.spec.ts new file mode 100644 index 0000000000..72970e5eaf --- /dev/null +++ b/packages/api/src/middleware/generationRetry.spec.ts @@ -0,0 +1,144 @@ +import express from 'express'; +import request from 'supertest'; +import { logger } from '@librechat/data-schemas'; +import type { NextFunction, Request, Response } from 'express'; +import { + detectGenerationRetry, + generationRetryLimiter, + generationRetryProbeLimiter, + GENERATION_RETRY_MAX, + GENERATION_RETRY_PROBE_MAX, + isConfirmedGenerationRetry, +} from './generationRetry'; +import { GenerationJobManager } from '~/stream/GenerationJobManager'; + +function generationRequest(overrides: Partial = {}): Request { + return { + method: 'POST', + path: '/', + body: { clientRequestId: 'request-1' }, + user: { id: 'user-1' }, + ...overrides, + } as Request; +} + +describe('generation retry admission', () => { + afterEach(() => { + jest.restoreAllMocks(); + }); + + it('marks only a submission with an existing durable claim as a retry', async () => { + const hasClaim = jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(true); + const req = generationRequest(); + const next = jest.fn() as NextFunction; + + await detectGenerationRetry(req, {} as Response, next); + + expect(hasClaim).toHaveBeenCalledWith('user-1', 'request-1'); + expect(isConfirmedGenerationRetry(req)).toBe(true); + expect(next).toHaveBeenCalledTimes(1); + }); + + it('leaves a new submission subject to the ordinary message limiters', async () => { + jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(false); + const req = generationRequest(); + + await detectGenerationRetry(req, {} as Response, jest.fn()); + + expect(isConfirmedGenerationRetry(req)).toBe(false); + }); + + it.each([ + ['a resume', { path: '/resume' }], + ['a resume with a trailing slash', { path: '/resume/' }], + ['a case-insensitive resume route', { path: '/Resume' }], + ['a request without an authenticated user', { user: undefined }], + ['an invalid idempotency key', { body: { clientRequestId: 'invalid key' } }], + ])('does not probe %s', async (_label, overrides) => { + const hasClaim = jest.spyOn(GenerationJobManager, 'hasGenerationClaim'); + const req = generationRequest(overrides as Partial); + const next = jest.fn() as NextFunction; + + await detectGenerationRetry(req, {} as Response, next); + + expect(hasClaim).not.toHaveBeenCalled(); + expect(isConfirmedGenerationRetry(req)).toBe(false); + expect(next).toHaveBeenCalledTimes(1); + }); + + it('fails closed to the ordinary limiters when the claim probe is unavailable', async () => { + jest + .spyOn(GenerationJobManager, 'hasGenerationClaim') + .mockRejectedValue(new Error('store unavailable')); + const warn = jest.spyOn(logger, 'warn').mockImplementation(() => logger); + const req = generationRequest(); + const next = jest.fn() as NextFunction; + + await detectGenerationRetry(req, {} as Response, next); + + expect(isConfirmedGenerationRetry(req)).toBe(false); + expect(warn).toHaveBeenCalledWith( + '[GenerationIdempotency] Failed to inspect start-generation claim', + expect.objectContaining({ userId: 'user-1', clientRequestId: 'request-1' }), + ); + expect(next).toHaveBeenCalledTimes(1); + }); + + it('bounds fresh claim probes before accessing the shared store', async () => { + const hasClaim = jest + .spyOn(GenerationJobManager, 'hasGenerationClaim') + .mockResolvedValue(false); + const downstream = jest.fn((_req, res) => res.sendStatus(204)); + const app = express(); + app.use(express.json()); + app.use((req, _res, next) => { + req.user = { id: 'bounded-probe-user' }; + next(); + }); + app.use(generationRetryProbeLimiter); + app.use(detectGenerationRetry); + app.post('/', downstream); + + for (let attempt = 0; attempt < GENERATION_RETRY_PROBE_MAX; attempt += 1) { + await request(app) + .post('/') + .send({ clientRequestId: `probe-${attempt}` }) + .expect(204); + } + const rejected = await request(app) + .post('/') + .send({ clientRequestId: 'probe-rejected' }) + .expect(503); + + expect(rejected.headers['retry-after']).toBeDefined(); + expect(rejected.body.code).toBe('SERVER_NOT_READY'); + expect(hasClaim).toHaveBeenCalledTimes(GENERATION_RETRY_PROBE_MAX); + expect(downstream).toHaveBeenCalledTimes(GENERATION_RETRY_PROBE_MAX); + }); + + it('makes a bounded confirmed retry delay participate in readiness recovery', async () => { + jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(true); + const downstream = jest.fn((_req, res) => res.sendStatus(204)); + const app = express(); + app.use(express.json()); + app.use((req, _res, next) => { + req.user = { id: 'bounded-retry-user' }; + next(); + }); + app.use(detectGenerationRetry); + app.use(generationRetryLimiter); + app.post('/', downstream); + + for (let attempt = 0; attempt < GENERATION_RETRY_MAX; attempt += 1) { + await request(app).post('/').send({ clientRequestId: 'bounded-request' }).expect(204); + } + const rejected = await request(app) + .post('/') + .send({ clientRequestId: 'bounded-request' }) + .expect(503); + + expect(rejected.headers['retry-after']).toBeDefined(); + expect(rejected.body.code).toBe('SERVER_NOT_READY'); + expect(downstream).toHaveBeenCalledTimes(GENERATION_RETRY_MAX); + }); +}); diff --git a/packages/api/src/middleware/generationRetry.ts b/packages/api/src/middleware/generationRetry.ts new file mode 100644 index 0000000000..b4e6416373 --- /dev/null +++ b/packages/api/src/middleware/generationRetry.ts @@ -0,0 +1,109 @@ +import { rateLimit } from 'express-rate-limit'; +import { logger } from '@librechat/data-schemas'; +import type { NextFunction, Request, RequestHandler, Response } from 'express'; +import { GenerationJobManager } from '~/stream/GenerationJobManager'; +import { limiterCache } from '~/cache/cacheFactory'; + +const CLIENT_REQUEST_ID_PATTERN = /^[A-Za-z0-9:_-]{1,128}$/; +const confirmedGenerationRetry: unique symbol = Symbol('confirmedGenerationRetry'); + +export const GENERATION_RETRY_WINDOW_MS = 60_000; +export const GENERATION_RETRY_MAX = 10; +export const GENERATION_RETRY_PROBE_MAX = 60; + +type GenerationRetryRequest = Request & { + user?: { id?: string }; + [confirmedGenerationRetry]?: boolean; +}; + +function isGenerationRetryCandidate(req: GenerationRetryRequest): boolean { + const clientRequestId = req.body?.clientRequestId; + const normalizedPath = req.path.replace(/\/+$/, '').toLowerCase(); + return ( + req.method === 'POST' && + normalizedPath !== '/resume' && + typeof req.user?.id === 'string' && + typeof clientRequestId === 'string' && + CLIENT_REQUEST_ID_PATTERN.test(clientRequestId) + ); +} + +/** + * Classifies only retries already represented by a durable generation claim. + * The controller still owns the authoritative claim/read transition. + */ +export async function detectGenerationRetry( + req: GenerationRetryRequest, + _res: Response, + next: NextFunction, +): Promise { + if (!isGenerationRetryCandidate(req)) { + next(); + return; + } + + const clientRequestId = req.body?.clientRequestId; + const userId = req.user?.id; + if (typeof userId !== 'string' || typeof clientRequestId !== 'string') { + next(); + return; + } + + try { + req[confirmedGenerationRetry] = await GenerationJobManager.hasGenerationClaim( + userId, + clientRequestId, + ); + } catch (error) { + logger.warn('[GenerationIdempotency] Failed to inspect start-generation claim', { + userId, + clientRequestId, + error: error instanceof Error ? error.message : String(error), + }); + } + next(); +} + +export function isConfirmedGenerationRetry(req: Request): boolean { + return (req as GenerationRetryRequest)[confirmedGenerationRetry] === true; +} + +const retryAdmissionHandler: RequestHandler = (_req, res) => { + res.status(503).type('application/json').json({ + code: 'SERVER_NOT_READY', + error: 'Generation retry admission is temporarily busy. Please retry shortly.', + }); +}; + +/** + * Bounds read-only claim probes before they touch the shared generation store. + * The retryable response participates in the client's existing 120-second + * readiness loop and express-rate-limit supplies its Retry-After header. + */ +export const generationRetryProbeLimiter: RequestHandler = rateLimit({ + windowMs: GENERATION_RETRY_WINDOW_MS, + max: GENERATION_RETRY_PROBE_MAX, + standardHeaders: true, + legacyHeaders: false, + skip: (req) => !isGenerationRetryCandidate(req as GenerationRetryRequest), + keyGenerator: (req) => String((req as GenerationRetryRequest).user?.id), + store: limiterCache('generation_retry_probe_limiter'), + handler: retryAdmissionHandler, +}); + +/** + * Confirmed retries bypass the ordinary message buckets so a lost response can + * be recovered, but they still receive a small user-scoped allowance before + * moderation and the rest of the request pipeline. This bounds replay costs + * while leaving the authoritative generation claim unchanged. + */ +export const generationRetryLimiter: RequestHandler = rateLimit({ + windowMs: GENERATION_RETRY_WINDOW_MS, + max: GENERATION_RETRY_MAX, + standardHeaders: true, + legacyHeaders: false, + skip: (req) => !isConfirmedGenerationRetry(req), + keyGenerator: (req) => String((req as GenerationRetryRequest).user?.id), + store: limiterCache('generation_retry_limiter'), + handler: retryAdmissionHandler, +}); diff --git a/packages/api/src/middleware/index.ts b/packages/api/src/middleware/index.ts index d6e06ea4dc..5b2ee6d8be 100644 --- a/packages/api/src/middleware/index.ts +++ b/packages/api/src/middleware/index.ts @@ -23,3 +23,4 @@ export * from './modelBoundContent'; export * from './messageFilterPii'; export * from './messageValidation'; export * from './feedback'; +export * from './generationRetry'; diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 699300e92b..8669b5fb4c 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -2923,6 +2923,21 @@ class GenerationJobManagerClass { return { claimed: false, existing: primary, source: 'primary' }; } + /** Checks the mixed-version admission key without creating or repairing a + * claim. This is deliberately weaker than `claimGeneration`: callers may + * use it only to exempt a confirmed retry from request rate limiting; the + * controller must still perform the authoritative claim transition. */ + async hasGenerationClaim(userId: string, clientRequestId: string): Promise { + if (!CLIENT_REQUEST_ID_PATTERN.test(clientRequestId)) { + return false; + } + return ( + (await this.jobStore.hasIdempotencyKey?.( + this.legacyGenerationClaimKey(userId, clientRequestId), + )) === true + ); + } + private generationClaimKey(userId: string, clientRequestId: string, streamId: string): string { return `{${streamId}}:${userId}:${clientRequestId}`; } 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 db9162721d..3259a910cb 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -4229,6 +4229,26 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('probes claim existence without creating a missing key', async () => { + if (!ioredisClient) { + return; + } + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const key = `user-1:req-probe-${Date.now()}`; + await expect(store.hasIdempotencyKey(key)).resolves.toBe(false); + + await store.claimIdempotencyKey(key, { streamId: 's1', conversationId: 'c1' }, 1200); + await expect(store.hasIdempotencyKey(key)).resolves.toBe(true); + + await store.releaseIdempotencyKey(key); + await expect(store.hasIdempotencyKey(key)).resolves.toBe(false); + + await store.destroy(); + }); + test('sets a bounded TTL on the claim', async () => { if (!ioredisClient) { return; diff --git a/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts index e9ed727c78..f1b0257451 100644 --- a/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts +++ b/packages/api/src/stream/__tests__/idempotencyClaim.spec.ts @@ -51,6 +51,33 @@ describe('InMemoryJobStore.claimIdempotencyKey', () => { expect(second).toEqual({ claimed: false, existing: { streamId: 's1', conversationId: 'c1' } }); }); + it('probes an existing claim without creating a missing key', async () => { + await expect(store.hasIdempotencyKey('user:missing')).resolves.toBe(false); + + await store.claimIdempotencyKey( + 'user:existing', + { streamId: 's1', conversationId: 'c1' }, + 1200, + ); + + await expect(store.hasIdempotencyKey('user:existing')).resolves.toBe(true); + await expect(store.hasIdempotencyKey('user:missing')).resolves.toBe(false); + }); + + it('does not report an expired claim as existing', async () => { + jest.useFakeTimers(); + try { + jest.setSystemTime(new Date('2026-08-29T00:00:00Z')); + await store.claimIdempotencyKey('user:expired', { streamId: 's1', conversationId: 'c1' }, 1); + await expect(store.hasIdempotencyKey('user:expired')).resolves.toBe(true); + + jest.setSystemTime(new Date('2026-08-29T00:00:02Z')); + await expect(store.hasIdempotencyKey('user:expired')).resolves.toBe(false); + } finally { + jest.useRealTimers(); + } + }); + it('lets a released key be claimed again', async () => { await store.claimIdempotencyKey('user:req', { streamId: 's1', conversationId: 'c1' }, 1200); await store.releaseIdempotencyKey('user:req'); @@ -197,6 +224,16 @@ describe('GenerationJobManager start-generation claim', () => { expect(typeof retry.existing?.claimedAt).toBe('number'); }); + it('detects only an already-claimed submission for pre-limiter retry admission', async () => { + await expect(manager.hasGenerationClaim('user-1', 'req-1')).resolves.toBe(false); + + await manager.claimGeneration('user-1', 'req-1', 'stream-a', 'convo-a'); + + await expect(manager.hasGenerationClaim('user-1', 'req-1')).resolves.toBe(true); + await expect(manager.hasGenerationClaim('user-1', 'req-2')).resolves.toBe(false); + await expect(manager.hasGenerationClaim('user-2', 'req-1')).resolves.toBe(false); + }); + it('claims the exact legacy key before the same-slot primary with staggered TTLs', async () => { const claimSpy = jest.spyOn(store, 'claimIdempotencyKey'); const result = await manager.claimGeneration( diff --git a/packages/api/src/stream/implementations/InMemoryJobStore.ts b/packages/api/src/stream/implementations/InMemoryJobStore.ts index f0ba55cfb7..888c708113 100644 --- a/packages/api/src/stream/implementations/InMemoryJobStore.ts +++ b/packages/api/src/stream/implementations/InMemoryJobStore.ts @@ -892,6 +892,18 @@ export class InMemoryJobStore implements IJobStoreV2 { return { claimed: true, existing: value }; } + async hasIdempotencyKey(key: string): Promise { + const existing = this.idempotencyClaims.get(key); + if (existing == null) { + return false; + } + if (existing.expiresAt > Date.now()) { + return true; + } + this.idempotencyClaims.delete(key); + return false; + } + async takeoverIdempotencyKey( key: string, expected: IdempotencyClaimValue, diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index f339fe2e49..73a01b034b 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -2713,6 +2713,10 @@ export class RedisJobStore implements IJobStoreV2 { } } + async hasIdempotencyKey(key: string): Promise { + return (await this.redis.exists(KEYS.idempotency(key))) === 1; + } + async takeoverIdempotencyKey( key: string, expected: IdempotencyClaimValue, diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index e23238b483..a885ae6371 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -885,6 +885,11 @@ export interface IJobStore { ): Promise; releaseIdempotencyKey(key: string): Promise; + /** Read-only existence probe used to identify a confirmed retry before + * request-rate admission. Optional stores keep the conservative behavior + * where every request remains subject to the limiter. */ + hasIdempotencyKey?(key: string): Promise; + deleteJob(streamId: string, expectedCreatedAt?: number): Promise; hasJob(streamId: string): Promise; getRunningJobs(): Promise;