Skip to content

Commit 7a00615

Browse files
authored
🪪 fix: Admit Confirmed Generation Retries Before Message Limits (#15341)
* fix: admit idempotent retries before message limits * fix: bound generation retry admission * fix: keep retry probe store-compatible * fix: exclude normalized resume routes * fix: preserve trusted retry exemptions * fix: bound retry claim admission * style: satisfy generation retry static checks
1 parent 773127b commit 7a00615

13 files changed

Lines changed: 537 additions & 4 deletions

File tree

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
const express = require('express');
2+
const request = require('supertest');
3+
4+
const mockHasGenerationClaim = jest.fn();
5+
const mockIpLimiter = jest.fn((_req, res) => res.status(429).json({ limited: 'ip' }));
6+
const mockUserLimiter = jest.fn((_req, res) => res.status(429).json({ limited: 'user' }));
7+
const mockRetryLimiter = jest.fn((_req, _res, next) => next());
8+
const mockRetryProbeLimiter = jest.fn((_req, _res, next) => next());
9+
const mockExemptAgentTrigger = jest.fn(() => false);
10+
const mockExemptSchedule = jest.fn(() => false);
11+
12+
jest.mock('@librechat/data-schemas', () => ({
13+
logger: {
14+
debug: jest.fn(),
15+
warn: jest.fn(),
16+
error: jest.fn(),
17+
info: jest.fn(),
18+
},
19+
}));
20+
21+
jest.mock('@librechat/api', () => ({
22+
isEnabled: jest.fn(() => true),
23+
detectGenerationRetry: async (req, _res, next) => {
24+
req._isConfirmedGenerationRetry = await mockHasGenerationClaim(
25+
req.user?.id,
26+
req.body?.clientRequestId,
27+
);
28+
next();
29+
},
30+
isConfirmedGenerationRetry: (req) => req._isConfirmedGenerationRetry === true,
31+
generationRetryProbeLimiter: (...args) => mockRetryProbeLimiter(...args),
32+
generationRetryLimiter: (...args) => mockRetryLimiter(...args),
33+
isAgentTriggerRequest: jest.fn(() => false),
34+
captureScheduleFireContext: jest.fn(),
35+
exemptAgentTriggerFromIpLimiter: (...args) => mockExemptAgentTrigger(...args),
36+
exemptFromUserLimiter: (...args) => mockExemptSchedule(...args),
37+
createMessageFilterPii: jest.fn(() => (_req, _res, next) => next()),
38+
}));
39+
40+
jest.mock('~/server/middleware', () => ({
41+
uaParser: (_req, _res, next) => next(),
42+
checkBan: (_req, _res, next) => next(),
43+
requireJwtAuth: (req, _res, next) => {
44+
req.user = { id: 'user-1' };
45+
next();
46+
},
47+
moderateText: (_req, _res, next) => next(),
48+
messageIpLimiter: (...args) => mockIpLimiter(...args),
49+
configMiddleware: (_req, _res, next) => next(),
50+
messageUserLimiter: (...args) => mockUserLimiter(...args),
51+
}));
52+
53+
jest.mock('~/server/routes/agents/chat', () => {
54+
const router = require('express').Router();
55+
router.post('/', (_req, res) => res.status(201).json({ admitted: true }));
56+
return router;
57+
});
58+
jest.mock('~/server/routes/agents/v1', () => ({
59+
v1: require('express').Router(),
60+
}));
61+
jest.mock('~/server/routes/agents/openai', () => require('express').Router());
62+
jest.mock('~/server/routes/agents/responses', () => require('express').Router());
63+
jest.mock('~/server/controllers/agents/steer', () => {
64+
const controller = (_req, _res, next) => next();
65+
controller.SteerDeliveryController = (_req, _res, next) => next();
66+
controller.SteerCancelController = (_req, _res, next) => next();
67+
controller.SteerArmController = (_req, _res, next) => next();
68+
return controller;
69+
});
70+
jest.mock('~/models', () => ({}));
71+
jest.mock('~/server/services/Schedules', () => ({}));
72+
73+
const agentsRouter = require('../index');
74+
const app = express();
75+
app.use(express.json());
76+
app.use('/agents', agentsRouter);
77+
78+
describe('start-generation idempotency before message limiters', () => {
79+
beforeEach(() => {
80+
jest.clearAllMocks();
81+
mockExemptAgentTrigger.mockReturnValue(false);
82+
mockExemptSchedule.mockReturnValue(false);
83+
});
84+
85+
it('lets a confirmed retry reach the controller without consuming either limiter', async () => {
86+
mockHasGenerationClaim.mockResolvedValue(true);
87+
88+
const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-1' });
89+
90+
expect(response.status).toBe(201);
91+
expect(mockRetryProbeLimiter).toHaveBeenCalledTimes(1);
92+
expect(mockRetryLimiter).toHaveBeenCalledTimes(1);
93+
expect(mockIpLimiter).not.toHaveBeenCalled();
94+
expect(mockUserLimiter).not.toHaveBeenCalled();
95+
});
96+
97+
it('keeps a new submission behind the configured message limiters', async () => {
98+
mockHasGenerationClaim.mockResolvedValue(false);
99+
100+
const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-2' });
101+
102+
expect(response.status).toBe(429);
103+
expect(response.body).toEqual({ limited: 'ip' });
104+
expect(mockRetryProbeLimiter).toHaveBeenCalledTimes(1);
105+
expect(mockRetryLimiter).toHaveBeenCalledTimes(1);
106+
expect(mockIpLimiter).toHaveBeenCalledTimes(1);
107+
expect(mockUserLimiter).not.toHaveBeenCalled();
108+
});
109+
110+
it('defers an excessive confirmed retry before the chat pipeline', async () => {
111+
mockHasGenerationClaim.mockResolvedValue(true);
112+
mockRetryLimiter.mockImplementationOnce((_req, res) =>
113+
res.status(503).json({ code: 'SERVER_NOT_READY' }),
114+
);
115+
116+
const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-3' });
117+
118+
expect(response.status).toBe(503);
119+
expect(response.body.code).toBe('SERVER_NOT_READY');
120+
expect(mockIpLimiter).not.toHaveBeenCalled();
121+
expect(mockUserLimiter).not.toHaveBeenCalled();
122+
});
123+
124+
it('bounds candidate probes before durable storage inspection', async () => {
125+
mockRetryProbeLimiter.mockImplementationOnce((_req, res) =>
126+
res.status(503).json({ code: 'SERVER_NOT_READY' }),
127+
);
128+
129+
const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-5' });
130+
131+
expect(response.status).toBe(503);
132+
expect(mockHasGenerationClaim).not.toHaveBeenCalled();
133+
expect(mockRetryLimiter).not.toHaveBeenCalled();
134+
expect(mockIpLimiter).not.toHaveBeenCalled();
135+
expect(mockUserLimiter).not.toHaveBeenCalled();
136+
});
137+
138+
it.each([
139+
['an agent-trigger delivery', mockExemptAgentTrigger],
140+
['a scheduled delivery', mockExemptSchedule],
141+
])('keeps %s outside the human retry bucket', async (_label, exemption) => {
142+
mockHasGenerationClaim.mockResolvedValue(true);
143+
exemption.mockReturnValue(true);
144+
145+
const response = await request(app).post('/agents/chat').send({ clientRequestId: 'request-4' });
146+
147+
expect(response.status).toBe(201);
148+
expect(mockRetryProbeLimiter).not.toHaveBeenCalled();
149+
expect(mockRetryLimiter).not.toHaveBeenCalled();
150+
expect(mockIpLimiter).not.toHaveBeenCalled();
151+
expect(mockUserLimiter).not.toHaveBeenCalled();
152+
});
153+
});

‎api/server/routes/agents/index.js‎

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,10 @@ const {
1818
exemptAgentTriggerFromIpLimiter,
1919
captureScheduleFireContext,
2020
exemptFromUserLimiter: exemptScheduleFromUserLimiter,
21+
detectGenerationRetry,
22+
isConfirmedGenerationRetry,
23+
generationRetryProbeLimiter,
24+
generationRetryLimiter,
2125
} = require('@librechat/api');
2226
const { createSseStreamTelemetry } = require('@librechat/api/telemetry');
2327
const { logger } = require('@librechat/data-schemas');
@@ -1075,14 +1079,41 @@ router.post(
10751079
router.use('/', v1);
10761080

10771081
const chatRouter = express.Router();
1082+
const useMessageIpLimiter = isEnabled(LIMIT_MESSAGE_IP);
1083+
const useMessageUserLimiter = isEnabled(LIMIT_MESSAGE_USER);
10781084
chatRouter.use(configMiddleware);
1085+
if (useMessageIpLimiter || useMessageUserLimiter) {
1086+
chatRouter.use(
1087+
unless(
1088+
(req) => exemptAgentTriggerFromIpLimiter(req) || exemptScheduleFromUserLimiter(req),
1089+
generationRetryProbeLimiter,
1090+
),
1091+
);
1092+
chatRouter.use(detectGenerationRetry);
1093+
chatRouter.use(
1094+
unless(
1095+
(req) => exemptAgentTriggerFromIpLimiter(req) || exemptScheduleFromUserLimiter(req),
1096+
generationRetryLimiter,
1097+
),
1098+
);
1099+
}
10791100

1080-
if (isEnabled(LIMIT_MESSAGE_IP)) {
1081-
chatRouter.use(unless(exemptAgentTriggerFromIpLimiter, messageIpLimiter));
1101+
if (useMessageIpLimiter) {
1102+
chatRouter.use(
1103+
unless(
1104+
(req) => exemptAgentTriggerFromIpLimiter(req) || isConfirmedGenerationRetry(req),
1105+
messageIpLimiter,
1106+
),
1107+
);
10821108
}
10831109

1084-
if (isEnabled(LIMIT_MESSAGE_USER)) {
1085-
chatRouter.use(unless(exemptScheduleFromUserLimiter, messageUserLimiter));
1110+
if (useMessageUserLimiter) {
1111+
chatRouter.use(
1112+
unless(
1113+
(req) => exemptScheduleFromUserLimiter(req) || isConfirmedGenerationRetry(req),
1114+
messageUserLimiter,
1115+
),
1116+
);
10861117
}
10871118

10881119
chatRouter.use('/', chat);

‎package-lock.json‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/api/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@
171171
"@langchain/langgraph-checkpoint-mongodb": "^1.4.0",
172172
"cluster-key-slot": "^1.1.2",
173173
"croner": "^10.0.1",
174+
"express-rate-limit": "^8.5.1",
174175
"helmet": "^8.3.0",
175176
"proxy-from-env": "^2.1.0",
176177
"re2js": "^2.8.6"
Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
import express from 'express';
2+
import request from 'supertest';
3+
import { logger } from '@librechat/data-schemas';
4+
import type { NextFunction, Request, Response } from 'express';
5+
import {
6+
detectGenerationRetry,
7+
generationRetryLimiter,
8+
generationRetryProbeLimiter,
9+
GENERATION_RETRY_MAX,
10+
GENERATION_RETRY_PROBE_MAX,
11+
isConfirmedGenerationRetry,
12+
} from './generationRetry';
13+
import { GenerationJobManager } from '~/stream/GenerationJobManager';
14+
15+
function generationRequest(overrides: Partial<Request> = {}): Request {
16+
return {
17+
method: 'POST',
18+
path: '/',
19+
body: { clientRequestId: 'request-1' },
20+
user: { id: 'user-1' },
21+
...overrides,
22+
} as Request;
23+
}
24+
25+
describe('generation retry admission', () => {
26+
afterEach(() => {
27+
jest.restoreAllMocks();
28+
});
29+
30+
it('marks only a submission with an existing durable claim as a retry', async () => {
31+
const hasClaim = jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(true);
32+
const req = generationRequest();
33+
const next = jest.fn() as NextFunction;
34+
35+
await detectGenerationRetry(req, {} as Response, next);
36+
37+
expect(hasClaim).toHaveBeenCalledWith('user-1', 'request-1');
38+
expect(isConfirmedGenerationRetry(req)).toBe(true);
39+
expect(next).toHaveBeenCalledTimes(1);
40+
});
41+
42+
it('leaves a new submission subject to the ordinary message limiters', async () => {
43+
jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(false);
44+
const req = generationRequest();
45+
46+
await detectGenerationRetry(req, {} as Response, jest.fn());
47+
48+
expect(isConfirmedGenerationRetry(req)).toBe(false);
49+
});
50+
51+
it.each([
52+
['a resume', { path: '/resume' }],
53+
['a resume with a trailing slash', { path: '/resume/' }],
54+
['a case-insensitive resume route', { path: '/Resume' }],
55+
['a request without an authenticated user', { user: undefined }],
56+
['an invalid idempotency key', { body: { clientRequestId: 'invalid key' } }],
57+
])('does not probe %s', async (_label, overrides) => {
58+
const hasClaim = jest.spyOn(GenerationJobManager, 'hasGenerationClaim');
59+
const req = generationRequest(overrides as Partial<Request>);
60+
const next = jest.fn() as NextFunction;
61+
62+
await detectGenerationRetry(req, {} as Response, next);
63+
64+
expect(hasClaim).not.toHaveBeenCalled();
65+
expect(isConfirmedGenerationRetry(req)).toBe(false);
66+
expect(next).toHaveBeenCalledTimes(1);
67+
});
68+
69+
it('fails closed to the ordinary limiters when the claim probe is unavailable', async () => {
70+
jest
71+
.spyOn(GenerationJobManager, 'hasGenerationClaim')
72+
.mockRejectedValue(new Error('store unavailable'));
73+
const warn = jest.spyOn(logger, 'warn').mockImplementation(() => logger);
74+
const req = generationRequest();
75+
const next = jest.fn() as NextFunction;
76+
77+
await detectGenerationRetry(req, {} as Response, next);
78+
79+
expect(isConfirmedGenerationRetry(req)).toBe(false);
80+
expect(warn).toHaveBeenCalledWith(
81+
'[GenerationIdempotency] Failed to inspect start-generation claim',
82+
expect.objectContaining({ userId: 'user-1', clientRequestId: 'request-1' }),
83+
);
84+
expect(next).toHaveBeenCalledTimes(1);
85+
});
86+
87+
it('bounds fresh claim probes before accessing the shared store', async () => {
88+
const hasClaim = jest
89+
.spyOn(GenerationJobManager, 'hasGenerationClaim')
90+
.mockResolvedValue(false);
91+
const downstream = jest.fn((_req, res) => res.sendStatus(204));
92+
const app = express();
93+
app.use(express.json());
94+
app.use((req, _res, next) => {
95+
req.user = { id: 'bounded-probe-user' };
96+
next();
97+
});
98+
app.use(generationRetryProbeLimiter);
99+
app.use(detectGenerationRetry);
100+
app.post('/', downstream);
101+
102+
for (let attempt = 0; attempt < GENERATION_RETRY_PROBE_MAX; attempt += 1) {
103+
await request(app)
104+
.post('/')
105+
.send({ clientRequestId: `probe-${attempt}` })
106+
.expect(204);
107+
}
108+
const rejected = await request(app)
109+
.post('/')
110+
.send({ clientRequestId: 'probe-rejected' })
111+
.expect(503);
112+
113+
expect(rejected.headers['retry-after']).toBeDefined();
114+
expect(rejected.body.code).toBe('SERVER_NOT_READY');
115+
expect(hasClaim).toHaveBeenCalledTimes(GENERATION_RETRY_PROBE_MAX);
116+
expect(downstream).toHaveBeenCalledTimes(GENERATION_RETRY_PROBE_MAX);
117+
});
118+
119+
it('makes a bounded confirmed retry delay participate in readiness recovery', async () => {
120+
jest.spyOn(GenerationJobManager, 'hasGenerationClaim').mockResolvedValue(true);
121+
const downstream = jest.fn((_req, res) => res.sendStatus(204));
122+
const app = express();
123+
app.use(express.json());
124+
app.use((req, _res, next) => {
125+
req.user = { id: 'bounded-retry-user' };
126+
next();
127+
});
128+
app.use(detectGenerationRetry);
129+
app.use(generationRetryLimiter);
130+
app.post('/', downstream);
131+
132+
for (let attempt = 0; attempt < GENERATION_RETRY_MAX; attempt += 1) {
133+
await request(app).post('/').send({ clientRequestId: 'bounded-request' }).expect(204);
134+
}
135+
const rejected = await request(app)
136+
.post('/')
137+
.send({ clientRequestId: 'bounded-request' })
138+
.expect(503);
139+
140+
expect(rejected.headers['retry-after']).toBeDefined();
141+
expect(rejected.body.code).toBe('SERVER_NOT_READY');
142+
expect(downstream).toHaveBeenCalledTimes(GENERATION_RETRY_MAX);
143+
});
144+
});

0 commit comments

Comments
 (0)