📔 feat: Preserve Owner Text During PII Redaction

api/app/clients/BaseClient.js

@@ -32,6 +32,12 @@ const {
   announceReply,
   needsRetentionConversation,
   getConversationWriteContext,
+  savePrivateTextMessage,
+  stampPrivateTextMessage,
+  deferPrivateTextStart,
+  requirePrivateTextPersistence,
+  rejectPrivateTextAdmission,
+  bindPrivateTextPersistenceAbort,
   persistedReasoningOverrideFields,
 } = require('@librechat/api');
 const {
@@ -589,13 +595,16 @@ class BaseClient {
     } = await this.setMessageOptions(opts);
     this.options.startupTelemetry?.mark('history_loaded');
 
-    const userMessage = this.resolveStartUserMessage({
-      opts,
-      message,
-      userMessageId,
-      parentMessageId,
-      conversationId,
-    });
+    const userMessage = stampPrivateTextMessage(
+      this.options.req,
+      this.resolveStartUserMessage({
+        opts,
+        message,
+        userMessageId,
+        parentMessageId,
+        conversationId,
+      }),
+    );
 
     /**
      * Attach quoted excerpts (the "Add to chat" selections from `req.body.quotes`)
@@ -630,10 +639,13 @@ class BaseClient {
       });
     }
 
-    if (typeof opts?.onStart === 'function') {
-      const isNewConvo = !requestConvoId && parentMessageId === Constants.NO_PARENT;
-      opts.onStart(userMessage, responseMessageId, isNewConvo);
-    }
+    this.privateTextStart = deferPrivateTextStart(
+      this.options.req,
+      opts?.onStart,
+      userMessage,
+      responseMessageId,
+      !requestConvoId && parentMessageId === Constants.NO_PARENT,
+    );
 
     return {
       ...opts,
@@ -985,17 +997,12 @@ class BaseClient {
           start,
           cancel,
         });
-        const requestAbortSignal = this.abortController?.signal;
-        if (requestAbortSignal?.aborted) {
-          /** Preserve the historical durability contract for Stop: abort
-           * persistence may publish the partial assistant response before the
-           * provider unwinds, so its parent write must already be underway. */
-          start();
-        } else if (requestAbortSignal != null) {
-          const startOnAbort = () => start();
-          requestAbortSignal.addEventListener('abort', startOnAbort, { once: true });
-          removeAbortListener = () => requestAbortSignal.removeEventListener('abort', startOnAbort);
-        }
+        removeAbortListener = bindPrivateTextPersistenceAbort(
+          this.options.req,
+          this.abortController?.signal,
+          start,
+          cancel,
+        );
         this.modelBoundUserMessagePersistence = userMessagePersistence;
         userMessagePromise = persistencePromise;
       } else {
@@ -1043,6 +1050,7 @@ class BaseClient {
 
       completionResult = await this.sendCompletion(payload, opts);
     } catch (error) {
+      rejectPrivateTextAdmission(this.options.req);
       if (userMessagePersistence?.isPending()) {
         if (isContentFilterError(error)) {
           userMessagePersistence.cancel();
@@ -1054,6 +1062,11 @@ class BaseClient {
     }
     /** A safe no-model completion (or a runtime that cannot expose the
      * admission callback) must not leave the parent-write gate pending. */
+    await requirePrivateTextPersistence(
+      this.options.req,
+      () => (userMessagePersistence != null ? userMessagePersistence.start() : userMessagePromise),
+      this.privateTextStart,
+    );
     userMessagePersistence?.start();
     const { completion, metadata } = completionResult;
     if (this.abortController) {
@@ -1272,7 +1285,8 @@ class BaseClient {
       return [];
     }
 
-    const messages = (await db.getMessages({ conversationId, user: this.user })) ?? [];
+    const messages =
+      (await db.getMessages({ conversationId, user: this.user }, '+privateTextTokens')) ?? [];
     /** A client that reads beyond the walk below (which stops at a checkpoint
      *  summary) receives every row here; the rest keep nothing. */
     this.onHistoryLoaded?.(messages);
@@ -1361,7 +1375,9 @@ class BaseClient {
       req.resolvedConversation = await db.getConvo(req.user.id, message.conversationId);
     }
     const reqCtx = getConversationWriteContext(req);
-    const savedMessage = await db.saveMessage(
+    const savedMessage = await savePrivateTextMessage(
+      db.saveMessage,
+      req,
       reqCtx,
       {
         ...message,

api/app/clients/specs/BaseClient.test.js

@@ -2,6 +2,11 @@ const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-prov
 const BaseClientClass = require('../BaseClient');
 const {
   ContentFilterError,
+  createPrivateTextIngress,
+  createModelBoundChatModelCallback,
+  getPrivateTextAdmission,
+  getPrivateTextInspectionTokens,
+  assertModelBoundContent,
   resolveTurnDeliveryRouting,
   buildSteerMedia,
   Tokenizer,
@@ -1597,7 +1602,7 @@ describe('BaseClient', () => {
 
       const chatMessages = await TestClient.loadHistory(conversationId, '1');
 
-      expect(getMessages).toHaveBeenCalledWith({ conversationId, user });
+      expect(getMessages).toHaveBeenCalledWith({ conversationId, user }, '+privateTextTokens');
       expect(chatMessages).toHaveLength(1);
       expect(chatMessages[0].text).toBe('Hello');
     });
@@ -1734,6 +1739,195 @@ describe('BaseClient', () => {
       );
     });
 
+    function protectedClient(history = [], legacyPii) {
+      const filters = {
+        messages: {
+          pii: {
+            action: 'redact',
+            fields: ['text'],
+            starterPatterns: [],
+            customPatterns: [
+              { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' },
+            ],
+          },
+        },
+      };
+      const req = {
+        user: { id: 'owner' },
+        path: '/',
+        body: { text: 'alice@example.com', clientRequestId: 'created-privacy' },
+        config: { filters, messageFilter: { pii: legacyPii } },
+      };
+      const next = jest.fn();
+      createPrivateTextIngress({
+        getFilters: () => filters,
+        getLegacyPii: () => legacyPii,
+        getKey: () => 'ab'.repeat(32),
+      })(req, { status: jest.fn().mockReturnThis(), json: jest.fn() }, next);
+      expect(next).toHaveBeenCalledTimes(1);
+      const client = initializeFakeClient(apiKey, { ...options, req }, history);
+      client.shouldDeferUserMessagePersistence = () => true;
+      client.assertStoredModelBoundContent = () =>
+        assertModelBoundContent({
+          legacyPii,
+          storedMessages: client.modelBoundStoredMessages,
+        });
+      client.assertBuiltModelBoundContent = () => {};
+      client.saveMessageToDatabase = jest.fn(async (message) => ({ message }));
+      const provider = jest.fn();
+      client.sendCompletion = jest.fn(async (payload) => {
+        const callback = createModelBoundChatModelCallback(
+          {
+            filters,
+            legacyPii,
+            storedMessages: client.modelBoundStoredMessages,
+            privateTextTokens: getPrivateTextInspectionTokens(client.modelBoundStoredMessages),
+          },
+          {
+            onContentRejected: client.modelBoundUserMessagePersistence?.cancel,
+            onContentAllowed: getPrivateTextAdmission(
+              req,
+              client.modelBoundUserMessagePersistence?.start,
+              client.privateTextStart,
+            ),
+          },
+        );
+        await callback.handleChatModelStart(undefined, [payload]);
+        provider();
+        return { completion: 'Safe reply' };
+      });
+      return { client, req, provider };
+    }
+
+    test('protected startup remains deferred until exact admission and the atomic write finishes', async () => {
+      const { client, req, provider } = protectedClient();
+      const committed = deferred();
+      let written;
+      client.saveMessageToDatabase.mockImplementationOnce((message) => {
+        written = message;
+        return committed.promise;
+      });
+      const onStart = jest.fn();
+      const sent = client.sendMessage(req.body.text, { onStart });
+      // Wait for the admission callback to begin the real deferred write.
+      for (let i = 0; i < 30 && !written; i++) {
+        await Promise.resolve();
+      }
+      expect(written).toBeDefined();
+      expect(onStart).not.toHaveBeenCalled();
+      expect(provider).not.toHaveBeenCalled();
+      committed.resolve({ message: written });
+      await sent;
+      expect(onStart).toHaveBeenCalledTimes(1);
+      expect(provider).toHaveBeenCalledTimes(1);
+      expect(onStart.mock.invocationCallOrder[0]).toBeLessThan(
+        provider.mock.invocationCallOrder[0],
+      );
+    });
+
+    test('protected write failure prevents created and the provider call', async () => {
+      const { client, req, provider } = protectedClient();
+      client.saveMessageToDatabase.mockResolvedValueOnce({});
+      const onStart = jest.fn();
+      await expect(client.sendMessage(req.body.text, { onStart })).rejects.toThrow();
+      expect(onStart).not.toHaveBeenCalled();
+      expect(provider).not.toHaveBeenCalled();
+    });
+
+    test.each([false, true])(
+      'cancels protected Stop before native admission, pre-aborted: %s',
+      async (preAborted) => {
+        const { client, req, provider } = protectedClient();
+        const controller = new AbortController();
+        if (preAborted) {
+          controller.abort();
+        }
+        const ready = deferred();
+        const completion = deferred();
+        const onStart = jest.fn();
+        client.sendCompletion = jest.fn(async () => {
+          ready.resolve();
+          return completion.promise;
+        });
+        const sent = client.sendMessage(req.body.text, { abortController: controller, onStart });
+        const observed = sent.catch((error) => error);
+        await ready.promise;
+        if (!preAborted) {
+          controller.abort();
+        }
+        completion.resolve({ completion: 'Stopped before model' });
+        expect(await observed).toEqual(expect.objectContaining({ code: 'content_filter_block' }));
+        expect(client.saveMessageToDatabase).not.toHaveBeenCalled();
+        expect(onStart).not.toHaveBeenCalled();
+        expect(provider).not.toHaveBeenCalled();
+      },
+    );
+
+    test.each(['user-id', 'user-id__1', 'user-id__invalid'])(
+      'rejects a protected persistence-skipping override %s before model invocation',
+      async (overrideUserMessageId) => {
+        const { client, req, provider } = protectedClient();
+        req.body.overrideUserMessageId = overrideUserMessageId;
+        const onStart = jest.fn();
+        await expect(client.sendMessage(req.body.text, { onStart })).rejects.toMatchObject({
+          code: 'content_filter_block',
+        });
+        expect(client.skipSaveUserMessage).toBe(true);
+        expect(client.saveMessageToDatabase).not.toHaveBeenCalled();
+        expect(onStart).not.toHaveBeenCalled();
+        expect(provider).not.toHaveBeenCalled();
+      },
+    );
+
+    test('retains protected admission for the normal browser override with writer index zero', async () => {
+      const { client, req, provider } = protectedClient();
+      req.body.overrideUserMessageId = 'normal-user-id__0';
+      const onStart = jest.fn();
+      await expect(client.sendMessage(req.body.text, { onStart })).resolves.toBeDefined();
+      expect(client.skipSaveUserMessage).toBe(false);
+      expect(onStart).toHaveBeenCalledWith(
+        expect.objectContaining({
+          messageId: 'normal-user-id',
+          privacyRevision: expect.any(String),
+        }),
+        expect.any(String),
+        true,
+      );
+      expect(provider).toHaveBeenCalledTimes(1);
+    });
+
+    test('a legacy history rejection leaves a transformed turn and conversation unsaved', async () => {
+      const history = [{ messageId: 'prior', text: 'LEGACY-SECRET', isCreatedByUser: true }];
+      const { client, req, provider } = protectedClient(history, {
+        starterPatterns: [],
+        customPatterns: [{ id: 'legacy', label: 'Legacy', regex: 'LEGACY-SECRET' }],
+      });
+      const onStart = jest.fn();
+      await expect(
+        client.sendMessage(req.body.text, {
+          conversationId: 'conversation',
+          parentMessageId: 'prior',
+          onStart,
+        }),
+      ).rejects.toThrow();
+      expect(client.saveMessageToDatabase).not.toHaveBeenCalled();
+      expect(onStart).not.toHaveBeenCalled();
+      expect(provider).not.toHaveBeenCalled();
+    });
+
+    test('an exact model-input rejection cancels the protected deferred write before created', async () => {
+      const { client, req, provider } = protectedClient();
+      const onStart = jest.fn();
+      client.buildMessages.mockResolvedValueOnce({
+        prompt: [{ role: 'user', content: 'alice@example.com' }],
+      });
+      await expect(client.sendMessage(req.body.text, { onStart })).rejects.toThrow();
+      expect(client.saveMessageToDatabase).not.toHaveBeenCalled();
+      expect(onStart).not.toHaveBeenCalled();
+      expect(provider).not.toHaveBeenCalled();
+      expect(client.modelBoundUserMessagePersistence.isPending()).toBe(false);
+    });
+
     test('onStart is called with the correct arguments', async () => {
       const onStart = jest.fn();
       const opts = { onStart };

api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js

@@ -47,6 +47,10 @@ jest.mock('@librechat/data-schemas', () => ({
 }));
 
 jest.mock('@librechat/api', () => ({
+  savePrivateTextMessage: (save, _req, ...args) => save(...args),
+  savePrivateTextErrorTurn: (...args) =>
+    jest.requireActual('@librechat/api').savePrivateTextErrorTurn(...args),
+  stampPreliminaryPrivateTextMessage: (_req, message) => message,
   getAgentErrorMetadata: (...args) =>
     jest.requireActual('@librechat/api').getAgentErrorMetadata(...args),
   applyForcedTemporaryRequest: jest.fn(),