🔒 fix: Cover the pre-migration queue, first tick, and /usage/

Codex review on 892a27d.

- The heartbeat watched only the active conversation id, but `drainNext`
  merges in the `NEW_CONVO` queue, which outlives the URL update: items
  queued during the first turn stay keyed there until that run ends. It now
  renews the union of both, deduped since they are the same atom before
  migration.
- The interval installed without firing, so returning to a conversation
  whose hold was nearly up waited out a full period before the first
  renewal. It now renews immediately, then on each tick.
- Express's non-strict routing sends `POST /files/usage/` to the same
  handler with `req.path === '/usage/'`, so the exact comparison pushed it
  onto both upload limiters. A trailing-slash client would have spent its
  upload quota, and collected file-upload violations, on metadata
  heartbeats. Matching now tolerates the trailing slash.

Firing on effect start also made the drain-time renewal redundant: popping
an item changes the held set, so the renewal effect re-runs on its own. The
one case it cannot see is a refused send, where restoring the item leaves
the set identical, so that branch keeps an explicit renewal and the rest is
removed. Net one request per transition instead of two.
This commit is contained in:
Danny Avila 2026-07-27 23:22:55 -04:00
parent 892a27d176
commit ed9bc92fff
4 changed files with 92 additions and 35 deletions

View file

@ -33,6 +33,11 @@ const initialize = async () => {
const { fileUploadIpLimiter, fileUploadUserLimiter } = createFileLimiters();
const fileUsageLimiter = createFileUsageLimiter();
/** Non-strict routing means `/usage/` reaches the same handler, so match the
* route the way Express does. An exact comparison would push a
* trailing-slash request onto the upload quota instead. */
const isUsagePath = (path) => path.replace(/\/+$/, '') === '/usage';
/** Apply rate limiters to all POST routes (excluding /speech which is handled
* above). `/usage` is a metadata touch, so it gets its own limiter rather
* than consuming upload quota, but it is never left unmetered. */
@ -40,7 +45,7 @@ const initialize = async () => {
if (req.method !== 'POST' || req.path.startsWith('/speech')) {
return next();
}
if (req.path === '/usage') {
if (isUsagePath(req.path)) {
return fileUsageLimiter(req, res, next);
}
return fileUploadIpLimiter(req, res, (err) => {

View file

@ -97,6 +97,20 @@ describe('file route limiter wiring', () => {
expect(hits.uploadUser).toBe(0);
});
/* Non-strict routing sends `/usage/` to the same handler, so an exact path
* comparison would bill a trailing-slash client's heartbeats to the upload
* quota and hand them file-upload violations. */
it('routes a trailing-slash /usage/ through the usage limiter too', async () => {
const response = await request(app)
.post('/api/files/usage/')
.send({ file_ids: ['f1'] });
expect(response.status).toBe(200);
expect(hits.usage).toBe(1);
expect(hits.uploadIp).toBe(0);
expect(hits.uploadUser).toBe(0);
});
it('still applies the upload limiters to real uploads', async () => {
await request(app).post('/api/files').send({});

View file

@ -150,6 +150,7 @@ describe('useQueueDrain', () => {
]);
});
mockMarkFilesUsage.mockClear();
act(() => {
setters.setRunEnd!(runEnd());
});
@ -168,6 +169,7 @@ describe('useQueueDrain', () => {
]);
});
mockMarkFilesUsage.mockClear();
act(() => {
setters.setRunEnd!(runEnd());
});
@ -189,6 +191,7 @@ describe('useQueueDrain', () => {
]);
});
mockMarkFilesUsage.mockClear();
act(() => {
setters.setRunEnd!(runEnd());
});
@ -213,6 +216,7 @@ describe('useQueueDrain', () => {
});
ask.mockReturnValue(false);
mockMarkFilesUsage.mockClear();
act(() => {
setters.setRunEnd!(runEnd());
});
@ -232,16 +236,19 @@ describe('useQueueDrain', () => {
]);
});
expect(mockMarkFilesUsage).not.toHaveBeenCalled();
act(() => {
jest.advanceTimersByTime(30 * 60 * 1000);
});
/* Immediate, then on each interval. */
expect(mockMarkFilesUsage).toHaveBeenCalledTimes(1);
expect(mockMarkFilesUsage).toHaveBeenCalledWith({ file_ids: ['held'] });
act(() => {
jest.advanceTimersByTime(30 * 60 * 1000);
});
expect(mockMarkFilesUsage).toHaveBeenCalledTimes(2);
act(() => {
jest.advanceTimersByTime(30 * 60 * 1000);
});
expect(mockMarkFilesUsage).toHaveBeenCalledTimes(3);
} finally {
jest.useRealTimers();
}
@ -263,6 +270,37 @@ describe('useQueueDrain', () => {
}
});
/** Items queued during the first turn stay keyed under NEW_CONVO until that
* run ends, so renewing only the active id would skip them for its whole
* duration. */
it('renews the pre-migration NEW_CONVO queue alongside the active one', async () => {
setup(({ set }) => {
set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [
{ ...queuedMessage('n1', 'queued pre-migration'), files: [{ file_id: 'pending-migrate' }] },
]);
set(store.queuedMessagesByConvoId(CONVO_ID), [
{ ...queuedMessage('q1', 'queued after'), files: [{ file_id: 'already-migrated' }] },
]);
});
await waitFor(() => expect(mockMarkFilesUsage).toHaveBeenCalled());
const sent = mockMarkFilesUsage.mock.calls.flatMap((call) => call[0].file_ids);
expect(sent).toContain('pending-migrate');
expect(sent).toContain('already-migrated');
});
it('does not double-count the queue before migration', async () => {
setup(({ set }) => {
set(store.queuedMessagesByConvoId(Constants.NEW_CONVO), [
{ ...queuedMessage('n1', 'new convo'), files: [{ file_id: 'only-once' }] },
]);
}, Constants.NEW_CONVO as string);
await waitFor(() => expect(mockMarkFilesUsage).toHaveBeenCalled());
const sent = mockMarkFilesUsage.mock.calls.flatMap((call) => call[0].file_ids);
expect(sent).toEqual(['only-once']);
});
it('passes carried quotes + manual skills through as overrides', async () => {
const { ask, setters } = setup(({ set }) => {
set(store.queuedMessagesByConvoId(CONVO_ID), [

View file

@ -1,4 +1,4 @@
import { useEffect } from 'react';
import { useEffect, useMemo } from 'react';
import { Constants } from 'librechat-data-provider';
import { useRecoilValue, useRecoilCallback } from 'recoil';
import type { QueuedMessage } from '~/store/families';
@ -65,6 +65,18 @@ export default function useQueueDrain(
const ownQueue = useRecoilValue(
store.queuedMessagesByConvoId(activeConversationId ?? Constants.NEW_CONVO),
);
/** `drainNext` merges this in, and it outlives the URL update: items queued
* during the first turn stay keyed here until that run ends. Renewing only
* the active id would skip them for the whole of that run. */
const newConvoQueue = useRecoilValue(store.queuedMessagesByConvoId(Constants.NEW_CONVO));
/* Deduped because the two subscriptions are the same atom before migration.
* Keyed by id list so the effect re-runs when the held set changes, not
* whenever Recoil hands back a new array for the same contents. */
const renewKey = [
...new Set([...collectQueuedFileIds(ownQueue), ...collectQueuedFileIds(newConvoQueue)]),
].join(',');
const queuedFileIds = useMemo(() => (renewKey ? renewKey.split(',') : []), [renewKey]);
/**
* Heartbeat renewal while anything is queued.
@ -79,29 +91,27 @@ export default function useQueueDrain(
* normally.
*/
useEffect(() => {
const fileIds = collectQueuedFileIds(ownQueue);
if (fileIds.length === 0) {
if (queuedFileIds.length === 0) {
return;
}
const renew = () => {
for (const file_ids of batchFileIds(fileIds)) {
for (const file_ids of batchFileIds(queuedFileIds)) {
markFilesUsage({ file_ids });
}
};
/* Immediately, not one interval later: returning to a conversation whose
* hold is nearly up would otherwise wait out a full period first. */
renew();
const timer = setInterval(renew, QUEUE_USAGE_RENEW_INTERVAL_MS);
return () => clearInterval(timer);
}, [ownQueue, markFilesUsage]);
}, [queuedFileIds, markFilesUsage]);
// Fully synchronous reads (getLoadable): a useRecoilCallback snapshot is
// only guaranteed valid for the callback's synchronous execution, so no
// awaits may interleave with the reads.
const drainNext = useRecoilCallback(
({ snapshot, set }) =>
(): {
next: QueuedMessage;
conversationId: string;
remainderFileIds: string[];
} | null => {
(): { next: QueuedMessage; conversationId: string } | null => {
let end = snapshot.getLoadable(store.runEndByIndex(index)).getValue();
let fromParked = false;
if (
@ -180,9 +190,7 @@ export default function useQueueDrain(
if (remainder.length !== ownQueue.length || shouldMigrate || next != null) {
set(store.queuedMessagesByConvoId(conversationId), remainder);
}
return next
? { next, conversationId, remainderFileIds: collectQueuedFileIds(remainder) }
: null;
return next ? { next, conversationId } : null;
},
[index, activeConversationId],
);
@ -205,7 +213,7 @@ export default function useQueueDrain(
if (drained == null) {
return;
}
const { next, conversationId, remainderFileIds } = drained;
const { next, conversationId } = drained;
// The queued item is the FULL submission context: explicit (possibly
// empty) overrides stop `ask` from vacuuming up files, quotes, or skill
// picks the user has staged in the composer for their NEXT message.
@ -220,24 +228,16 @@ export default function useQueueDrain(
if (accepted === false) {
// `ask` refused without sending (e.g. the conversation history is not
// in the query cache yet, right after navigating back). Restore the
// item so the user's text is never silently dropped the chip stays
// item so the user's text is never silently dropped, the chip stays
// available for manual send.
restoreQueued(conversationId, next);
}
/** Renew the TTL hold on everything still queued. Items behind this one
* wait another full run (which may itself pause for approval), so a hold
* taken once at enqueue would lapse before a deep queue drains. Runs
* after `ask` so a refused send, whose item goes back on the queue with
* its run-end signal already consumed, is renewed too rather than left
* on its original hold. The server clamps every renewal against the
* upload time, so this cannot extend a file indefinitely.
* Fire-and-forget: send-time marking is the backstop. */
const toRenew =
accepted === false
? [...collectQueuedFileIds([next]), ...remainderFileIds]
: remainderFileIds;
for (const file_ids of batchFileIds(toRenew)) {
markFilesUsage({ file_ids });
/** Popping and restoring leaves the held set identical, so the renewal
* effect sees no change and will not re-run. Draining normally does
* change the set, and renews itself. Fire-and-forget: send-time
* marking is the backstop. */
for (const file_ids of batchFileIds(collectQueuedFileIds([next]))) {
markFilesUsage({ file_ids });
}
}
}, [
runEnd,