diff --git a/api/server/services/Threads/manage.js b/api/server/services/Threads/manage.js index 772cbd977b..27068a60fd 100644 --- a/api/server/services/Threads/manage.js +++ b/api/server/services/Threads/manage.js @@ -8,7 +8,7 @@ const { AnnotationTypes, defaultOrderQuery, } = require('librechat-data-provider'); -const { recordMessage, getMessages, spendTokens, saveConvo } = require('~/models'); +const { saveMessage, getMessages, spendTokens, saveConvo } = require('~/models'); const { retrieveAndProcessFile } = require('~/server/services/Files/process'); /** @@ -90,16 +90,17 @@ async function saveUserMessage(req, params) { convo.file_ids = params.file_ids; } - const message = await recordMessage(userMessage); - await saveConvo( - { - userId: req?.user?.id, - isTemporary: req?.body?.isTemporary, - interfaceConfig: req?.config?.interfaceConfig, - }, - convo, - { context: 'api/server/services/Threads/manage.js #saveUserMessage' }, - ); + const retentionContext = { + userId: req?.user?.id, + isTemporary: req?.body?.isTemporary, + interfaceConfig: req?.config?.interfaceConfig, + }; + const message = await saveMessage(retentionContext, userMessage, { + context: 'api/server/services/Threads/manage.js #saveUserMessage', + }); + await saveConvo(retentionContext, convo, { + context: 'api/server/services/Threads/manage.js #saveUserMessage', + }); return message; } @@ -128,31 +129,36 @@ async function saveUserMessage(req, params) { async function saveAssistantMessage(req, params) { // const tokenCount = // TODO: need to count each content part - const message = await recordMessage({ - user: params.user, - endpoint: params.endpoint, - messageId: params.messageId, - conversationId: params.conversationId, - parentMessageId: params.parentMessageId, - thread_id: params.thread_id, - /* For messages, use the assistant_id instead of model */ - model: params.assistant_id, - content: params.content, - sender: 'Assistant', - isCreatedByUser: false, - text: params.text, - unfinished: false, - // tokenCount, - iconURL: params.iconURL, - spec: params.spec, - }); + const retentionContext = { + userId: req?.user?.id, + isTemporary: req?.body?.isTemporary, + interfaceConfig: req?.config?.interfaceConfig, + }; + const message = await saveMessage( + retentionContext, + { + user: params.user, + endpoint: params.endpoint, + messageId: params.messageId, + conversationId: params.conversationId, + parentMessageId: params.parentMessageId, + thread_id: params.thread_id, + /* For messages, use the assistant_id instead of model */ + model: params.assistant_id, + content: params.content, + sender: 'Assistant', + isCreatedByUser: false, + text: params.text, + unfinished: false, + // tokenCount, + iconURL: params.iconURL, + spec: params.spec, + }, + { context: 'api/server/services/Threads/manage.js #saveAssistantMessage' }, + ); await saveConvo( - { - userId: req?.user?.id, - isTemporary: req?.body?.isTemporary, - interfaceConfig: req?.config?.interfaceConfig, - }, + retentionContext, { endpoint: params.endpoint, conversationId: params.conversationId, @@ -236,7 +242,17 @@ async function syncMessages({ * @param {dbMessage} params.apiMessage */ const processNewMessage = async ({ dbMessage, apiMessage }) => { - recordPromises.push(recordMessage({ ...dbMessage, user: openai.req.user.id })); + recordPromises.push( + saveMessage( + { + userId: openai.req?.user?.id, + isTemporary: openai.req?.body?.isTemporary, + interfaceConfig: openai.req?.config?.interfaceConfig, + }, + { ...dbMessage, user: openai.req.user.id }, + { context: 'api/server/services/Threads/manage.js #syncMessages' }, + ), + ); if (!apiMessage.id.includes('msg_')) { return; diff --git a/api/server/services/Threads/manage.retention.spec.js b/api/server/services/Threads/manage.retention.spec.js new file mode 100644 index 0000000000..71ce6249e5 --- /dev/null +++ b/api/server/services/Threads/manage.retention.spec.js @@ -0,0 +1,115 @@ +/** + * Integration test for forced-retention enforcement on the Assistants chat path. + * + * `saveUserMessage` / `saveAssistantMessage` (and the `syncMessages` reconciliation) + * used to persist messages through `recordMessage`, which bypasses retention. The + * conversation was forced temporary by `saveConvo` while its messages were stored with + * no `isTemporary`/`expiredAt`, leaving them uncovered by the forced-retention TTL. + * These tests exercise the real save path against an in-memory MongoDB to prove the + * messages now inherit the same retention as their conversation. + */ + +jest.mock('~/server/services/Files/process', () => ({ + retrieveAndProcessFile: jest.fn(), +})); + +// `countTokens` lazy-loads its tokenizer via a dynamic ESM import that jest cannot +// execute under CommonJS; stub the token count while leaving the rest of the module real. +jest.mock('@librechat/api', () => ({ + ...jest.requireActual('@librechat/api'), + countTokens: jest.fn().mockResolvedValue(1), +})); + +const { v4 } = require('uuid'); +const mongoose = require('mongoose'); +const { MongoMemoryServer } = require('mongodb-memory-server'); +const { createModels } = require('@librechat/data-schemas'); +const { RetentionMode } = require('librechat-data-provider'); + +const { saveUserMessage, saveAssistantMessage } = require('./manage'); + +describe('Threads/manage forced retention', () => { + let mongoServer; + let Message; + const userId = new mongoose.Types.ObjectId().toString(); + + const buildReq = (interfaceConfig) => ({ + user: { id: userId }, + body: {}, + config: { interfaceConfig }, + }); + + beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + Message = mongoose.models.Message; + }, 30000); + + afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); + }); + + afterEach(async () => { + await Message.deleteMany({}); + }); + + it('forces user messages temporary with an expiration in ephemeral mode', async () => { + const conversationId = v4(); + await saveUserMessage(buildReq({ retentionMode: RetentionMode.EPHEMERAL }), { + user: userId, + endpoint: 'assistants', + messageId: v4(), + conversationId, + text: 'hello', + assistant_id: 'asst_1', + thread_id: 'thread_1', + }); + + const saved = await Message.findOne({ conversationId }).lean(); + expect(saved).toBeTruthy(); + expect(saved.isTemporary).toBe(true); + expect(saved.expiredAt).toBeInstanceOf(Date); + expect(saved.expiredAt.getTime()).toBeGreaterThan(Date.now()); + }); + + it('forces assistant messages temporary with an expiration in ephemeral mode', async () => { + const conversationId = v4(); + await saveAssistantMessage(buildReq({ retentionMode: RetentionMode.EPHEMERAL }), { + user: userId, + endpoint: 'assistants', + messageId: v4(), + conversationId, + parentMessageId: v4(), + text: 'hi there', + assistant_id: 'asst_1', + thread_id: 'thread_1', + content: [{ type: 'text', text: 'hi there' }], + }); + + const saved = await Message.findOne({ conversationId }).lean(); + expect(saved).toBeTruthy(); + expect(saved.isCreatedByUser).toBe(false); + expect(saved.isTemporary).toBe(true); + expect(saved.expiredAt).toBeInstanceOf(Date); + expect(saved.expiredAt.getTime()).toBeGreaterThan(Date.now()); + }); + + it('does not force an expiration when retention is not ephemeral', async () => { + const conversationId = v4(); + await saveUserMessage(buildReq({ retentionMode: RetentionMode.TEMPORARY }), { + user: userId, + endpoint: 'assistants', + messageId: v4(), + conversationId, + text: 'hello', + assistant_id: 'asst_1', + thread_id: 'thread_1', + }); + + const saved = await Message.findOne({ conversationId }).lean(); + expect(saved).toBeTruthy(); + expect(saved.expiredAt ?? null).toBeNull(); + }); +}); diff --git a/packages/api/src/files/retention.spec.ts b/packages/api/src/files/retention.spec.ts index 1fb8bc213e..c2d49b8e66 100644 --- a/packages/api/src/files/retention.spec.ts +++ b/packages/api/src/files/retention.spec.ts @@ -419,5 +419,59 @@ describe('retention helpers', () => { ).resolves.toBe(expiredAt); expect(dependencies.createExpirationDate).not.toHaveBeenCalled(); }); + + it('caps the share at an active source conversation expiration', async () => { + const conversationExpiredAt = new Date(Date.now() + 60 * 60 * 1000); + dependencies.getConvo.mockResolvedValue({ expiredAt: conversationExpiredAt }); + + await expect( + getSharedLinkExpiration( + { + req: request({ + config: { interfaceConfig: { retentionMode: RetentionMode.EPHEMERAL } }, + }), + conversationId: 'convo-1', + }, + dependencies, + ), + ).resolves.toBe(conversationExpiredAt); + }); + + it('uses the fresh window when it expires before the source conversation', async () => { + const conversationExpiredAt = new Date('2031-01-01T00:00:00.000Z'); + dependencies.getConvo.mockResolvedValue({ expiredAt: conversationExpiredAt }); + + await expect( + getSharedLinkExpiration( + { + req: request({ + config: { interfaceConfig: { retentionMode: RetentionMode.EPHEMERAL } }, + }), + conversationId: 'convo-1', + }, + dependencies, + ), + ).resolves.toBe(expirationDate); + }); + + it('falls back to the active source expiration when creating a window throws', async () => { + const conversationExpiredAt = new Date(Date.now() + 60 * 60 * 1000); + dependencies.getConvo.mockResolvedValue({ expiredAt: conversationExpiredAt }); + dependencies.createExpirationDate.mockImplementation(() => { + throw new Error('boom'); + }); + + await expect( + getSharedLinkExpiration( + { + req: request({ + config: { interfaceConfig: { retentionMode: RetentionMode.EPHEMERAL } }, + }), + conversationId: 'convo-1', + }, + dependencies, + ), + ).resolves.toBe(conversationExpiredAt); + }); }); }); diff --git a/packages/api/src/files/retention.ts b/packages/api/src/files/retention.ts index f8a2ff12d7..96a510eadb 100644 --- a/packages/api/src/files/retention.ts +++ b/packages/api/src/files/retention.ts @@ -202,6 +202,11 @@ export async function getAgentFileRetentionExpiry( * - `undefined`: no decision can be made because the conversation id or row is missing. * - `null`: the share should be stored without an expiration. * - `Date`: the share should expire at that date; callers reject already-expired dates. + * + * A share embeds a snapshot of the source conversation's messages, so it must never + * outlive the conversation it was created from. When the source conversation still has an + * active expiration, the share is capped at the earlier of that deadline and a freshly + * created retention window rather than starting a brand-new window. */ export async function getSharedLinkExpiration( { @@ -234,10 +239,14 @@ export async function getSharedLinkExpiration( } try { - return dependencies.createExpirationDate(req?.config?.interfaceConfig); + const createdExpiration = dependencies.createExpirationDate(req?.config?.interfaceConfig); + if (conversationExpiredAt != null && conversationExpiredAt < createdExpiration) { + return conversationExpiredAt; + } + return createdExpiration; } catch (err) { dependencies.logger?.error('[getSharedLinkExpiration] Error creating expiration date:', err); - return null; + return conversationExpiredAt ?? null; } }