mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-02-27 12:54:09 +01:00
📌 fix: Populate userMessage.files Before First DB Save (#11939)
* fix: populate userMessage.files before first DB save * fix: ESLint error fixed * fix: deduplicate file-population logic and add test coverage Extract `buildMessageFiles` helper into `packages/api/src/utils/message` to replace three near-identical loops in BaseClient and both agent controllers. Fixes set poisoning from undefined file_id entries, moves file population inside the skipSaveUserMessage guard to avoid wasted work, and adds full unit test coverage for the new behavior. * chore: reorder import statements in openIdJwtStrategy.js for consistency --------- Co-authored-by: Danny Avila <danny@librechat.ai>
This commit is contained in:
parent
13df8ed67c
commit
3a079b980a
6 changed files with 258 additions and 48 deletions
|
|
@ -4,6 +4,7 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const {
|
||||
countTokens,
|
||||
getBalanceConfig,
|
||||
buildMessageFiles,
|
||||
extractFileContext,
|
||||
encodeAndFormatAudios,
|
||||
encodeAndFormatVideos,
|
||||
|
|
@ -670,6 +671,14 @@ class BaseClient {
|
|||
}
|
||||
|
||||
if (!isEdited && !this.skipSaveUserMessage) {
|
||||
const reqFiles = this.options.req?.body?.files;
|
||||
if (reqFiles && Array.isArray(this.options.attachments)) {
|
||||
const files = buildMessageFiles(reqFiles, this.options.attachments);
|
||||
if (files.length > 0) {
|
||||
userMessage.files = files;
|
||||
}
|
||||
delete userMessage.image_urls;
|
||||
}
|
||||
userMessagePromise = this.saveMessageToDatabase(userMessage, saveOptions, user);
|
||||
this.savedMessageIds.add(userMessage.messageId);
|
||||
if (typeof opts?.getReqData === 'function') {
|
||||
|
|
|
|||
|
|
@ -928,4 +928,123 @@ describe('BaseClient', () => {
|
|||
expect(result.remainingContextTokens).toBe(2); // 25 - 20 - 3(assistant label)
|
||||
});
|
||||
});
|
||||
|
||||
describe('sendMessage file population', () => {
|
||||
const attachment = {
|
||||
file_id: 'file-abc',
|
||||
filename: 'image.png',
|
||||
filepath: '/uploads/image.png',
|
||||
type: 'image/png',
|
||||
bytes: 1024,
|
||||
object: 'file',
|
||||
user: 'user-1',
|
||||
embedded: false,
|
||||
usage: 0,
|
||||
text: 'large ocr blob that should be stripped',
|
||||
_id: 'mongo-id-1',
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
TestClient.options.req = { body: { files: [{ file_id: 'file-abc' }] } };
|
||||
TestClient.options.attachments = [attachment];
|
||||
});
|
||||
|
||||
test('populates userMessage.files before saveMessageToDatabase is called', async () => {
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockImplementation((msg) => {
|
||||
return Promise.resolve({ message: msg });
|
||||
});
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave).toBeDefined();
|
||||
expect(userSave[0].files).toBeDefined();
|
||||
expect(userSave[0].files).toHaveLength(1);
|
||||
expect(userSave[0].files[0].file_id).toBe('file-abc');
|
||||
});
|
||||
|
||||
test('strips text and _id from files before saving', async () => {
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave[0].files[0].text).toBeUndefined();
|
||||
expect(userSave[0].files[0]._id).toBeUndefined();
|
||||
expect(userSave[0].files[0].filename).toBe('image.png');
|
||||
});
|
||||
|
||||
test('deletes image_urls from userMessage when files are present', async () => {
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
TestClient.options.attachments = [
|
||||
{ ...attachment, image_urls: ['data:image/png;base64,...'] },
|
||||
];
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave[0].image_urls).toBeUndefined();
|
||||
});
|
||||
|
||||
test('does not set files when no attachments match request file IDs', async () => {
|
||||
TestClient.options.req = { body: { files: [{ file_id: 'file-nomatch' }] } };
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave[0].files).toBeUndefined();
|
||||
});
|
||||
|
||||
test('skips file population when attachments is not an array (Promise case)', async () => {
|
||||
TestClient.options.attachments = Promise.resolve([attachment]);
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave[0].files).toBeUndefined();
|
||||
});
|
||||
|
||||
test('skips file population when skipSaveUserMessage is true', async () => {
|
||||
TestClient.skipSaveUserMessage = true;
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg?.isCreatedByUser,
|
||||
);
|
||||
expect(userSave).toBeUndefined();
|
||||
});
|
||||
|
||||
test('ignores file_id: undefined entries in req.body.files (no set poisoning)', async () => {
|
||||
TestClient.options.req = {
|
||||
body: { files: [{ file_id: undefined }, { file_id: 'file-abc' }] },
|
||||
};
|
||||
TestClient.options.attachments = [
|
||||
{ ...attachment, file_id: undefined },
|
||||
{ ...attachment, file_id: 'file-abc' },
|
||||
];
|
||||
TestClient.saveMessageToDatabase = jest.fn().mockResolvedValue({ message: {} });
|
||||
|
||||
await TestClient.sendMessage('Hello');
|
||||
|
||||
const userSave = TestClient.saveMessageToDatabase.mock.calls.find(
|
||||
([msg]) => msg.isCreatedByUser,
|
||||
);
|
||||
expect(userSave[0].files).toHaveLength(1);
|
||||
expect(userSave[0].files[0].file_id).toBe('file-abc');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -3,9 +3,9 @@ const { Constants, ViolationTypes } = require('librechat-data-provider');
|
|||
const {
|
||||
sendEvent,
|
||||
getViolationInfo,
|
||||
buildMessageFiles,
|
||||
GenerationJobManager,
|
||||
decrementPendingRequest,
|
||||
sanitizeFileForTransmit,
|
||||
sanitizeMessageForTransmit,
|
||||
checkAndIncrementPendingRequest,
|
||||
} = require('@librechat/api');
|
||||
|
|
@ -252,13 +252,10 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
conversation.title =
|
||||
conversation && !conversation.title ? null : conversation?.title || 'New Chat';
|
||||
|
||||
if (req.body.files && client.options?.attachments) {
|
||||
userMessage.files = [];
|
||||
const messageFiles = new Set(req.body.files.map((file) => file.file_id));
|
||||
for (const attachment of client.options.attachments) {
|
||||
if (messageFiles.has(attachment.file_id)) {
|
||||
userMessage.files.push(sanitizeFileForTransmit(attachment));
|
||||
}
|
||||
if (req.body.files && Array.isArray(client.options.attachments)) {
|
||||
const files = buildMessageFiles(req.body.files, client.options.attachments);
|
||||
if (files.length > 0) {
|
||||
userMessage.files = files;
|
||||
}
|
||||
delete userMessage.image_urls;
|
||||
}
|
||||
|
|
@ -639,14 +636,10 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle
|
|||
conversation.title =
|
||||
conversation && !conversation.title ? null : conversation?.title || 'New Chat';
|
||||
|
||||
// Process files if needed (sanitize to remove large text fields before transmission)
|
||||
if (req.body.files && client.options?.attachments) {
|
||||
userMessage.files = [];
|
||||
const messageFiles = new Set(req.body.files.map((file) => file.file_id));
|
||||
for (const attachment of client.options.attachments) {
|
||||
if (messageFiles.has(attachment.file_id)) {
|
||||
userMessage.files.push(sanitizeFileForTransmit(attachment));
|
||||
}
|
||||
if (req.body.files && Array.isArray(client.options.attachments)) {
|
||||
const files = buildMessageFiles(req.body.files, client.options.attachments);
|
||||
if (files.length > 0) {
|
||||
userMessage.files = files;
|
||||
}
|
||||
delete userMessage.image_urls;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,8 +4,8 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { HttpsProxyAgent } = require('https-proxy-agent');
|
||||
const { SystemRoles } = require('librechat-data-provider');
|
||||
const { isEnabled, findOpenIDUser, math } = require('@librechat/api');
|
||||
const { getOpenIdEmail } = require('./openidStrategy');
|
||||
const { Strategy: JwtStrategy, ExtractJwt } = require('passport-jwt');
|
||||
const { getOpenIdEmail } = require('./openidStrategy');
|
||||
const { updateUser, findUser } = require('~/models');
|
||||
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue