mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 06:52:47 +00:00
🔒 fix: Bound /files/usage TTL Hold Instead of Clearing It (#14470)
* 🔒 fix: Bound `/files/usage` TTL Hold Instead of Clearing It `POST /files/usage` marks queued attachments so the 1-hour upload-window TTL cannot reap them before the client queue drains. It did this by calling `updateFilesUsage`, which unsets `expiresAt` outright, turning every touched upload into a permanently retained file. The client queue is ephemeral browser state, so this also leaks in normal use: a closed tab or cleared queue leaves nothing referencing the files, but their TTL is already gone. The same mechanism let an authenticated user pin arbitrary owned uploads indefinitely, and the route was excluded from the file limiters, so the touch was entirely unmetered. Make the operation match its intent, a renewable hold rather than a release: - Add `extendFilesTTL`, which pushes `expiresAt` forward by a bounded window in a single owner-scoped `updateMany`. Two filter guards keep it safe under client-supplied ids: `$exists: true` so an already-released file never has a TTL re-added (that would schedule a live file for deletion), and `$lt` so a hold only ever moves the deadline later. The owner scope is a required argument, so an unscoped call is a no-op rather than a cross-user update. - `handleFilesUsageRequest` now holds for 24h instead of clearing, and no longer increments `usage`, since a queue touch is not a send. The real release still happens at drain, where `updateFilesUsage` marks the files used against an actual message. - Give `/usage` its own per-user limiter. Keeping it off the upload quota was intentional, leaving it unmetered was not. Abandoned queues are now reaped on schedule, and a replayed touch can only ever re-assert the same bounded window. * 🔒 fix: Anchor the `/files/usage` hold to upload time Codex review onb687922. The hold derived each new deadline from `Date.now()`, so a caller touching once a day advanced it by another 24h every time, far below the rate limit. That left indefinite preservation reachable and made the PR's replay claim wrong: the window was bounded per call but not in aggregate. Anchor the deadline to the file's immutable `createdAt` instead of the request clock. `extendFilesTTL` now takes a lifetime and sets `expiresAt = max(expiresAt, createdAt + holdMs)` in an aggregation pipeline, so the target is a fixed point per file and replay is inert rather than merely bounded. `$max` keeps the widen-only property and the `expiresAt: {$exists: true}` filter still refuses to resurrect a released TTL; `createdAt: {$exists: true}` fail-closes when the anchor is absent. The update runs with `timestamps: false`: a hold is TTL bookkeeping, not a content write, and bumping `updatedAt` also made every re-touch count as a modification, hiding whether the deadline actually moved. Also drop four `.node_modules-*` symlinks that `git add -A` swept in from an npm install. They pointed at absolute paths on one machine, so every other checkout got dangling entries. Added the pattern to .gitignore so a workspace install cannot reintroduce them. * 🔒 fix: Track the configured approval window in the `/files/usage` hold Codex review on9277620. `endpoints.agents.checkpointer.ttl` is a positive int with no upper bound, and its docs invite raising it for longer review windows. It drives the pending-action expiry, so a run can legitimately stay paused past 24h. The fixed 24h lifetime would then let Mongo reap an attachment while its approval was still live, and the later queue drain would send a file that no longer exists. Replace the fixed constant with `resolveFilesUsageHoldMs`, which adds the configured approval window to a 24h baseline covering upload, enqueue, and the run reaching its pause. The route reads the window from the same `getApprovalTtlMs(checkpointerCfg)` the pending action uses, so the two stay in lockstep. The replay bound is unaffected: the window is a per-deployment constant and the deadline is still `createdAt + holdMs`, so a replayed touch re-asserts the same instant and `$max` skips the write. Only an operator config change moves it, never a client. * 🔒 fix: Renew the `/files/usage` hold across queued runs, under a ceiling Codex review on2bd3c52. The drain sends one queued item per run completion, and each item starts a run that may itself pause for the full approval window. Since the hold was taken once at enqueue and pinned to the upload time, an item several places back could sit through multiple approval windows and lose its attachment while its chip and the live approval were still there. Another regression from this PR: the old `$unset` made retention permanent, so deep queues happened to work. The queue is unbounded, so no fixed lifetime covers it. Split the hold into a renewable window and a ceiling: expiresAt = max(expiresAt, min(now + renewMs, createdAt + maxLifetimeMs)) `renewMs` covers one run's wait and is granted from now, so a queue that is still draining re-asserts it at each transition; `useQueueDrain` now marks the remaining items' files whenever it pops one. `maxLifetimeMs` is measured from the immutable upload time and clamps every renewal, so repeated touches converge on a ceiling instead of advancing per call, which keeps the replay bound from the previous round intact. This also tightens abandonment: a queue nobody drains now lapses one `renewMs` after its last touch instead of surviving to the ceiling. `useQueueDrain`'s spec gained a QueryClientProvider, since the renewal goes through react-query. * 🔒 fix: Renew queued holds on a heartbeat, and stop dropping batches Codex review onf616bed. Three gaps in the renewal added last commit: - `collectQueuedFileIds` returned early at the server's 10-id cap, so a remainder holding more than one batch renewed only its first message and left the rest on their enqueue-time hold. Collect everything and split into capped requests instead of truncating. - A refused `ask()` restores the popped item, but renewal ran before the send and covered only the pre-existing remainder. Since the run-end signal is already consumed, nothing would touch that item again. Renewal now runs after `ask` and includes the restored item. - A single run can interrupt for approval more than once, each pause running to the configured window, so renewing only at drain transitions leaves a gap longer than `renewMs` with no renewal in it. The ceiling cannot help when nothing renews. The third is the same structural gap as the previous round along a new axis: renewal tied to discrete events loses the file whenever two events are further apart than the hold. Rather than hook each transition, renew on a 30 minute heartbeat while anything is queued, which is far below the smallest hold (24h) and so covers any single gap regardless of cause. Still bounded: every renewal is clamped against the file's upload time, so the ceiling is unchanged. A queue nobody has open emits no heartbeat and lapses one `renewMs` after its last touch, preserving the abandonment behaviour. * 🔒 fix: Cover the pre-migration queue, first tick, and `/usage/` Codex review on892a27d. - 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
GitHub
parent
ea643e8c9c
commit
728fc1276e
15 changed files with 975 additions and 71 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -52,6 +52,7 @@ __blobstorage__/**/*
|
|||
# Deployed apps should consider commenting these lines out:
|
||||
# see https://npmjs.org/doc/faq.html#Should-I-check-my-node_modules-folder-into-git
|
||||
node_modules/
|
||||
.node_modules-*
|
||||
meili_data/
|
||||
api/node_modules/
|
||||
client/node_modules/
|
||||
|
|
|
|||
|
|
@ -80,6 +80,39 @@ const createFileLimiters = () => {
|
|||
return { fileUploadIpLimiter, fileUploadUserLimiter };
|
||||
};
|
||||
|
||||
/**
|
||||
* Per-user limiter for the `/files/usage` TTL hold. Deliberately separate from
|
||||
* the upload limiters: a metadata touch must not consume upload quota, but it
|
||||
* still writes to the DB and so cannot go unmetered. Sized well above the
|
||||
* enqueue-driven call rate a real client produces.
|
||||
*/
|
||||
const createFileUsageLimiter = () => {
|
||||
const windowMinutes = parseInt(process.env.FILE_USAGE_USER_WINDOW) || 15;
|
||||
const max = parseInt(process.env.FILE_USAGE_USER_MAX) || 120;
|
||||
const windowMs = windowMinutes * 60 * 1000;
|
||||
|
||||
return rateLimit({
|
||||
windowMs,
|
||||
max,
|
||||
handler: async (req, res) => {
|
||||
const type = ViolationTypes.FILE_UPLOAD_LIMIT;
|
||||
await logViolation(
|
||||
req,
|
||||
res,
|
||||
type,
|
||||
{ type, max, limiter: 'user', windowInMinutes: windowMinutes },
|
||||
process.env.FILE_UPLOAD_VIOLATION_SCORE,
|
||||
);
|
||||
res.status(429).json({ message: 'Too many file usage requests. Try again later' });
|
||||
},
|
||||
keyGenerator: function (req) {
|
||||
return req.user?.id;
|
||||
},
|
||||
store: limiterCache('file_usage_user_limiter'),
|
||||
});
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
createFileLimiters,
|
||||
createFileUsageLimiter,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ const express = require('express');
|
|||
const { logger, SystemCapabilities } = require('@librechat/data-schemas');
|
||||
const {
|
||||
logAxiosError,
|
||||
getApprovalTtlMs,
|
||||
refreshS3FileUrls,
|
||||
handleFilesUsageRequest,
|
||||
shouldUseUploadSse,
|
||||
|
|
@ -144,15 +145,19 @@ router.get('/config', async (req, res) => {
|
|||
/**
|
||||
* POST /files/usage
|
||||
*
|
||||
* Owner-scoped TTL touch for uploads held in a client-side queue (mid-run
|
||||
* Owner-scoped TTL hold for uploads sitting in a client-side queue (mid-run
|
||||
* queued messages), so the upload-window TTL cannot reap them before drain.
|
||||
* Thin wrapper: validation, cap, and best-effort semantics live in
|
||||
* `@librechat/api` (`handleFilesUsageRequest`).
|
||||
* Extends the deadline rather than clearing it; the real release happens at
|
||||
* send. The approval window is passed through so a queue waiting on a paused
|
||||
* run outlives that pause. Thin wrapper: validation, cap, hold window, and
|
||||
* best-effort semantics live in `@librechat/api` (`handleFilesUsageRequest`).
|
||||
*/
|
||||
router.post('/usage', async (req, res) => {
|
||||
try {
|
||||
const checkpointerCfg = req.config?.endpoints?.[EModelEndpoint.agents]?.checkpointer;
|
||||
const { status, body } = await handleFilesUsageRequest(req.user ?? {}, req.body ?? {}, {
|
||||
updateFilesUsage: db.updateFilesUsage,
|
||||
extendFilesTTL: db.extendFilesTTL,
|
||||
approvalTtlMs: getApprovalTtlMs(checkpointerCfg),
|
||||
});
|
||||
return res.status(status).json(body);
|
||||
} catch (error) {
|
||||
|
|
|
|||
|
|
@ -938,7 +938,7 @@ describe('File Routes - Delete with Agent Access', () => {
|
|||
});
|
||||
|
||||
describe('POST /files/usage', () => {
|
||||
it('marks owned files used and clears the upload TTL', async () => {
|
||||
const createQueuedFile = async (expiresAt) => {
|
||||
const ownFileId = uuidv4();
|
||||
await createFile({
|
||||
user: otherUserId,
|
||||
|
|
@ -948,31 +948,110 @@ describe('File Routes - Delete with Agent Access', () => {
|
|||
bytes: 10,
|
||||
type: 'image/png',
|
||||
});
|
||||
await File.updateOne({ file_id: ownFileId }, { $set: { expiresAt: new Date() } });
|
||||
await File.updateOne({ file_id: ownFileId }, { $set: { expiresAt } });
|
||||
return ownFileId;
|
||||
};
|
||||
|
||||
it('extends the upload TTL of owned files without clearing it', async () => {
|
||||
const soon = new Date(Date.now() + 60 * 1000);
|
||||
const ownFileId = await createQueuedFile(soon);
|
||||
|
||||
const response = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [ownFileId] });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ marked: 1 });
|
||||
const marked = await File.findOne({ file_id: ownFileId }).lean();
|
||||
expect(marked.usage).toBe(1);
|
||||
expect(marked.expiresAt).toBeUndefined();
|
||||
expect(response.body).toEqual({ held: 1 });
|
||||
const held = await File.findOne({ file_id: ownFileId }).lean();
|
||||
/* The hold must remain a hold: still reapable, just later. */
|
||||
expect(held.expiresAt).toBeDefined();
|
||||
expect(held.expiresAt.getTime()).toBeGreaterThan(soon.getTime());
|
||||
/* The 24h baseline plus the default 24h approval window, so a queue
|
||||
* waiting on a paused run outlives that pause. Renewed from now, but
|
||||
* never past the ceiling measured from upload time. */
|
||||
const HOUR = 60 * 60 * 1000;
|
||||
expect(held.expiresAt.getTime()).toBeGreaterThan(Date.now() + 47 * HOUR);
|
||||
expect(held.expiresAt.getTime()).toBeLessThanOrEqual(
|
||||
held.createdAt.getTime() + 24 * HOUR + 8 * 24 * HOUR,
|
||||
);
|
||||
/* A queue touch is not a send, so it must not inflate usage. */
|
||||
expect(held.usage).toBe(0);
|
||||
});
|
||||
|
||||
it('cannot be replayed to preserve a file indefinitely', async () => {
|
||||
const ownFileId = await createQueuedFile(new Date(Date.now() + 60 * 1000));
|
||||
|
||||
const first = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [ownFileId] });
|
||||
expect(first.body).toEqual({ held: 1 });
|
||||
|
||||
for (let i = 0; i < 5; i++) {
|
||||
const repeat = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [ownFileId] });
|
||||
expect(repeat.status).toBe(200);
|
||||
}
|
||||
|
||||
/* Every renewal is clamped to the ceiling measured from upload time, so
|
||||
* replay converges there instead of advancing a window per call. */
|
||||
const HOUR = 60 * 60 * 1000;
|
||||
const held = await File.findOne({ file_id: ownFileId }).lean();
|
||||
expect(held.expiresAt).toBeDefined();
|
||||
expect(held.expiresAt.getTime()).toBeLessThanOrEqual(
|
||||
held.createdAt.getTime() + 24 * HOUR + 8 * 24 * HOUR,
|
||||
);
|
||||
});
|
||||
|
||||
it('never re-adds a TTL to a file that was already sent', async () => {
|
||||
const sentFileId = uuidv4();
|
||||
await createFile({
|
||||
user: otherUserId,
|
||||
file_id: sentFileId,
|
||||
filename: 'sent.png',
|
||||
filepath: '/uploads/sent.png',
|
||||
bytes: 10,
|
||||
type: 'image/png',
|
||||
});
|
||||
await File.updateOne({ file_id: sentFileId }, { $unset: { expiresAt: '' } });
|
||||
|
||||
const response = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [sentFileId] });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ held: 0 });
|
||||
const permanent = await File.findOne({ file_id: sentFileId }).lean();
|
||||
expect(permanent.expiresAt).toBeUndefined();
|
||||
});
|
||||
|
||||
it('never shortens an existing hold', async () => {
|
||||
const farOut = new Date(Date.now() + 90 * 24 * 60 * 60 * 1000);
|
||||
const ownFileId = await createQueuedFile(farOut);
|
||||
|
||||
const response = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [ownFileId] });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ held: 0 });
|
||||
const untouched = await File.findOne({ file_id: ownFileId }).lean();
|
||||
expect(untouched.expiresAt.getTime()).toBe(farOut.getTime());
|
||||
});
|
||||
|
||||
it("is owner-scoped: another user's file stays untouched (best-effort 200)", async () => {
|
||||
await File.updateOne({ file_id: fileId }, { $set: { expiresAt: new Date() } });
|
||||
const soon = new Date(Date.now() + 60 * 1000);
|
||||
await File.updateOne({ file_id: fileId }, { $set: { expiresAt: soon } });
|
||||
|
||||
const response = await request(app)
|
||||
.post('/files/usage')
|
||||
.send({ file_ids: [fileId] });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ marked: 0 });
|
||||
expect(response.body).toEqual({ held: 0 });
|
||||
const untouched = await File.findOne({ file_id: fileId }).lean();
|
||||
expect(untouched.usage).toBe(0);
|
||||
expect(untouched.expiresAt).toBeDefined();
|
||||
expect(untouched.expiresAt.getTime()).toBe(soon.getTime());
|
||||
});
|
||||
|
||||
it('rejects a list over the cap', async () => {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
const express = require('express');
|
||||
const {
|
||||
createFileUsageLimiter,
|
||||
createFileLimiters,
|
||||
configMiddleware,
|
||||
requireJwtAuth,
|
||||
|
|
@ -30,19 +31,29 @@ const initialize = async () => {
|
|||
router.use('/speech', speech);
|
||||
|
||||
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, and /usage — a metadata touch that must not consume upload quota) */
|
||||
* above). `/usage` is a metadata touch, so it gets its own limiter rather
|
||||
* than consuming upload quota, but it is never left unmetered. */
|
||||
router.use((req, res, next) => {
|
||||
if (req.method === 'POST' && !req.path.startsWith('/speech') && req.path !== '/usage') {
|
||||
return fileUploadIpLimiter(req, res, (err) => {
|
||||
if (err) {
|
||||
return next(err);
|
||||
}
|
||||
return fileUploadUserLimiter(req, res, next);
|
||||
});
|
||||
if (req.method !== 'POST' || req.path.startsWith('/speech')) {
|
||||
return next();
|
||||
}
|
||||
next();
|
||||
if (isUsagePath(req.path)) {
|
||||
return fileUsageLimiter(req, res, next);
|
||||
}
|
||||
return fileUploadIpLimiter(req, res, (err) => {
|
||||
if (err) {
|
||||
return next(err);
|
||||
}
|
||||
return fileUploadUserLimiter(req, res, next);
|
||||
});
|
||||
});
|
||||
|
||||
router.post('/', upload.single('file'), restoreTenantContextFromReq);
|
||||
|
|
|
|||
129
api/server/routes/files/index.limiters.test.js
Normal file
129
api/server/routes/files/index.limiters.test.js
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
const express = require('express');
|
||||
const request = require('supertest');
|
||||
|
||||
const hits = { uploadIp: 0, uploadUser: 0, usage: 0 };
|
||||
|
||||
jest.mock('~/server/middleware', () => ({
|
||||
createFileLimiters: jest.fn(() => ({
|
||||
fileUploadIpLimiter: (req, res, next) => {
|
||||
hits.uploadIp++;
|
||||
next();
|
||||
},
|
||||
fileUploadUserLimiter: (req, res, next) => {
|
||||
hits.uploadUser++;
|
||||
next();
|
||||
},
|
||||
})),
|
||||
createFileUsageLimiter: jest.fn(() => (req, res, next) => {
|
||||
hits.usage++;
|
||||
next();
|
||||
}),
|
||||
configMiddleware: (req, res, next) => {
|
||||
req.config = { fileStrategy: 'local', paths: { uploads: '/tmp', images: '/tmp' } };
|
||||
next();
|
||||
},
|
||||
requireJwtAuth: (req, res, next) => {
|
||||
req.user = { id: 'user-1', role: 'USER' };
|
||||
next();
|
||||
},
|
||||
uaParser: (req, res, next) => next(),
|
||||
checkBan: (req, res, next) => next(),
|
||||
}));
|
||||
|
||||
jest.mock('./multer', () => ({
|
||||
createMulterInstance: jest.fn(async () => ({
|
||||
single: jest.fn(() => (req, res, next) => next()),
|
||||
})),
|
||||
}));
|
||||
|
||||
const okRouter = (paths) => {
|
||||
const router = express.Router();
|
||||
for (const path of paths) {
|
||||
router.post(path, (req, res) => res.status(200).json({ ok: true }));
|
||||
}
|
||||
return router;
|
||||
};
|
||||
|
||||
jest.mock('./files', () => {
|
||||
const express = require('express');
|
||||
const router = express.Router();
|
||||
router.post('/', (req, res) => res.status(200).json({ ok: true }));
|
||||
router.post('/usage', (req, res) => res.status(200).json({ ok: true }));
|
||||
return router;
|
||||
});
|
||||
|
||||
jest.mock('./images', () => okRouter(['/']));
|
||||
jest.mock('./avatar', () => okRouter(['/']));
|
||||
jest.mock('./speech', () => okRouter(['/stt']));
|
||||
|
||||
jest.mock('~/server/routes/agents/v1', () => ({
|
||||
avatar: okRouter(['/:agent_id/avatar/']),
|
||||
}));
|
||||
jest.mock('~/server/routes/assistants/v1', () => ({
|
||||
avatar: okRouter(['/:assistant_id/avatar/']),
|
||||
}));
|
||||
|
||||
describe('file route limiter wiring', () => {
|
||||
let app;
|
||||
|
||||
beforeAll(async () => {
|
||||
const { initialize } = require('./index');
|
||||
app = express();
|
||||
app.use(express.json());
|
||||
app.use('/api/files', await initialize());
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
hits.uploadIp = 0;
|
||||
hits.uploadUser = 0;
|
||||
hits.usage = 0;
|
||||
});
|
||||
|
||||
it('meters POST /usage with the usage limiter, never leaving it unlimited', async () => {
|
||||
const response = await request(app)
|
||||
.post('/api/files/usage')
|
||||
.send({ file_ids: ['f1'] });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(hits.usage).toBe(1);
|
||||
});
|
||||
|
||||
it('keeps POST /usage off the upload quota', async () => {
|
||||
await request(app)
|
||||
.post('/api/files/usage')
|
||||
.send({ file_ids: ['f1'] });
|
||||
|
||||
expect(hits.uploadIp).toBe(0);
|
||||
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({});
|
||||
|
||||
expect(hits.uploadIp).toBe(1);
|
||||
expect(hits.uploadUser).toBe(1);
|
||||
expect(hits.usage).toBe(0);
|
||||
});
|
||||
|
||||
it('leaves /speech exempt from both limiters', async () => {
|
||||
await request(app).post('/api/files/speech/stt').send({});
|
||||
|
||||
expect(hits.uploadIp).toBe(0);
|
||||
expect(hits.uploadUser).toBe(0);
|
||||
expect(hits.usage).toBe(0);
|
||||
});
|
||||
});
|
||||
|
|
@ -21,6 +21,7 @@ jest.mock('~/server/middleware', () => ({
|
|||
fileUploadIpLimiter: (req, res, next) => next(),
|
||||
fileUploadUserLimiter: (req, res, next) => next(),
|
||||
})),
|
||||
createFileUsageLimiter: jest.fn(() => (req, res, next) => next()),
|
||||
configMiddleware: (req, res, next) => {
|
||||
req.config = {
|
||||
fileStrategy: 'local',
|
||||
|
|
|
|||
|
|
@ -1,11 +1,18 @@
|
|||
import React from 'react';
|
||||
import { Constants } from 'librechat-data-provider';
|
||||
import { act, renderHook, waitFor } from '@testing-library/react';
|
||||
import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
|
||||
import { RecoilRoot, useRecoilValue, useSetRecoilState, type MutableSnapshot } from 'recoil';
|
||||
import type { RunEnd, QueuedMessage } from '~/store/families';
|
||||
import useQueueDrain from '../useQueueDrain';
|
||||
import store from '~/store';
|
||||
|
||||
const mockMarkFilesUsage = jest.fn();
|
||||
jest.mock('~/data-provider', () => ({
|
||||
...jest.requireActual('~/data-provider'),
|
||||
useMarkFilesUsageMutation: () => ({ mutate: mockMarkFilesUsage }),
|
||||
}));
|
||||
|
||||
const INDEX = 0;
|
||||
const CONVO_ID = 'convo-drain';
|
||||
|
||||
|
|
@ -35,11 +42,18 @@ function setup(
|
|||
return null;
|
||||
}
|
||||
|
||||
/* The drain renews the TTL hold on what stays queued, which goes through
|
||||
* react-query, so the harness needs a client. */
|
||||
const queryClient = new QueryClient({
|
||||
defaultOptions: { mutations: { retry: false } },
|
||||
});
|
||||
const wrapper = ({ children }: { children: React.ReactNode }) => (
|
||||
<RecoilRoot initializeState={initialize}>
|
||||
<Harness />
|
||||
{children}
|
||||
</RecoilRoot>
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<RecoilRoot initializeState={initialize}>
|
||||
<Harness />
|
||||
{children}
|
||||
</RecoilRoot>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
renderHook(() => null, { wrapper });
|
||||
return { ask, setters };
|
||||
|
|
@ -57,6 +71,10 @@ const runEnd = (overrides: Partial<RunEnd> = {}): RunEnd => ({
|
|||
});
|
||||
|
||||
describe('useQueueDrain', () => {
|
||||
beforeEach(() => {
|
||||
mockMarkFilesUsage.mockClear();
|
||||
});
|
||||
|
||||
it('drains exactly one queued message on clean completion', async () => {
|
||||
const { ask, setters } = setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
|
|
@ -117,6 +135,172 @@ describe('useQueueDrain', () => {
|
|||
);
|
||||
});
|
||||
|
||||
/** Items behind the drained one wait another full run, which may itself
|
||||
* pause for approval, so a hold taken once at enqueue would lapse before a
|
||||
* deep queue finishes draining. */
|
||||
it('renews the TTL hold on attachments that stay queued', async () => {
|
||||
const { setters } = setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
{ ...queuedMessage('q1', 'first'), files: [{ file_id: 'sent', type: 'image/png' }] },
|
||||
{
|
||||
...queuedMessage('q2', 'second'),
|
||||
files: [{ file_id: 'still-queued', type: 'image/png' }],
|
||||
},
|
||||
{ ...queuedMessage('q3', 'third'), files: [{ file_id: 'also-queued', type: 'image/png' }] },
|
||||
]);
|
||||
});
|
||||
|
||||
mockMarkFilesUsage.mockClear();
|
||||
act(() => {
|
||||
setters.setRunEnd!(runEnd());
|
||||
});
|
||||
|
||||
await waitFor(() => expect(mockMarkFilesUsage).toHaveBeenCalledTimes(1));
|
||||
/* Only what remains: the drained item's file is released at send. */
|
||||
expect(mockMarkFilesUsage).toHaveBeenCalledWith({
|
||||
file_ids: ['still-queued', 'also-queued'],
|
||||
});
|
||||
});
|
||||
|
||||
it('does not renew when nothing with attachments stays queued', async () => {
|
||||
const { setters } = setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
{ ...queuedMessage('q1', 'only one'), files: [{ file_id: 'sent', type: 'image/png' }] },
|
||||
]);
|
||||
});
|
||||
|
||||
mockMarkFilesUsage.mockClear();
|
||||
act(() => {
|
||||
setters.setRunEnd!(runEnd());
|
||||
});
|
||||
|
||||
await waitFor(() => expect(mockMarkFilesUsage).not.toHaveBeenCalled());
|
||||
});
|
||||
|
||||
/** The server caps ids per request, so a queue holding more than one batch
|
||||
* must renew in several; truncating would leave later messages on their
|
||||
* enqueue-time hold. */
|
||||
it('renews every queued attachment across multiple capped batches', async () => {
|
||||
const manyFiles = (prefix: string, n: number) =>
|
||||
Array.from({ length: n }, (_, i) => ({ file_id: `${prefix}-${i}`, type: 'image/png' }));
|
||||
const { setters } = setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
{ ...queuedMessage('q1', 'first'), files: manyFiles('sent', 2) },
|
||||
{ ...queuedMessage('q2', 'second'), files: manyFiles('a', 10) },
|
||||
{ ...queuedMessage('q3', 'third'), files: manyFiles('b', 4) },
|
||||
]);
|
||||
});
|
||||
|
||||
mockMarkFilesUsage.mockClear();
|
||||
act(() => {
|
||||
setters.setRunEnd!(runEnd());
|
||||
});
|
||||
|
||||
await waitFor(() => expect(mockMarkFilesUsage).toHaveBeenCalledTimes(2));
|
||||
const sent = mockMarkFilesUsage.mock.calls.flatMap((call) => call[0].file_ids);
|
||||
expect(sent).toHaveLength(14);
|
||||
expect(sent).toContain('b-3');
|
||||
expect(sent).not.toContain('sent-0');
|
||||
for (const call of mockMarkFilesUsage.mock.calls) {
|
||||
expect(call[0].file_ids.length).toBeLessThanOrEqual(10);
|
||||
}
|
||||
});
|
||||
|
||||
/** A refused send puts the item back with its run-end signal already
|
||||
* consumed, so nothing else would touch it before the next drain. */
|
||||
it('renews a restored item when the send is refused', async () => {
|
||||
const { ask, setters } = setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
{ ...queuedMessage('q1', 'refused'), files: [{ file_id: 'restored', type: 'image/png' }] },
|
||||
]);
|
||||
});
|
||||
ask.mockReturnValue(false);
|
||||
|
||||
mockMarkFilesUsage.mockClear();
|
||||
act(() => {
|
||||
setters.setRunEnd!(runEnd());
|
||||
});
|
||||
|
||||
await waitFor(() => expect(mockMarkFilesUsage).toHaveBeenCalledTimes(1));
|
||||
expect(mockMarkFilesUsage).toHaveBeenCalledWith({ file_ids: ['restored'] });
|
||||
});
|
||||
|
||||
/** A single run can pause for approval more than once, so renewal cannot
|
||||
* depend on catching drain transitions alone. */
|
||||
it('renews on a heartbeat while items stay queued', async () => {
|
||||
jest.useFakeTimers();
|
||||
try {
|
||||
setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [
|
||||
{ ...queuedMessage('q1', 'waiting'), files: [{ file_id: 'held', type: 'image/png' }] },
|
||||
]);
|
||||
});
|
||||
|
||||
/* 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();
|
||||
}
|
||||
});
|
||||
|
||||
it('emits no heartbeat when the queue holds no attachments', async () => {
|
||||
jest.useFakeTimers();
|
||||
try {
|
||||
setup(({ set }) => {
|
||||
set(store.queuedMessagesByConvoId(CONVO_ID), [queuedMessage('q1', 'no files')]);
|
||||
});
|
||||
|
||||
act(() => {
|
||||
jest.advanceTimersByTime(2 * 60 * 60 * 1000);
|
||||
});
|
||||
expect(mockMarkFilesUsage).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
jest.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
/** 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), [
|
||||
|
|
@ -235,11 +419,16 @@ describe('useQueueDrain', () => {
|
|||
return null;
|
||||
}
|
||||
|
||||
const queryClient = new QueryClient({
|
||||
defaultOptions: { mutations: { retry: false } },
|
||||
});
|
||||
const wrapper = ({ children }: { children: React.ReactNode }) => (
|
||||
<RecoilRoot initializeState={initialize}>
|
||||
<Harness />
|
||||
{children}
|
||||
</RecoilRoot>
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<RecoilRoot initializeState={initialize}>
|
||||
<Harness />
|
||||
{children}
|
||||
</RecoilRoot>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
renderHook(() => null, { wrapper });
|
||||
return { ask, setters, state };
|
||||
|
|
|
|||
|
|
@ -1,10 +1,41 @@
|
|||
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';
|
||||
import type { TAskFunction } from '~/common';
|
||||
import { useMarkFilesUsageMutation } from '~/data-provider';
|
||||
import store from '~/store';
|
||||
|
||||
/** Mirrors the server's per-request cap on a usage touch. */
|
||||
const QUEUE_USAGE_MAX_FILES = 10;
|
||||
|
||||
/** Well under the server's smallest hold (24h), so no gap between renewals
|
||||
* can outlive one, however long a run pauses. */
|
||||
const QUEUE_USAGE_RENEW_INTERVAL_MS = 30 * 60 * 1000;
|
||||
|
||||
const collectQueuedFileIds = (items: QueuedMessage[]): string[] => {
|
||||
const fileIds: string[] = [];
|
||||
for (const item of items) {
|
||||
for (const file of item.files ?? []) {
|
||||
if (typeof file.file_id === 'string' && file.file_id.length > 0) {
|
||||
fileIds.push(file.file_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
return fileIds;
|
||||
};
|
||||
|
||||
/** The server caps ids per request, so a queue holding more than one batch
|
||||
* has to renew in several. Truncating instead would leave everything after
|
||||
* the first batch on its enqueue-time hold. */
|
||||
const batchFileIds = (fileIds: string[]): string[][] => {
|
||||
const batches: string[][] = [];
|
||||
for (let i = 0; i < fileIds.length; i += QUEUE_USAGE_MAX_FILES) {
|
||||
batches.push(fileIds.slice(i, i + QUEUE_USAGE_MAX_FILES));
|
||||
}
|
||||
return batches;
|
||||
};
|
||||
|
||||
/**
|
||||
* Auto-sends queued follow-up messages when a run finishes.
|
||||
*
|
||||
|
|
@ -30,6 +61,50 @@ export default function useQueueDrain(
|
|||
store.pendingRunEndByConvoId(activeConversationId ?? Constants.NEW_CONVO),
|
||||
);
|
||||
const isSubmitting = useRecoilValue(store.isSubmittingFamily(index));
|
||||
const { mutate: markFilesUsage } = useMarkFilesUsageMutation();
|
||||
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.
|
||||
*
|
||||
* Renewing only at drain transitions ties an attachment's survival to
|
||||
* catching every state change, and a single run can stretch well past one
|
||||
* hold: it may interrupt for approval more than once, and each pause can
|
||||
* run to the configured window. Rather than hook every transition, renew on
|
||||
* a cadence far shorter than the hold itself, so no single gap can outlive
|
||||
* it. Bounded regardless: the server clamps each renewal against the file's
|
||||
* upload time. A queue nobody has open stops emitting these and lapses
|
||||
* normally.
|
||||
*/
|
||||
useEffect(() => {
|
||||
if (queuedFileIds.length === 0) {
|
||||
return;
|
||||
}
|
||||
const renew = () => {
|
||||
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);
|
||||
}, [queuedFileIds, markFilesUsage]);
|
||||
|
||||
// Fully synchronous reads (getLoadable): a useRecoilCallback snapshot is
|
||||
// only guaranteed valid for the callback's synchronous execution, so no
|
||||
|
|
@ -153,9 +228,25 @@ 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);
|
||||
/** 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, parkedRunEnd, isSubmitting, activeConversationId, drainNext, restoreQueued, ask]);
|
||||
}, [
|
||||
runEnd,
|
||||
parkedRunEnd,
|
||||
isSubmitting,
|
||||
activeConversationId,
|
||||
drainNext,
|
||||
restoreQueued,
|
||||
markFilesUsage,
|
||||
ask,
|
||||
]);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -251,10 +251,11 @@ export default function useSteering({
|
|||
[],
|
||||
);
|
||||
|
||||
/** Fire-and-forget TTL touch for uploads entering the client queue: a
|
||||
/** Fire-and-forget TTL hold for uploads entering the client queue: a
|
||||
* queued message can outlive the upload window (long run, approval pause)
|
||||
* and send-time marking only happens at drain. Failure is tolerated —
|
||||
* the send-time marking remains the backstop. */
|
||||
* and send-time marking only happens at drain. The server extends the
|
||||
* deadline rather than clearing it, so a queue this tab never drains still
|
||||
* gets reaped. Failure is tolerated; send-time marking is the backstop. */
|
||||
const markQueuedFilesUsage = useCallback(
|
||||
(files?: TMessage['files']) => {
|
||||
if (files == null || files.length === 0) {
|
||||
|
|
@ -285,8 +286,8 @@ export default function useSteering({
|
|||
files?: TMessage['files'];
|
||||
quotes?: string[];
|
||||
manualSkills?: string[];
|
||||
/** Set when the files were ALREADY queued/steered — their usage was
|
||||
* marked when they first entered the queue (or at the steer 202). */
|
||||
/** Set when the files were ALREADY queued/steered: their TTL was
|
||||
* held when they first entered the queue (or at the steer 202). */
|
||||
skipUsageMark?: boolean;
|
||||
},
|
||||
) => {
|
||||
|
|
@ -694,7 +695,7 @@ export default function useSteering({
|
|||
}
|
||||
return;
|
||||
}
|
||||
// Files were already marked used when the item first entered the queue.
|
||||
// Files were already TTL-held when the item first entered the queue.
|
||||
enqueue(taken.text, {
|
||||
front: true,
|
||||
files: taken.files,
|
||||
|
|
|
|||
|
|
@ -1,17 +1,23 @@
|
|||
import { handleFilesUsageRequest, FILES_USAGE_MAX_IDS } from './usage';
|
||||
import {
|
||||
handleFilesUsageRequest,
|
||||
resolveFilesUsageHold,
|
||||
FILES_USAGE_QUEUED_RUN_ALLOWANCE,
|
||||
FILES_USAGE_BASE_HOLD_MS,
|
||||
FILES_USAGE_MAX_IDS,
|
||||
} from './usage';
|
||||
|
||||
describe('handleFilesUsageRequest', () => {
|
||||
const user = { id: 'user-1', tenantId: 'tenant-1' };
|
||||
|
||||
const createDeps = (marked: unknown[] = []) => ({
|
||||
updateFilesUsage: jest.fn().mockResolvedValue(marked),
|
||||
const createDeps = (held = 0) => ({
|
||||
extendFilesTTL: jest.fn().mockResolvedValue(held),
|
||||
});
|
||||
|
||||
it('rejects unauthenticated requests without touching the DB', async () => {
|
||||
const deps = createDeps();
|
||||
const result = await handleFilesUsageRequest({}, { file_ids: ['f1'] }, deps);
|
||||
expect(result).toEqual({ status: 401, body: { code: 'UNAUTHORIZED' } });
|
||||
expect(deps.updateFilesUsage).not.toHaveBeenCalled();
|
||||
expect(deps.extendFilesTTL).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it.each([
|
||||
|
|
@ -25,7 +31,7 @@ describe('handleFilesUsageRequest', () => {
|
|||
const result = await handleFilesUsageRequest(user, body, deps);
|
||||
expect(result.status).toBe(400);
|
||||
expect(result.body).toEqual({ code: 'INVALID_FILE_IDS' });
|
||||
expect(deps.updateFilesUsage).not.toHaveBeenCalled();
|
||||
expect(deps.extendFilesTTL).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('caps the list at FILES_USAGE_MAX_IDS', async () => {
|
||||
|
|
@ -34,24 +40,80 @@ describe('handleFilesUsageRequest', () => {
|
|||
const result = await handleFilesUsageRequest(user, { file_ids }, deps);
|
||||
expect(result.status).toBe(400);
|
||||
expect(result.body).toEqual({ code: 'TOO_MANY_FILES', max: FILES_USAGE_MAX_IDS });
|
||||
expect(deps.updateFilesUsage).not.toHaveBeenCalled();
|
||||
expect(deps.extendFilesTTL).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('marks usage owner-scoped and returns best-effort 200', async () => {
|
||||
const deps = createDeps([{ file_id: 'f1' }]);
|
||||
it('requests an owner-scoped hold of the resolved window', async () => {
|
||||
const deps = createDeps(1);
|
||||
const result = await handleFilesUsageRequest(user, { file_ids: ['f1', 'f2'] }, deps);
|
||||
expect(deps.updateFilesUsage).toHaveBeenCalledTimes(1);
|
||||
expect(deps.updateFilesUsage).toHaveBeenCalledWith(
|
||||
[{ file_id: 'f1' }, { file_id: 'f2' }],
|
||||
undefined,
|
||||
expect(deps.extendFilesTTL).toHaveBeenCalledTimes(1);
|
||||
expect(deps.extendFilesTTL).toHaveBeenCalledWith(
|
||||
['f1', 'f2'],
|
||||
{ renewMs: FILES_USAGE_BASE_HOLD_MS, maxLifetimeMs: FILES_USAGE_BASE_HOLD_MS },
|
||||
{ user: 'user-1', tenantId: 'tenant-1' },
|
||||
);
|
||||
expect(result).toEqual({ status: 200, body: { marked: 1 } });
|
||||
expect(result).toEqual({ status: 200, body: { held: 1 } });
|
||||
});
|
||||
|
||||
it('returns 200 with zero marked when no id resolves to an owned file', async () => {
|
||||
const deps = createDeps([]);
|
||||
/** The ceiling must be a constant the data layer clamps against the upload
|
||||
* time, not a deadline derived here: a request-clock deadline with no
|
||||
* ceiling would let a caller walk the file's lifetime forward per call. */
|
||||
it('passes a constant ceiling, never a request-derived deadline', async () => {
|
||||
const deps = createDeps(1);
|
||||
await handleFilesUsageRequest(user, { file_ids: ['f1'] }, deps);
|
||||
await handleFilesUsageRequest(user, { file_ids: ['f1'] }, deps);
|
||||
|
||||
const [, firstHold] = deps.extendFilesTTL.mock.calls[0];
|
||||
const [, secondHold] = deps.extendFilesTTL.mock.calls[1];
|
||||
expect(secondHold).toEqual(firstHold);
|
||||
expect(typeof firstHold.maxLifetimeMs).toBe('number');
|
||||
});
|
||||
|
||||
/** A paused run's approval window has no configured upper bound, so a fixed
|
||||
* window shorter than it would reap an attachment whose approval is still
|
||||
* live and leave the drain sending a missing file. */
|
||||
it('stretches the hold to outlast a long configured approval window', async () => {
|
||||
const sevenDaysMs = 7 * 24 * 60 * 60 * 1000;
|
||||
const deps = { ...createDeps(1), approvalTtlMs: sevenDaysMs };
|
||||
await handleFilesUsageRequest(user, { file_ids: ['f1'] }, deps);
|
||||
|
||||
const [, hold] = deps.extendFilesTTL.mock.calls[0];
|
||||
expect(hold.renewMs).toBe(FILES_USAGE_BASE_HOLD_MS + sevenDaysMs);
|
||||
expect(hold.renewMs).toBeGreaterThan(sevenDaysMs);
|
||||
});
|
||||
|
||||
describe('resolveFilesUsageHold', () => {
|
||||
it('covers one approval window per renewal', () => {
|
||||
expect(resolveFilesUsageHold(1000).renewMs).toBe(FILES_USAGE_BASE_HOLD_MS + 1000);
|
||||
});
|
||||
|
||||
/** The drain sends one queued item per run completion, so a deep queue
|
||||
* waits behind several approval windows; the ceiling has to span them. */
|
||||
it('scales the ceiling to the queued run chain', () => {
|
||||
const hold = resolveFilesUsageHold(1000);
|
||||
expect(hold.maxLifetimeMs).toBe(
|
||||
FILES_USAGE_BASE_HOLD_MS + 1000 * FILES_USAGE_QUEUED_RUN_ALLOWANCE,
|
||||
);
|
||||
expect(hold.maxLifetimeMs).toBeGreaterThan(hold.renewMs);
|
||||
});
|
||||
|
||||
it.each([
|
||||
['undefined', undefined],
|
||||
['zero', 0],
|
||||
['negative', -1],
|
||||
['NaN', Number.NaN],
|
||||
['Infinity', Number.POSITIVE_INFINITY],
|
||||
])('falls back to the baseline for %s', (_label, value) => {
|
||||
expect(resolveFilesUsageHold(value)).toEqual({
|
||||
renewMs: FILES_USAGE_BASE_HOLD_MS,
|
||||
maxLifetimeMs: FILES_USAGE_BASE_HOLD_MS,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
it('returns 200 with zero held when no id resolves to an owned file', async () => {
|
||||
const deps = createDeps(0);
|
||||
const result = await handleFilesUsageRequest(user, { file_ids: ['not-owned'] }, deps);
|
||||
expect(result).toEqual({ status: 200, body: { marked: 0 } });
|
||||
expect(result).toEqual({ status: 200, body: { held: 0 } });
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,6 +1,53 @@
|
|||
/** Cap per usage touch, mirroring the composer's practical attachment limit. */
|
||||
export const FILES_USAGE_MAX_IDS: number = 10;
|
||||
|
||||
/**
|
||||
* Baseline lifetime a held upload gets, measured from its upload time. Covers
|
||||
* the span a queued attachment spends before a run can even pause: upload,
|
||||
* enqueue, and the run itself reaching an approval point.
|
||||
*/
|
||||
export const FILES_USAGE_BASE_HOLD_MS: number = 24 * 60 * 60 * 1000;
|
||||
|
||||
/**
|
||||
* Queued runs a held attachment is assumed to wait behind. The queue is
|
||||
* unbounded, so no multiple is provably sufficient; this is the depth the
|
||||
* ceiling covers before a still-queued attachment can lapse.
|
||||
*/
|
||||
export const FILES_USAGE_QUEUED_RUN_ALLOWANCE: number = 8;
|
||||
|
||||
export interface FilesUsageHold {
|
||||
/** Granted from now, so an actively draining queue keeps renewing. */
|
||||
renewMs: number;
|
||||
/** Hard ceiling from upload time, so renewals converge instead of walking. */
|
||||
maxLifetimeMs: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the hold window for one touch.
|
||||
*
|
||||
* A queued message legitimately outlives the upload window when a run pauses
|
||||
* for approval, and that pause is bounded by the deployment's configured
|
||||
* approval window (`endpoints.agents.checkpointer.ttl`), which has no upper
|
||||
* limit. Deeper queues wait behind several such runs, since the drain sends
|
||||
* one item per run completion.
|
||||
*
|
||||
* Hence two numbers rather than one. `renewMs` covers a single run's wait and
|
||||
* is granted from now, so a queue that is still draining re-asserts it at each
|
||||
* transition. `maxLifetimeMs` is measured from the immutable upload time and
|
||||
* caps every renewal, so repeated touches converge on a fixed ceiling instead
|
||||
* of advancing a window per call. An abandoned queue therefore lapses one
|
||||
* `renewMs` after its last touch rather than surviving to the ceiling.
|
||||
*
|
||||
* @param approvalTtlMs - Configured approval window; falsy means none configured
|
||||
*/
|
||||
export function resolveFilesUsageHold(approvalTtlMs?: number): FilesUsageHold {
|
||||
const approval = Number.isFinite(approvalTtlMs) && approvalTtlMs! > 0 ? approvalTtlMs! : 0;
|
||||
return {
|
||||
renewMs: FILES_USAGE_BASE_HOLD_MS + approval,
|
||||
maxLifetimeMs: FILES_USAGE_BASE_HOLD_MS + approval * FILES_USAGE_QUEUED_RUN_ALLOWANCE,
|
||||
};
|
||||
}
|
||||
|
||||
export interface FilesUsageUser {
|
||||
id?: string;
|
||||
tenantId?: string;
|
||||
|
|
@ -17,20 +64,35 @@ export interface FilesUsageResult {
|
|||
}
|
||||
|
||||
export interface FilesUsageDeps {
|
||||
/** Owner-scoped usage marker (`db.updateFilesUsage`-shaped). */
|
||||
updateFilesUsage: (
|
||||
files: Array<{ file_id: string }>,
|
||||
fileIds?: string[],
|
||||
options?: { user?: string; tenantId?: string | null },
|
||||
) => Promise<unknown[]>;
|
||||
/** Owner-scoped TTL hold (`db.extendFilesTTL`-shaped). */
|
||||
extendFilesTTL: (
|
||||
fileIds: string[],
|
||||
hold: FilesUsageHold,
|
||||
owner: { user: string; tenantId?: string | null },
|
||||
) => Promise<number>;
|
||||
/** Configured approval window, so the hold outlasts a paused run. */
|
||||
approvalTtlMs?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Owner-scoped usage touch for attachments entering a client-side queue: a
|
||||
* queued message can outlive the upload window (long run, approval pause), so
|
||||
* marking at queue time stops the TTL from reaping files the drain will send.
|
||||
* Best-effort 200 — ids that do not resolve to an owned file are not errors
|
||||
* (send-time marking remains the backstop).
|
||||
* Owner-scoped TTL hold for attachments entering a client-side queue: a
|
||||
* queued message can outlive the upload window (long run, approval pause),
|
||||
* so holding at queue time stops the TTL from reaping files the drain will
|
||||
* send.
|
||||
*
|
||||
* A hold, not a release. The client queue is ephemeral browser state, so a
|
||||
* closed tab or a cleared queue leaves nothing referencing these files, and
|
||||
* clearing the TTL outright would strand them in storage permanently. The
|
||||
* hold only ever widens, and every renewal is capped against the file's
|
||||
* upload time, so replaying it converges on a ceiling instead of advancing
|
||||
* indefinitely; the real release happens at send, where `updateFilesUsage`
|
||||
* marks the files used against an actual message.
|
||||
*
|
||||
* The window tracks the deployment's configured approval TTL so a queue
|
||||
* waiting on a paused run cannot be reaped while that approval is live.
|
||||
*
|
||||
* Best-effort 200: ids that do not resolve to a held file are not errors
|
||||
* (they may be already-sent files, or not owned).
|
||||
*/
|
||||
export async function handleFilesUsageRequest(
|
||||
user: FilesUsageUser,
|
||||
|
|
@ -47,16 +109,16 @@ export async function handleFilesUsageRequest(
|
|||
if (raw.length > FILES_USAGE_MAX_IDS) {
|
||||
return { status: 400, body: { code: 'TOO_MANY_FILES', max: FILES_USAGE_MAX_IDS } };
|
||||
}
|
||||
const files: Array<{ file_id: string }> = [];
|
||||
const fileIds: string[] = [];
|
||||
for (const value of raw) {
|
||||
if (typeof value !== 'string' || value.length === 0) {
|
||||
return { status: 400, body: { code: 'INVALID_FILE_IDS' } };
|
||||
}
|
||||
files.push({ file_id: value });
|
||||
fileIds.push(value);
|
||||
}
|
||||
const marked = await deps.updateFilesUsage(files, undefined, {
|
||||
const held = await deps.extendFilesTTL(fileIds, resolveFilesUsageHold(deps.approvalTtlMs), {
|
||||
user: user.id,
|
||||
tenantId: user.tenantId,
|
||||
});
|
||||
return { status: 200, body: { marked: marked.length } };
|
||||
return { status: 200, body: { held } };
|
||||
}
|
||||
|
|
|
|||
|
|
@ -242,7 +242,8 @@ export type TFilesUsageBody = {
|
|||
};
|
||||
|
||||
export type TFilesUsageResponse = {
|
||||
marked: number;
|
||||
/** Count of queued uploads whose TTL hold was extended. */
|
||||
held: number;
|
||||
};
|
||||
|
||||
export type DeleteFilesResponse = {
|
||||
|
|
|
|||
|
|
@ -1123,6 +1123,167 @@ describe('File Methods', () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe('extendFilesTTL', () => {
|
||||
const HOUR = 3_600_000;
|
||||
const HOLD = { renewMs: 24 * HOUR, maxLifetimeMs: 48 * HOUR };
|
||||
|
||||
const seedTempFile = async (
|
||||
userId: mongoose.Types.ObjectId,
|
||||
expiresAt: Date,
|
||||
createdAt?: Date,
|
||||
) => {
|
||||
const fileId = uuidv4();
|
||||
await fileMethods.createFile({
|
||||
file_id: fileId,
|
||||
user: userId,
|
||||
filename: `${fileId}.txt`,
|
||||
filepath: `/uploads/${fileId}.txt`,
|
||||
type: 'text/plain',
|
||||
bytes: 100,
|
||||
});
|
||||
await mongoose.models.File.updateOne(
|
||||
{ file_id: fileId },
|
||||
{ $set: { expiresAt, ...(createdAt ? { createdAt } : {}) } },
|
||||
{ timestamps: false },
|
||||
);
|
||||
return fileId;
|
||||
};
|
||||
|
||||
const readFile = async (fileId: string) =>
|
||||
(await mongoose.models.File.findOne({ file_id: fileId })
|
||||
.lean<{ createdAt: Date; expiresAt?: Date }>()
|
||||
.exec())!;
|
||||
|
||||
it('widens the TTL toward the renewal window without unsetting it', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const fileId = await seedTempFile(userId, new Date(Date.now() + 60_000));
|
||||
|
||||
const count = await fileMethods.extendFilesTTL([fileId], HOLD, { user: String(userId) });
|
||||
|
||||
expect(count).toBe(1);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt).toBeDefined();
|
||||
expect(file.expiresAt!.getTime()).toBeGreaterThan(Date.now() + 23 * HOUR);
|
||||
expect(file.expiresAt!.getTime()).toBeLessThanOrEqual(
|
||||
file.createdAt.getTime() + HOLD.maxLifetimeMs,
|
||||
);
|
||||
});
|
||||
|
||||
/** The bound that makes the endpoint safe to expose: every renewal is
|
||||
* clamped to createdAt + maxLifetimeMs, so replaying the call converges
|
||||
* on a ceiling instead of advancing a window at a time. */
|
||||
it('clamps every renewal to the ceiling, however often it is replayed', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const fileId = await seedTempFile(userId, new Date(Date.now() + 60_000));
|
||||
const ceilingHold = { renewMs: 24 * HOUR, maxLifetimeMs: 10 * 60_000 };
|
||||
const { createdAt } = await readFile(fileId);
|
||||
|
||||
const first = await fileMethods.extendFilesTTL([fileId], ceilingHold, {
|
||||
user: String(userId),
|
||||
});
|
||||
expect(first).toBe(1);
|
||||
|
||||
for (let i = 0; i < 5; i++) {
|
||||
expect(
|
||||
await fileMethods.extendFilesTTL([fileId], ceilingHold, { user: String(userId) }),
|
||||
).toBe(0);
|
||||
}
|
||||
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt!.getTime()).toBe(createdAt.getTime() + ceilingHold.maxLifetimeMs);
|
||||
/* The renewal window alone would have granted a full day. */
|
||||
expect(file.expiresAt!.getTime()).toBeLessThan(Date.now() + HOUR);
|
||||
});
|
||||
|
||||
/** Renewal is measured from now, so a queue still draining across
|
||||
* successive runs keeps its attachments past the first hold. */
|
||||
it('renews from now while the ceiling still allows it', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const twoHoursAgo = new Date(Date.now() - 2 * HOUR);
|
||||
const fileId = await seedTempFile(userId, new Date(Date.now() + 60_000), twoHoursAgo);
|
||||
|
||||
const count = await fileMethods.extendFilesTTL(
|
||||
[fileId],
|
||||
{ renewMs: HOUR, maxLifetimeMs: 24 * HOUR },
|
||||
{ user: String(userId) },
|
||||
);
|
||||
|
||||
expect(count).toBe(1);
|
||||
const file = await readFile(fileId);
|
||||
/* Anchoring to createdAt alone would have expired this an hour ago. */
|
||||
expect(file.expiresAt!.getTime()).toBeGreaterThan(Date.now() + 50 * 60_000);
|
||||
});
|
||||
|
||||
it('does not resurrect a TTL on an already-released file', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const fileId = await seedTempFile(userId, new Date(Date.now() + 60_000));
|
||||
await fileMethods.updateFileUsage({ file_id: fileId, user: String(userId) });
|
||||
|
||||
const count = await fileMethods.extendFilesTTL([fileId], HOLD, { user: String(userId) });
|
||||
|
||||
expect(count).toBe(0);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt).toBeUndefined();
|
||||
});
|
||||
|
||||
it('never moves an expiry earlier', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const farOut = new Date(Date.now() + 7 * 24 * HOUR);
|
||||
const fileId = await seedTempFile(userId, farOut);
|
||||
|
||||
const count = await fileMethods.extendFilesTTL([fileId], HOLD, { user: String(userId) });
|
||||
|
||||
expect(count).toBe(0);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt!.getTime()).toBe(farOut.getTime());
|
||||
});
|
||||
|
||||
it("leaves another user's file untouched", async () => {
|
||||
const ownerId = new mongoose.Types.ObjectId();
|
||||
const attackerId = new mongoose.Types.ObjectId();
|
||||
const soon = new Date(Date.now() + 60_000);
|
||||
const fileId = await seedTempFile(ownerId, soon);
|
||||
|
||||
const count = await fileMethods.extendFilesTTL([fileId], HOLD, {
|
||||
user: String(attackerId),
|
||||
});
|
||||
|
||||
expect(count).toBe(0);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt!.getTime()).toBe(soon.getTime());
|
||||
});
|
||||
|
||||
it('is a no-op without an owner scope rather than a cross-user update', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const soon = new Date(Date.now() + 60_000);
|
||||
const fileId = await seedTempFile(userId, soon);
|
||||
|
||||
const count = await fileMethods.extendFilesTTL([fileId], HOLD, { user: '' } as {
|
||||
user: string;
|
||||
});
|
||||
|
||||
expect(count).toBe(0);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt!.getTime()).toBe(soon.getTime());
|
||||
});
|
||||
|
||||
it('is a no-op for a non-positive hold', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const soon = new Date(Date.now() + 60_000);
|
||||
const fileId = await seedTempFile(userId, soon);
|
||||
const owner = { user: String(userId) };
|
||||
|
||||
expect(
|
||||
await fileMethods.extendFilesTTL([fileId], { renewMs: 0, maxLifetimeMs: 0 }, owner),
|
||||
).toBe(0);
|
||||
expect(
|
||||
await fileMethods.extendFilesTTL([fileId], { renewMs: -1, maxLifetimeMs: -1 }, owner),
|
||||
).toBe(0);
|
||||
const file = await readFile(fileId);
|
||||
expect(file.expiresAt!.getTime()).toBe(soon.getTime());
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteFile', () => {
|
||||
it('should delete a file by file_id', async () => {
|
||||
const fileId = uuidv4();
|
||||
|
|
|
|||
|
|
@ -82,6 +82,11 @@ export function createFileMethods(mongoose: typeof import('mongoose')): {
|
|||
fileIds?: string[],
|
||||
options?: { user?: string; tenantId?: string | null },
|
||||
) => Promise<IMongoFile[]>;
|
||||
extendFilesTTL: (
|
||||
fileIds: string[],
|
||||
hold: { renewMs: number; maxLifetimeMs: number },
|
||||
owner: { user: string; tenantId?: string | null },
|
||||
) => Promise<number>;
|
||||
sweepOrphanedPreviews: (maxAgeMs?: number) => Promise<number>;
|
||||
} {
|
||||
/**
|
||||
|
|
@ -546,6 +551,78 @@ export function createFileMethods(mongoose: typeof import('mongoose')): {
|
|||
return results.filter((result): result is IMongoFile => result != null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Widens the upload-window TTL of owned, still-temporary files to
|
||||
* `min(now + renewMs, createdAt + maxLifetimeMs)`.
|
||||
*
|
||||
* A renewable hold, not a release: unlike `updateFileUsage` this never
|
||||
* unsets `expiresAt`, so a file that is held but never actually sent is
|
||||
* still reaped once the hold lapses. Four properties hold by
|
||||
* construction, which is what makes the write safe to drive from a
|
||||
* client-supplied id list:
|
||||
* - `$min` against `createdAt + maxLifetimeMs` caps every renewal against
|
||||
* an immutable anchor, so repeated calls converge on a fixed ceiling
|
||||
* instead of walking a file's lifetime forward a window at a time;
|
||||
* - renewing from `now` up to that ceiling lets a queue that is still
|
||||
* draining keep its attachments alive across successive runs, while an
|
||||
* abandoned queue lapses a single `renewMs` after its last touch rather
|
||||
* than surviving to the ceiling;
|
||||
* - `$max` against the current value means a hold only ever widens;
|
||||
* - `expiresAt: { $exists: true }` means a file whose TTL was already
|
||||
* cleared by a real send stays permanent. Re-adding `expiresAt` there
|
||||
* would schedule a live file for deletion.
|
||||
*
|
||||
* `createdAt` is required rather than defaulted: without the anchor there
|
||||
* is no ceiling to enforce, so such a file is skipped instead of held.
|
||||
*
|
||||
* The owner scope is required, not optional: an unscoped call would hold
|
||||
* every user's matching file. A missing owner is a no-op, not a wide
|
||||
* update.
|
||||
*
|
||||
* @param fileIds - File IDs to hold
|
||||
* @param hold - `renewMs` granted from now, capped at `maxLifetimeMs` from upload
|
||||
* @param owner - Owner scope; mismatches leave the TTL unchanged
|
||||
* @returns Number of files whose hold was widened
|
||||
*/
|
||||
async function extendFilesTTL(
|
||||
fileIds: string[],
|
||||
hold: { renewMs: number; maxLifetimeMs: number },
|
||||
owner: { user: string; tenantId?: string | null },
|
||||
): Promise<number> {
|
||||
const renewMs = hold?.renewMs;
|
||||
const maxLifetimeMs = hold?.maxLifetimeMs;
|
||||
if (fileIds.length === 0 || !owner?.user || !(renewMs > 0) || !(maxLifetimeMs > 0)) {
|
||||
return 0;
|
||||
}
|
||||
const File = mongoose.models.File as Model<IMongoFile>;
|
||||
const filter = withOwnerScope(
|
||||
{
|
||||
file_id: { $in: [...new Set(fileIds)] },
|
||||
expiresAt: { $exists: true },
|
||||
createdAt: { $exists: true },
|
||||
},
|
||||
{ userId: owner.user, tenantId: owner.tenantId },
|
||||
);
|
||||
const renewUntil = new Date(Date.now() + renewMs);
|
||||
const result = await File.updateMany(
|
||||
filter,
|
||||
[
|
||||
{
|
||||
$set: {
|
||||
expiresAt: {
|
||||
$max: ['$expiresAt', { $min: [renewUntil, { $add: ['$createdAt', maxLifetimeMs] }] }],
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
/** `timestamps: false`: a hold is TTL bookkeeping, not a content write.
|
||||
* Bumping `updatedAt` would also make every re-touch count as a
|
||||
* modification, hiding whether the deadline actually moved. */
|
||||
{ timestamps: false },
|
||||
);
|
||||
return result.modifiedCount ?? 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark stale `status: 'pending'` file records as `'failed'` with
|
||||
* `previewError: 'orphaned'`. Recovers from the one case the
|
||||
|
|
@ -594,6 +671,7 @@ export function createFileMethods(mongoose: typeof import('mongoose')): {
|
|||
deleteFileByFilter,
|
||||
batchUpdateFiles,
|
||||
updateFilesUsage,
|
||||
extendFilesTTL,
|
||||
sweepOrphanedPreviews,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue