From 7117f1d40c016fbe5bdc370221806d95ce28ec52 Mon Sep 17 00:00:00 2001 From: Lia Date: Thu, 24 Sep 2026 13:28:16 +0000 Subject: [PATCH 01/18] =?UTF-8?q?=F0=9F=93=94=20feat:=20Preserve=20Owner?= =?UTF-8?q?=20Text=20During=20PII=20Redaction?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/app/clients/BaseClient.js | 27 +- .../request.partialDisconnect.spec.js | 1 + .../__tests__/request.resumeMetadata.spec.js | 1 + api/server/controllers/agents/request.js | 21 +- api/server/routes/agents/chat.js | 8 + api/server/routes/messages.js | 8 + client/src/components/Chat/ChatView.tsx | 9 +- .../Chat/Messages/PrivateText.spec.tsx | 144 +++++++++++ .../components/Chat/Messages/PrivateText.tsx | 128 ++++++++++ .../Chat/Messages/ui/MessageRender.tsx | 41 +-- .../src/components/Messages/ContentRender.tsx | 45 ++-- client/src/hooks/Chat/useChatFunctions.ts | 9 +- client/src/locales/en/translation.json | 3 + client/src/utils/messages.ts | 1 + e2e/setup/fake-model.js | 16 ++ e2e/specs/mock/private-text.spec.ts | 121 +++++++++ librechat.example.yaml | 8 +- packages/api/src/protection/index.ts | 2 + packages/api/src/protection/private/crypto.ts | 53 ++++ .../src/protection/private/submission.spec.ts | 219 ++++++++++++++++ .../api/src/protection/private/submission.ts | 238 ++++++++++++++++++ packages/api/src/protection/private/view.ts | 65 +++++ packages/data-provider/src/data-service.ts | 15 ++ packages/data-provider/src/schemas.ts | 2 + .../src/methods/message.privateText.spec.ts | 160 ++++++++++++ packages/data-schemas/src/methods/message.ts | 64 ++++- packages/data-schemas/src/schema/message.ts | 2 + packages/data-schemas/src/types/message.ts | 3 + 28 files changed, 1357 insertions(+), 57 deletions(-) create mode 100644 client/src/components/Chat/Messages/PrivateText.spec.tsx create mode 100644 client/src/components/Chat/Messages/PrivateText.tsx create mode 100644 e2e/specs/mock/private-text.spec.ts create mode 100644 packages/api/src/protection/private/crypto.ts create mode 100644 packages/api/src/protection/private/submission.spec.ts create mode 100644 packages/api/src/protection/private/submission.ts create mode 100644 packages/api/src/protection/private/view.ts create mode 100644 packages/data-schemas/src/methods/message.privateText.spec.ts diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 4a4f6975bda..62eacccb4be 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -30,6 +30,9 @@ const { seedTurnConversation, needsRetentionConversation, getConversationWriteContext, + savePrivateTextMessage, + stampPrivateTextMessage, + requirePrivateTextPersistence, } = require('@librechat/api'); const { Constants, @@ -586,13 +589,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`) @@ -1029,6 +1035,9 @@ class BaseClient { await balanceReservations.track(balanceAdmission); } + await requirePrivateTextPersistence(this.options.req, () => + userMessagePersistence != null ? userMessagePersistence.start() : userMessagePromise, + ); completionResult = await this.sendCompletion(payload, opts); } catch (error) { if (userMessagePersistence?.isPending()) { @@ -1349,7 +1358,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, diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 437e74c45d7..ac82e669ac5 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -47,6 +47,7 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/api', () => ({ + savePrivateTextMessage: (save, _req, ...args) => save(...args), sendEvent: jest.fn(), isScheduleFireRequest: jest.fn(() => false), exemptFromConcurrencyLimiter: jest.fn(() => false), diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index 058d3302dc8..e4689b832f0 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -252,6 +252,7 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/api', () => ({ + savePrivateTextMessage: (save, _req, ...args) => save(...args), sendEvent: jest.fn(), logAgentMemorySnapshot: jest.fn(), isScheduleFireRequest: (...args) => mockIsScheduleFireRequest(...args), diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 3cfa9d71aed..9730cbf06cd 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -50,6 +50,7 @@ const { resolvePersistableCodeEnvironmentDecision, getFailedTurnTraceFields, resolveFailedTurnContent, + savePrivateTextMessage, } = require('@librechat/api'); const { disposeClient } = require('~/server/cleanup'); const { @@ -458,7 +459,9 @@ async function saveErrorTurn( const iconURL = getEndpointIconURL(req, endpointOption); if (userMessage) { - const savedUserMessage = await saveMessage( + const savedUserMessage = await savePrivateTextMessage( + saveMessage, + req, reqCtx, { ...userMessage, @@ -2729,7 +2732,9 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit convoSignal.observeMessageWrite(Promise.resolve(savedUserTurn)); } else { // Custom clients used by integrations/tests may not inherit BaseClient. - const savedUserMessage = await saveMessage( + const savedUserMessage = await savePrivateTextMessage( + saveMessage, + req, { userId, isTemporary: @@ -2982,9 +2987,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit if (!userMessage) { throw new Error('User message was unavailable before terminal persistence'); } - const savedUserMessage = await saveMessage(reqCtx, userMessage, { - context: 'api/server/controllers/agents/request.js - resumable user message', - }); + const savedUserMessage = await savePrivateTextMessage( + saveMessage, + req, + reqCtx, + userMessage, + { + context: 'api/server/controllers/agents/request.js - resumable user message', + }, + ); if (!savedUserMessage) { throw new Error('User message could not be persisted before terminal publication'); } diff --git a/api/server/routes/agents/chat.js b/api/server/routes/agents/chat.js index 16ecf54fd04..d25d5e353b2 100644 --- a/api/server/routes/agents/chat.js +++ b/api/server/routes/agents/chat.js @@ -2,6 +2,7 @@ const express = require('express'); const { logger } = require('@librechat/data-schemas'); const { createMessageFilterPii, + createPrivateTextIngress, reportLocatorTraversalFailure, generateCheckAccess, skipAgentCheck, @@ -71,6 +72,13 @@ const restoreResumeContext = async (req, res, next) => { }; router.use(restoreResumeContext); +router.use( + createPrivateTextIngress({ + getFilters: (req) => req.config?.filters, + getLegacyPii: (req) => req.config?.messageFilter?.pii, + getKey: () => process.env.CREDS_KEY ?? '', + }), +); router.use( createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, diff --git a/api/server/routes/messages.js b/api/server/routes/messages.js index 6ac391c18d4..fb56d950363 100644 --- a/api/server/routes/messages.js +++ b/api/server/routes/messages.js @@ -27,6 +27,7 @@ const { mergeUserSubmittedMessageFieldPaths, isContentFilterError, withoutTraceRefs, + createPrivateTextView, } = require('@librechat/api'); const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/Artifacts/update'); @@ -63,6 +64,13 @@ const storedMessageMutationMiddleware = [ ]; router.use(requireJwtAuth); +router.post( + '/:conversationId/owner-text', + createPrivateTextView({ + read: db.getPrivateMessageTexts, + getKey: () => process.env.CREDS_KEY ?? '', + }), +); async function rejectSubagentThreadWrite(req, res, conversationId) { const blocked = await isSubagentThreadWriteBlocked( diff --git a/client/src/components/Chat/ChatView.tsx b/client/src/components/Chat/ChatView.tsx index a7bd0361f1e..fd43bf554a9 100644 --- a/client/src/components/Chat/ChatView.tsx +++ b/client/src/components/Chat/ChatView.tsx @@ -22,6 +22,7 @@ import { QueuedTurnPortalProvider } from './Steering/QueuedTurnPortal'; import ApprovalProvider from './Messages/Content/ApprovalContext'; import ConversationStarters from './Input/ConversationStarters'; import { pendingApprovalActionFamily } from './approval/state'; +import { OwnerTextProvider } from './Messages/PrivateText'; import { useGetMessagesByConvoId } from '~/data-provider'; import Footer, { useConfiguredFooter } from './Footer'; import { AskAnswerHostProvider } from './ask/state'; @@ -187,7 +188,13 @@ function ChatView({ index = 0, project }: { index?: number; project?: TChatProje : 'h-full overflow-y-auto', )} > - {content} + + {content} + {/* Named + opaque so a view transition (the ask_user_question popover ⇄ chat-card morph) paints the whole composer band over the travelling card instead of letting it show diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx new file mode 100644 index 00000000000..c8a3be3de5d --- /dev/null +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -0,0 +1,144 @@ +import { dataService } from 'librechat-data-provider'; +import { render, screen, waitFor, act } from '@testing-library/react'; +import type { TMessage } from 'librechat-data-provider'; +import { OwnerTextProvider, PrivateText } from './PrivateText'; + +let mockOwnerId = 'owner'; +let mockTenantId = 'tenant-a'; +jest.mock('~/hooks/AuthContext', () => ({ + useAuthContext: () => ({ user: { id: mockOwnerId, tenantId: mockTenantId } }), +})); +jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); +jest.mock('librechat-data-provider', () => ({ dataService: { getOwnerMessageTexts: jest.fn() } })); + +const canonical = Object.freeze({ + messageId: 'message', + conversationId: 'conversation', + isCreatedByUser: true, + text: '[EMAIL_1_turn]', + privacyRevision: 'turn', +}) as TMessage; +const load = dataService.getOwnerMessageTexts as jest.Mock; +const original = { + canonicalText: canonical.text, + messageId: 'message', + revision: 'turn', + text: 'alice@example.com', +}; +function View({ + conversationId = 'conversation', + messages = [canonical], +}: { + conversationId?: string; + messages?: TMessage[]; +}) { + return ( + + + 
{JSON.stringify(messages)}
+ + ); +} +beforeEach(() => { + mockOwnerId = 'owner'; + mockTenantId = 'tenant-a'; + load.mockReset(); +}); + +it('renders originals without mutating canonical model/export input, and reloads from the private API', async () => { + load.mockResolvedValue({ messages: [original] }); + const first = render(); + expect(await screen.findByText('alice@example.com')).toBeInTheDocument(); + expect(screen.getByTestId('canonical')).not.toHaveTextContent('alice@example.com'); + expect(canonical.text).toBe('[EMAIL_1_turn]'); + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_hidden'); + first.unmount(); + render(); + expect(await screen.findByText('alice@example.com')).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(2); +}); + +it('renders only filtered text without an owner provider, as on external viewers', () => { + render(); + expect(screen.getByText(canonical.text)).toBeInTheDocument(); + expect(load).not.toHaveBeenCalled(); +}); + +it('shows loading then safe unavailable text when decryption or authorization fails', async () => { + let finish!: (value: { messages: [] }) => void; + load.mockReturnValue( + new Promise((resolve) => { + finish = resolve; + }), + ); + render(); + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_loading'); + await act(async () => { + finish({ messages: [] }); + }); + expect(screen.getByText(canonical.text)).toBeInTheDocument(); + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'); +}); + +it('rejects stale revisions instead of restoring a previous original', async () => { + load.mockResolvedValue({ messages: [{ ...original, revision: 'old-revision' }] }); + render(); + await waitFor(() => + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), + ); + expect(screen.queryByText(original.text)).not.toBeInTheDocument(); +}); + +it('clears the visible original immediately on account switching, ignoring late responses', async () => { + let finish!: (value: { messages: (typeof original)[] }) => void; + load.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + load.mockResolvedValue({ messages: [] }); + const view = render(); + mockOwnerId = 'another-owner'; + view.rerender(); + await act(async () => { + finish({ messages: [original] }); + }); + expect(screen.queryByText(original.text)).not.toBeInTheDocument(); + expect(screen.getByText(canonical.text)).toBeInTheDocument(); +}); + +it('batches selected private rows and never loads ordinary messages', async () => { + load.mockResolvedValue({ messages: [] }); + const messages: TMessage[] = Array.from({ length: 51 }, (_, index) => ({ + ...canonical, + messageId: `message-${index}`, + })); + messages.push({ ...canonical, messageId: 'plain', privacyRevision: undefined }); + render(); + await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); + expect(load.mock.calls.map(([, ids]) => ids.length)).toEqual([50, 1]); + expect(load.mock.calls.flatMap(([, ids]) => ids)).not.toContain('plain'); +}); + +it('invalidates an already rendered original when the canonical message changes', async () => { + load.mockResolvedValue({ messages: [original] }); + const view = render(); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + view.rerender(); + expect(screen.queryByText(original.text)).not.toBeInTheDocument(); + await waitFor(() => + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), + ); +}); + +it('clears originals when tenant identity changes even if the user ID is unchanged', async () => { + load.mockResolvedValueOnce({ messages: [original] }); + const view = render(); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + load.mockResolvedValue({ messages: [] }); + mockTenantId = 'tenant-b'; + view.rerender(); + expect(screen.queryByText(original.text)).not.toBeInTheDocument(); + await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); +}); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx new file mode 100644 index 00000000000..e10d3876320 --- /dev/null +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -0,0 +1,128 @@ +import { createContext, useContext, useEffect, useMemo, useState } from 'react'; +import { dataService } from 'librechat-data-provider'; +import type { TMessage } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import { useAuthContext } from '~/hooks/AuthContext'; +import { useLocalize } from '~/hooks'; + +interface Original { + canonicalText: string; + revision: string; + text?: string; +} +interface OwnerTextState { + scope: string; + messages: ReadonlyMap; + loading: boolean; +} +const empty: OwnerTextState = { scope: '', messages: new Map(), loading: false }; +const OwnerTextContext = createContext(empty); + +export function OwnerTextProvider({ + messages, + conversationId, + isSubmitting, + children, +}: { + messages: readonly TMessage[] | null; + conversationId?: string; + isSubmitting: boolean; + children: ReactNode; +}) { + const { user } = useAuthContext(); + const selection = useMemo( + () => + JSON.stringify( + (messages ?? []) + .filter((message) => message.isCreatedByUser && message.privacyRevision) + .map((message) => [message.messageId, message.privacyRevision, message.text]) + .sort(), + ), + [messages], + ); + const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); + const [state, setState] = useState(empty); + useEffect(() => { + let cancelled = false; + const selected = JSON.parse(selection) as Array<[string, string, string]>; + if (!user?.id || !conversationId || selected.length === 0) { + setState(empty); + return; + } + setState({ scope, messages: new Map(), loading: true }); + const load = async () => { + const originals = new Map(); + try { + for (let index = 0; index < selected.length; index += 50) { + if (cancelled) { + return; + } + const batch = selected.slice(index, index + 50); + const result = await dataService.getOwnerMessageTexts( + conversationId, + batch.map(([id]) => id), + ); + for (const message of result.messages) { + if ( + batch.some( + ([id, revision, text]) => + id === message.messageId && + revision === message.revision && + text === message.canonicalText, + ) + ) { + originals.set(message.messageId, { + revision: message.revision, + text: message.text, + canonicalText: message.canonicalText, + }); + } + } + } + if (!cancelled) { + setState({ scope, messages: originals, loading: false }); + } + } catch { + if (!cancelled) { + setState({ scope, messages: new Map(), loading: false }); + } + } + }; + void load(); + return () => { + cancelled = true; + }; + }, [scope, selection, conversationId, user?.id, user?.tenantId, isSubmitting]); + const visible = state.scope === scope ? state : empty; + return {children}; +} + +/** No owner-view data is passed to edit, copy/export, retry, or prompt-building callbacks. */ +export function PrivateText({ message }: { message: TMessage }) { + const localize = useLocalize(); + const state = useContext(OwnerTextContext); + const original = state.messages.get(message.messageId); + const text = + original != null && + original.revision === message.privacyRevision && + original.canonicalText === message.text + ? original.text + : undefined; + return ( +
+
{text ?? message.text}
+

+ {localize('com_ui_private_text_hidden')} + {text == null && ( + + {' '} + ·{' '} + {localize( + state.loading ? 'com_ui_private_text_loading' : 'com_ui_private_text_unavailable', + )} + + )} +

+
+ ); +} diff --git a/client/src/components/Chat/Messages/ui/MessageRender.tsx b/client/src/components/Chat/Messages/ui/MessageRender.tsx index 37a3a9f8d07..0d5dd896842 100644 --- a/client/src/components/Chat/Messages/ui/MessageRender.tsx +++ b/client/src/components/Chat/Messages/ui/MessageRender.tsx @@ -15,6 +15,7 @@ import { getHeaderHoverLabel } from '~/components/Chat/Messages/ui/HeaderLabel'; import MessageContent from '~/components/Chat/Messages/Content/MessageContent'; import { useLocalize, useMessageActions, useContentMetadata } from '~/hooks'; import SiblingSwitch from '~/components/Chat/Messages/SiblingSwitch'; +import { PrivateText } from '~/components/Chat/Messages/PrivateText'; import HoverButtons from '~/components/Chat/Messages/HoverButtons'; import MessageRow from '~/components/Chat/Messages/ui/MessageRow'; import MessageIcon from '~/components/Chat/Messages/MessageIcon'; @@ -153,6 +154,8 @@ const MessageRender = memo(function MessageRender({ return null; } + const showOwnerText = !edit && msg.isCreatedByUser && Boolean(msg.privacyRevision); + return ( - {wakeupDisplay != null && !edit ? ( - - ) : ( - ({}))} - /> - )} + {showOwnerText && } + {!showOwnerText && + (wakeupDisplay != null && !edit ? ( + + ) : ( + ({}))} + /> + ))} ); diff --git a/client/src/components/Messages/ContentRender.tsx b/client/src/components/Messages/ContentRender.tsx index 8207d594402..6d3facb10bf 100644 --- a/client/src/components/Messages/ContentRender.tsx +++ b/client/src/components/Messages/ContentRender.tsx @@ -18,6 +18,7 @@ import { ErrorSourceProvider } from '~/components/Messages/Content/Error/source' import Elapsed, { shouldShowElapsed } from '~/components/Chat/Messages/Elapsed'; import { getHeaderHoverLabel } from '~/components/Chat/Messages/ui/HeaderLabel'; import ContentParts from '~/components/Chat/Messages/Content/ContentParts'; +import { PrivateText } from '~/components/Chat/Messages/PrivateText'; import SiblingSwitch from '~/components/Chat/Messages/SiblingSwitch'; import HoverButtons from '~/components/Chat/Messages/HoverButtons'; import MessageRow from '~/components/Chat/Messages/ui/MessageRow'; @@ -221,26 +222,30 @@ const ContentRender = memo(function ContentRender({ > - } - /> + {!edit && msg.isCreatedByUser && msg.privacyRevision ? ( + + ) : ( + } + /> + )} {/** A turn that ran out of agent steps is incomplete, not broken. Rendered diff --git a/client/src/hooks/Chat/useChatFunctions.ts b/client/src/hooks/Chat/useChatFunctions.ts index d3b370da95b..ccb14e3aed4 100644 --- a/client/src/hooks/Chat/useChatFunctions.ts +++ b/client/src/hooks/Chat/useChatFunctions.ts @@ -56,7 +56,14 @@ const STALE_SEND_REVALIDATION_MS = 5_000; const logChatRequest = (request: Record) => { logger.log('=====================================\nAsk function called with:'); - logger.dir(request); + logger.dir({ + conversationId: request.conversationId, + messageId: request.messageId, + parentMessageId: request.parentMessageId, + isEdited: request.isEdited, + isContinued: request.isContinued, + isRegenerate: request.isRegenerate, + }); logger.log('====================================='); }; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 9df188df8a1..40a040d55ef 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1,4 +1,7 @@ { + "com_ui_private_text_hidden": "Private details hidden from the model", + "com_ui_private_text_loading": "Loading your original message…", + "com_ui_private_text_unavailable": "Original unavailable. Showing the filtered message.", "com_ui_repository_instructions": "Repository instructions", "com_ui_repository_instructions_prefer": "Prefer repository conventions", "com_ui_repository_instructions_defer": "Prefer agent conventions", diff --git a/client/src/utils/messages.ts b/client/src/utils/messages.ts index 02e05d5f74e..2c59cee4a4b 100644 --- a/client/src/utils/messages.ts +++ b/client/src/utils/messages.ts @@ -852,6 +852,7 @@ export function areMessageFieldsEqual( return ( prevMsg.messageId === nextMsg.messageId && prevMsg.text === nextMsg.text && + prevMsg.privacyRevision === nextMsg.privacyRevision && prevMsg.error === nextMsg.error && prevMsg.unfinished === nextMsg.unfinished && /** Read by the row: `useGenerationsByLatest` gates the Continue button on it and diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index fbd658c5fec..2b5238b8f61 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -2976,6 +2976,22 @@ function resolveResponses({ graph, messages, text, toolNames }) { return askUserQuestionResponses(askUserQuestionLabel, toolNames); } + if (text.includes('E2E_PRIVATE_TEXT:')) { + return { + responses: [MOCK_REPLY], + resolveOnStream: (streamMessages) => { + const prompt = JSON.stringify(streamMessages); + const protectedText = + !prompt.includes('alice@example.com') && /EMAIL_1_[a-f0-9]{32}/.test(prompt); + return { + responses: [ + protectedText ? 'E2E private model input verified' : 'E2E private model input failed', + ], + }; + }, + }; + } + if (text.includes(ASSERT_AGENT_CONTEXT_MARKER)) { return { responses: [MOCK_REPLY], diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts new file mode 100644 index 00000000000..214d79665d9 --- /dev/null +++ b/e2e/specs/mock/private-text.spec.ts @@ -0,0 +1,121 @@ +import { expect, test } from '@playwright/test'; +import type { TMessage } from 'librechat-data-provider'; +import { withMongo } from './db'; +import { + loginAdmin, + setRuntimeFilters, + restoreRuntimeFilters, + requestResult, +} from './content-filters.helpers'; +import { + MOCK_ENDPOINTS, + selectMockEndpoint, + sendMessageAndWaitForCompletion, + messagesView, + fetchJson, +} from './helpers'; + +const original = 'E2E_PRIVATE_TEXT: alice@example.com'; + +test('owner sees original after reload while provider, sharing, and canonical reads stay filtered', async ({ + page, + request, +}, testInfo) => { + test.setTimeout(120000); + const token = await loginAdmin(request); + await setRuntimeFilters(request, token, { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }); + let conversationId: string | undefined; + try { + await page.goto('/c/new'); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const response = await sendMessageAndWaitForCompletion(page, original); + conversationId = (await response.json()).conversationId as string; + expect(conversationId).toBeTruthy(); + await expect( + messagesView(page).getByText('E2E private model input verified', { exact: true }), + ).toBeVisible(); + await expect(messagesView(page).getByText(original, { exact: true })).toBeVisible(); + await expect( + messagesView(page).getByText('Private details hidden from the model', { exact: true }), + ).toBeVisible(); + + const replay = await requestResult(request, { + path: new URL(response.url()).pathname, + token, + method: 'POST', + data: response.request().postDataJSON(), + }); + expect(replay.ok).toBe(true); + expect((replay.body as { conversationId: string }).conversationId).toBe(conversationId); + + const canonical = await fetchJson(page, `/api/messages/${conversationId}`, token); + expect(JSON.stringify(canonical)).not.toContain('alice@example.com'); + expect(JSON.stringify(canonical)).not.toContain('privateText'); + const user = canonical.find((message) => message.isCreatedByUser)!; + expect(user.text).toMatch(/\[EMAIL_1_[a-f0-9]{32}\]/); + await withMongo(async (db) => { + const row = await db + .collection('messages') + .findOne({ conversationId, messageId: user.messageId }); + expect(row?.privateText).toMatch(/^v1:/); + expect(JSON.stringify(row)).not.toContain('alice@example.com'); + }); + + await page.reload(); + await expect(messagesView(page).getByText(original, { exact: true })).toBeVisible(); + for (const theme of ['light', 'dark']) { + await page.evaluate( + (dark) => document.documentElement.classList.toggle('dark', dark), + theme === 'dark', + ); + await page.screenshot({ + path: testInfo.outputPath(`owner-text-${theme}.png`), + fullPage: true, + }); + } + const share = await requestResult(request, { + path: `/api/share/${conversationId}`, + token, + method: 'POST', + data: {}, + }); + expect(share.ok).toBe(true); + const shared = await requestResult(request, { + path: `/api/share/${(share.body as { shareId: string }).shareId}`, + token, + }); + expect(shared.ok).toBe(true); + expect(shared.text).not.toContain('alice@example.com'); + expect(shared.text).not.toContain('privateText'); + expect(shared.text).toContain('EMAIL_1_'); + + const unauthorized = await request.post(`/api/messages/${conversationId}/owner-text`, { + data: { messageIds: [user.messageId] }, + }); + expect(unauthorized.status()).toBe(401); + } finally { + await restoreRuntimeFilters(request, token); + if (conversationId) { + await requestResult(request, { + path: '/api/convos', + token, + method: 'DELETE', + data: { arg: { conversationId } }, + }); + await withMongo(async (db) => { + expect(await db.collection('messages').countDocuments({ conversationId })).toBe(0); + }); + } + } +}); diff --git a/librechat.example.yaml b/librechat.example.yaml index 3f25483564b..b24e913d37a 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -1390,8 +1390,12 @@ endpoints: # messages: # unattributedAssistantContent: model_output # `model_output` (default) or `inspect` # pii: -# action: audit # `block` (default), `audit`, or `redact` (reserved for transformation-aware ingress; current routes still block) -# # `redact` requires the later encrypted owner-view and canonical-message integration. +# action: audit # `block` (default), `audit`, or opt-in `redact` +# # `redact` transforms fresh, text-only interactive Agent chat submissions with a clientRequestId. +# # Edits, resumes, files, quotes, other endpoints, and other fields still use blocking checks. +# # Requires a valid CREDS_KEY. Preserve that key for owner-original recovery; changing it +# # makes existing originals unavailable. Text remains canonical/filtered in ordinary exports. +# # Originals are an encrypted, unindexed sidecar field, deleted/expired with the message. # # maxCharacters: 65536 # bounded total inspected text per transformation session # # maxMatches: 256 # bounded matches per transformation session # # Custom patterns may set category: email, phone, name, credential, or custom. diff --git a/packages/api/src/protection/index.ts b/packages/api/src/protection/index.ts index 264f49b381b..bfc6d02532b 100644 --- a/packages/api/src/protection/index.ts +++ b/packages/api/src/protection/index.ts @@ -10,3 +10,5 @@ export * from './adapters/chat'; export * from './adapters/nested'; export * from './adapters/messages'; export * from './adapters/submissions'; +export * from './private/submission'; +export * from './private/view'; diff --git a/packages/api/src/protection/private/crypto.ts b/packages/api/src/protection/private/crypto.ts new file mode 100644 index 00000000000..20aa47451cf --- /dev/null +++ b/packages/api/src/protection/private/crypto.ts @@ -0,0 +1,53 @@ +import { createCipheriv, createDecipheriv, createHmac, hkdfSync, randomBytes } from 'node:crypto'; + +export interface PrivateTextCipher { + seal(text: string, binding: readonly string[]): string; + open(envelope: string, binding: readonly string[]): string; + revision(value: readonly string[]): string; +} + +/** Uses the deployment credential key, domain-separated from credential encryption. */ +export function createPrivateTextCipher(hexKey: string): PrivateTextCipher { + if (!/^[a-fA-F0-9]{64}$/.test(hexKey)) { + throw new Error('Private message encryption is unavailable.'); + } + const key = Buffer.from( + hkdfSync('sha256', Buffer.from(hexKey, 'hex'), '', 'librechat-owner-text-v1', 32), + ); + const aad = (binding: readonly string[]) => + Buffer.from(JSON.stringify(['owner-text-v1', ...binding])); + return { + revision(value) { + return createHmac('sha256', key).update(aad(value)).digest('hex').slice(0, 32); + }, + seal(text, binding) { + const nonce = randomBytes(12); + const cipher = createCipheriv('aes-256-gcm', key, nonce); + cipher.setAAD(aad(binding)); + const ciphertext = Buffer.concat([cipher.update(text, 'utf8'), cipher.final()]); + return [ + 'v1', + nonce.toString('base64'), + cipher.getAuthTag().toString('base64'), + ciphertext.toString('base64'), + ].join(':'); + }, + open(envelope, binding) { + try { + const [version, nonce, tag, ciphertext, extra] = envelope.split(':'); + if (version !== 'v1' || extra != null || !nonce || !tag || ciphertext == null) { + throw new Error(); + } + const decipher = createDecipheriv('aes-256-gcm', key, Buffer.from(nonce, 'base64')); + decipher.setAAD(aad(binding)); + decipher.setAuthTag(Buffer.from(tag, 'base64')); + return Buffer.concat([ + decipher.update(Buffer.from(ciphertext, 'base64')), + decipher.final(), + ]).toString('utf8'); + } catch { + throw new Error('Private message text is unavailable.'); + } + }, + }; +} diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts new file mode 100644 index 00000000000..44bda5f0eae --- /dev/null +++ b/packages/api/src/protection/private/submission.spec.ts @@ -0,0 +1,219 @@ +import type { MessageMethods, IMessage } from '@librechat/data-schemas'; +import type { FiltersConfig } from 'librechat-data-provider'; +import type { Request, Response } from 'express'; +import { + createPrivateTextIngress, + savePrivateTextMessage, + stampPrivateTextMessage, + requirePrivateTextPersistence, + privateTextBinding, +} from './submission'; +import { createPrivateTextCipher } from './crypto'; +import { createPrivateTextView } from './view'; + +jest.mock('@librechat/data-schemas', () => ({ + logger: { warn: jest.fn(), info: jest.fn(), error: jest.fn(), debug: jest.fn() }, +})); + +const key = 'ab'.repeat(32); +const original = 'Email alice@example.com'; +const filters: FiltersConfig = { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: '[a-z]+@[a-z]+\\.[a-z]+', category: 'email' }, + ], + }, + }, +}; +function submit(overrides: object = {}, encryptionKey = key) { + const req = { + path: '/', + user: { id: 'owner', tenantId: 'tenant-a' }, + body: { text: original, clientRequestId: 'request-1', ...overrides }, + } as unknown as Request; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn(), setHeader: jest.fn() }; + const next = jest.fn(); + createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => encryptionKey, + })(req, res as unknown as Response, next); + const message = stampPrivateTextMessage(req, { + messageId: 'message-1', + conversationId: 'conversation-1', + isCreatedByUser: true, + text: req.body.text, + }); + return { req, res, next, message }; +} + +describe('private text submission boundary', () => { + it('replaces request text before consumers and exposes no original in metadata or serialization', () => { + const { req, message, next } = submit(); + expect(next).toHaveBeenCalledTimes(1); + expect(req.body.text).toMatch(/^Email \[EMAIL_1_[a-f0-9]{32}\]$/); + expect(message).toHaveProperty('privacyRevision'); + expect(JSON.stringify({ req, message })).not.toContain('alice@example.com'); + }); + + it('uses stable retry revisions and distinct namespaces for different turns or originals', () => { + const first = submit().message; + expect(submit().message).toEqual(first); + expect(submit({ clientRequestId: 'request-2' }).message.text).not.toBe(first.text); + expect(submit({ text: 'Email bob@example.com' }).message.text).not.toBe(first.text); + }); + + it.each([ + { files: [{}] }, + { quotes: ['quote'] }, + { isRegenerate: true }, + { isEdited: true }, + { isContinued: true }, + { editedContent: {} }, + { recoverySteerId: 'steer' }, + ])('leaves unsupported submissions to the existing blocking inspector: %j', (extra) => { + const { req, next, message } = submit(extra); + expect(req.body.text).toBe(original); + expect(message).not.toHaveProperty('privacyRevision'); + expect(next).toHaveBeenCalledTimes(1); + }); + + it('fails closed without a valid key and does not echo matched text', () => { + const { next, res } = submit({}, ''); + expect(next).not.toHaveBeenCalled(); + expect(res.status).toHaveBeenCalledWith(400); + expect(JSON.stringify(res.json.mock.calls)).not.toContain(original); + }); + + it('commits canonical text and ciphertext together and binds the owner view to final identity', async () => { + const { req, message } = submit(); + let envelope = ''; + const save: MessageMethods['saveMessage'] = jest.fn(async (_ctx, value, metadata) => { + envelope = metadata?.privateText?.envelope ?? ''; + expect(JSON.stringify({ value, metadata })).not.toContain(original); + return { ...value, privacyRevision: metadata?.privateText?.revision } as IMessage; + }); + const stored = await savePrivateTextMessage(save, req, { userId: 'owner' }, message); + expect(save).toHaveBeenCalledTimes(1); + const cipher = createPrivateTextCipher(key); + const binding = privateTextBinding('owner', 'tenant-a', stored!); + expect(cipher.open(envelope, binding)).toBe(original); + for (let index = 0; index < binding.length; index++) { + const tampered = [...binding]; + tampered[index] += '-changed'; + expect(() => cipher.open(envelope, tampered)).toThrow('Private message text is unavailable.'); + } + expect(() => createPrivateTextCipher('cd'.repeat(32)).open(envelope, binding)).toThrow(); + expect(() => cipher.open(envelope.slice(0, -4) + 'abcd', binding)).toThrow(); + await expect(savePrivateTextMessage(save, req, { userId: 'other' }, message)).rejects.toThrow(); + }); + + it('does not release main provider admission until persistence finishes', async () => { + const { req, message } = submit(); + let finish!: (result: { message: typeof message }) => void; + const pending = new Promise<{ message: typeof message }>((resolve) => { + finish = resolve; + }); + const provider = jest.fn(); + const started = requirePrivateTextPersistence(req, () => pending).then(provider); + await Promise.resolve(); + expect(provider).not.toHaveBeenCalled(); + finish({ message }); + await started; + expect(provider).toHaveBeenCalledTimes(1); + await expect(requirePrivateTextPersistence(req, async () => ({}))).rejects.toThrow(); + await expect( + requirePrivateTextPersistence(req, async () => { + throw new Error('write failed'); + }), + ).rejects.toThrow(); + }); + + it('rejects stale or swallowed persistence results, including a duplicate ID with different text', async () => { + const { req, message } = submit(); + const save: MessageMethods['saveMessage'] = jest.fn(async () => undefined); + await expect(savePrivateTextMessage(save, req, { userId: 'owner' }, message)).rejects.toThrow(); + await expect( + requirePrivateTextPersistence(req, async () => ({ message: { ...message, text: 'stale' } })), + ).rejects.toThrow(); + }); + + it('returns original text only from the authenticated private view, with no-store headers', async () => { + const { req, message } = submit(); + const cipher = createPrivateTextCipher(key); + const row = { + ...message, + privacyRevision: message.privacyRevision!, + privateText: cipher.seal(original, privateTextBinding('owner', 'tenant-a', message)), + }; + const read = jest.fn(async () => [row]); + const handler = createPrivateTextView({ read, getKey: () => key }); + req.params = { conversationId: 'conversation-1' }; + req.body = { messageIds: ['message-1'] }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn(), setHeader: jest.fn() }; + await handler(req, res as unknown as Response, jest.fn()); + expect(read).toHaveBeenCalledWith({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId: 'conversation-1', + messageIds: ['message-1'], + }); + expect(res.setHeader).toHaveBeenCalledWith('Cache-Control', 'private, no-store'); + expect(res.json).toHaveBeenCalledWith({ + messages: [ + { + messageId: 'message-1', + revision: row.privacyRevision, + canonicalText: row.text, + text: original, + }, + ], + }); + req.body.messageIds = Array(51).fill('message-1'); + await handler(req, res as unknown as Response, jest.fn()); + expect(res.status).toHaveBeenLastCalledWith(400); + expect(read).toHaveBeenCalledTimes(1); + await handler( + { params: req.params, body: { messageIds: ['message-1'] } } as Request, + res as unknown as Response, + jest.fn(), + ); + expect(res.status).toHaveBeenLastCalledWith(401); + expect(read).toHaveBeenCalledTimes(1); + }); + + it('returns no original when a stored row has changed canonical text or the key is unavailable', async () => { + const { message, req } = submit(); + const cipher = createPrivateTextCipher(key); + const privateText = cipher.seal(original, privateTextBinding('owner', 'tenant-a', message)); + const read = jest.fn(async () => [ + { + ...message, + text: 'edited canonical', + privacyRevision: message.privacyRevision!, + privateText, + }, + ]); + req.params = { conversationId: 'conversation-1' }; + req.body = { messageIds: ['message-1'] }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn(), setHeader: jest.fn() }; + await createPrivateTextView({ read, getKey: () => key })( + req, + res as unknown as Response, + jest.fn(), + ); + expect(JSON.stringify(res.json.mock.calls)).not.toContain(original); + expect(res.status).toHaveBeenLastCalledWith(200); + await createPrivateTextView({ read, getKey: () => '' })( + req, + res as unknown as Response, + jest.fn(), + ); + expect(res.status).toHaveBeenLastCalledWith(503); + expect(JSON.stringify(res.json.mock.calls)).not.toContain(original); + }); +}); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts new file mode 100644 index 00000000000..98c0bb8fdda --- /dev/null +++ b/packages/api/src/protection/private/submission.ts @@ -0,0 +1,238 @@ +import type { FiltersConfig, MessageFilterPiiConfig } from 'librechat-data-provider'; +import type { MessageMethods } from '@librechat/data-schemas'; +import type { RequestHandler, Request } from 'express'; +import type { PrivateTextCipher } from './crypto'; +import { ContentFilterError } from '../../middleware/contentFilter'; +import { createPiiTextTransformer } from '../transform'; +import { createPrivateTextCipher } from './crypto'; +import { inspectContent } from '../runtime'; + +interface PrivateTextMessage { + messageId?: string; + conversationId?: string | null; + isCreatedByUser?: boolean; + text?: string; + privacyRevision?: string; +} + +interface Capture { + readonly userId: string; + readonly tenantId: string; + readonly revision: string; + readonly text: string; + readonly envelope: string; + readonly cipher: PrivateTextCipher; +} + +const captures = new WeakMap(); + +function unavailable(): ContentFilterError { + return new ContentFilterError({ + detectorId: 'pii-pattern', + ruleId: 'private-text', + label: 'private value that could not be protected', + source: 'message', + field: 'text', + provenance: 'user', + fragmentId: 'chat.text', + fragmentPath: '/text', + }); +} + +export function privateTextBinding( + userId: string, + tenantId: string, + message: PrivateTextMessage, +): string[] { + return [ + userId, + tenantId, + message.conversationId ?? '', + message.messageId ?? '', + message.privacyRevision ?? '', + message.text ?? '', + ]; +} + +/** Installed only on the authenticated interactive Agent chat router, before any content consumer. */ +export function createPrivateTextIngress(options: { + getFilters(req: Request): FiltersConfig | undefined; + getLegacyPii(req: Request): MessageFilterPiiConfig | undefined; + getKey(): string; +}): RequestHandler { + return (req, res, next) => { + const rule = options.getFilters(req)?.messages?.pii; + if (rule?.action !== 'redact' || typeof req.body?.text !== 'string') { + next(); + return; + } + const body = req.body as { + text: string; + clientRequestId?: string; + files?: object[]; + quotes?: string[]; + isEdited?: boolean; + isContinued?: boolean; + isRegenerate?: boolean; + compact?: boolean; + overrideParentMessageId?: string; + overrideConvoId?: string; + addedConvo?: boolean; + editedContent?: unknown; + recoverySteerId?: string; + responseMessageId?: string; + }; + const request = req as Request & { + user?: { id?: string; tenantId?: string | null }; + _isAgentTrigger?: boolean; + }; + if ( + req.path === '/resume' || + request._isAgentTrigger === true || + body.isEdited || + body.isContinued || + body.isRegenerate || + body.compact || + body.editedContent != null || + body.recoverySteerId != null || + (typeof body.clientRequestId === 'string' && + body.clientRequestId.startsWith('steer-recovery:')) || + body.overrideParentMessageId || + body.overrideConvoId || + body.addedConvo || + body.files?.length || + body.quotes?.length + ) { + next(); + return; + } + try { + const fragment = { + id: 'chat.text', + path: '/text', + text: body.text, + source: 'message', + field: 'text', + format: 'plain', + treatment: 'replaceable', + provenance: 'user', + } as const; + const legacy = inspectContent([fragment], { legacyPii: options.getLegacyPii(req) }); + if (legacy != null) { + throw new ContentFilterError(legacy); + } + const result = createPiiTextTransformer(rule).createSession().transform(fragment); + if (result.replacements === 0) { + next(); + return; + } + const userId = request.user?.id; + const tenantId = request.user?.tenantId ?? ''; + if ( + !userId || + typeof body.clientRequestId !== 'string' || + body.clientRequestId.length > 256 + ) { + throw unavailable(); + } + const cipher = createPrivateTextCipher(options.getKey()); + const revision = cipher.revision([userId, tenantId, body.clientRequestId, body.text]); + // A keyed, turn-specific namespace prevents unrelated historical placeholders aliasing. + const text = result.content.replace( + /\[(EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_(\d+)\]/g, + (marker, category: string, index: string) => + body.text.includes(marker) ? marker : `[${category}_${index}_${revision}]`, + ); + const envelope = cipher.seal(body.text, [userId, tenantId, revision]); + captures.set(req, { userId, tenantId, revision, text, envelope, cipher }); + body.text = text; + next(); + } catch { + res.status(400).json({ + error: 'content_filter_block', + message: 'Private details could not be protected. Nothing was sent to the model.', + }); + } + }; +} + +/** Only safe metadata joins the user-message projection sent to events and ordinary readers. */ +export function stampPrivateTextMessage( + req: object | undefined, + message: T, +): T & { privacyRevision?: string } { + const capture = req == null ? undefined : captures.get(req); + if (capture != null && message.isCreatedByUser === true && message.text === capture.text) { + message.privacyRevision = capture.revision; + } + return message; +} + +/** Encrypts against final server-resolved message identity, then commits both views in one write. */ +export async function savePrivateTextMessage( + save: MessageMethods['saveMessage'], + req: object | undefined, + ...args: Parameters +): ReturnType { + const [ctx, message, metadata] = args; + const capture = req == null ? undefined : captures.get(req); + if (capture == null || message.isCreatedByUser !== true) { + return save(...args); + } + if (message.text !== capture.text) { + throw unavailable(); + } + if ( + ctx.userId !== capture.userId || + !message.messageId || + !message.conversationId || + message.newMessageId + ) { + throw unavailable(); + } + const revision = capture.revision; + const original = capture.cipher.open(capture.envelope, [ + capture.userId, + capture.tenantId, + revision, + ]); + const envelope = capture.cipher.seal( + original, + privateTextBinding(capture.userId, capture.tenantId, { ...message, privacyRevision: revision }), + ); + const saved = await save( + ctx, + { ...message, tenantId: capture.tenantId || undefined }, + { + ...metadata, + privateText: { envelope, revision }, + }, + ); + if ( + saved?.privacyRevision !== revision || + saved.text !== capture.text || + saved.messageId !== message.messageId || + saved.conversationId !== message.conversationId + ) { + throw unavailable(); + } + return saved; +} + +/** Must complete before sendCompletion. Existing cancellation/deletion still owns the run. */ +export async function requirePrivateTextPersistence( + req: object | undefined, + start: () => Promise<{ message?: PrivateTextMessage | null } | undefined>, +): Promise { + const capture = req == null ? undefined : captures.get(req); + if (capture == null) { + return; + } + const result = await start(); + if ( + result?.message?.privacyRevision !== capture.revision || + result.message.text !== capture.text + ) { + throw unavailable(); + } +} diff --git a/packages/api/src/protection/private/view.ts b/packages/api/src/protection/private/view.ts new file mode 100644 index 00000000000..ebfdef77f0f --- /dev/null +++ b/packages/api/src/protection/private/view.ts @@ -0,0 +1,65 @@ +import { z } from 'zod'; +import type { MessageMethods } from '@librechat/data-schemas'; +import type { RequestHandler, Request } from 'express'; +import { createPrivateTextCipher } from './crypto'; +import { privateTextBinding } from './submission'; + +const inputSchema = z + .object({ + messageIds: z.array(z.string().min(1).max(256)).min(1).max(50), + }) + .strict(); + +export function createPrivateTextView(options: { + read: MessageMethods['getPrivateMessageTexts']; + getKey(): string; +}): RequestHandler { + return async (req, res) => { + res.setHeader('Cache-Control', 'private, no-store'); + const request = req as Request & { user?: { id?: string; tenantId?: string | null } }; + const userId = request.user?.id; + if (!userId) { + res.status(401).json({ error: 'Unauthorized' }); + return; + } + const input = inputSchema.safeParse(req.body); + const conversationId = req.params.conversationId; + if (!input.success || typeof conversationId !== 'string') { + res.status(400).json({ error: 'Invalid private message request' }); + return; + } + try { + const tenantId = request.user?.tenantId ?? ''; + const rows = await options.read({ + userId, + tenantId: tenantId || undefined, + conversationId, + messageIds: input.data.messageIds, + }); + const cipher = createPrivateTextCipher(options.getKey()); + const messages = rows.map((row) => { + try { + const text = cipher.open( + row.privateText, + privateTextBinding(userId, tenantId, { ...row, conversationId }), + ); + return { + messageId: row.messageId, + revision: row.privacyRevision, + canonicalText: row.text, + text, + }; + } catch { + return { + messageId: row.messageId, + revision: row.privacyRevision, + canonicalText: row.text, + }; + } + }); + res.status(200).json({ messages }); + } catch { + res.status(503).json({ error: 'Private message text is unavailable.' }); + } + }; +} diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index cf5ade19a1c..d972978c06b 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -1094,6 +1094,21 @@ export const branchMessage = async ( return request.post(endpoints.messagesBranch(), payload); }; +export interface OwnerMessageText { + canonicalText: string; + messageId: string; + revision: string; + text?: string; +} + +/** Private display data; never merge into ordinary message/query-cache objects. */ +export function getOwnerMessageTexts( + conversationId: string, + messageIds: string[], +): Promise<{ messages: OwnerMessageText[] }> { + return request.post(`${endpoints.messages({ conversationId })}/owner-text`, { messageIds }); +} + export function getMessagesByConvoId(conversationId: string): Promise { if ( conversationId === config.Constants.NEW_CONVO || diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index a43dfab2b11..d8379852bb9 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -898,6 +898,8 @@ export const tMessageSchema = z.object({ /** @deprecated */ generation: z.string().nullable().optional(), isCreatedByUser: z.boolean(), + /** Opaque revision of the separately authorized owner display. */ + privacyRevision: z.string().optional(), /** True when the complete stored row came from outside the model. */ isUserSubmitted: z.boolean().optional(), /** JSON pointers to caller-authored fields in an otherwise mixed model response. */ diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts new file mode 100644 index 00000000000..57f6ce60538 --- /dev/null +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -0,0 +1,160 @@ +import mongoose from 'mongoose'; +import { v4 as uuid } from 'uuid'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import { createMessageMethods, CLIENT_MESSAGE_SELECT } from './message'; +import { tenantStorage, runAsSystem } from '~/config/tenantContext'; +import { createModels } from '../models'; + +jest.mock('~/config/winston', () => ({ + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), +})); + +let server: MongoMemoryServer; +const methods = createMessageMethods(mongoose); +const tenant = (id: string, fn: () => T) => tenantStorage.run({ tenantId: id }, fn); + +beforeAll(async () => { + server = await MongoMemoryServer.create(); + Object.assign(mongoose.models, createModels(mongoose)); + await mongoose.connect(server.getUri()); +}); +afterAll(async () => { + await mongoose.disconnect(); + await server?.stop(); +}); +afterEach(async () => { + await runAsSystem(() => mongoose.models.Message.deleteMany({})); +}); + +it('stores both views atomically and excludes ciphertext from ordinary and client reads', async () => { + const conversationId = uuid(); + const messageId = uuid(); + await tenant('tenant-a', async () => { + const saved = await methods.saveMessage( + { userId: 'owner' }, + { + messageId, + conversationId, + text: '[EMAIL_1_turn]', + isCreatedByUser: true, + }, + { privateText: { envelope: 'v1:ciphertext', revision: 'turn' } }, + ); + expect(saved?.privacyRevision).toBe('turn'); + expect(saved).not.toHaveProperty('privateText'); + for (const projection of [undefined, CLIENT_MESSAGE_SELECT]) { + const rows = await methods.getMessages({ conversationId, user: 'owner' }, projection); + expect(rows[0].text).toBe('[EMAIL_1_turn]'); + expect(rows[0]).not.toHaveProperty('privateText'); + } + const own = await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }); + expect(own[0]).toMatchObject({ + privateText: 'v1:ciphertext', + privacyRevision: 'turn', + text: '[EMAIL_1_turn]', + }); + expect( + await methods.getPrivateMessageTexts({ + userId: 'other', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + expect( + await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-b', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + }); + expect( + await tenant('tenant-b', () => + methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-b', + conversationId, + messageIds: [messageId], + }), + ), + ).toEqual([]); + await tenant('tenant-a', async () => { + await methods.deleteMessages({ conversationId, user: 'owner' }); + expect( + await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + }); +}); + +it('does not accept sidecar writes from message parameters or generic edits', async () => { + await tenant('tenant-a', async () => { + const messageId = uuid(); + const conversationId = uuid(); + const saved = await methods.saveMessage( + { userId: 'owner' }, + { + messageId, + conversationId, + isCreatedByUser: true, + text: 'clean', + privateText: 'untrusted', + privacyRevision: 'untrusted', + }, + ); + expect(saved).not.toHaveProperty('privateText'); + expect(saved).not.toHaveProperty('privacyRevision'); + await methods.updateMessage('owner', { + messageId, + privateText: 'forged', + privacyRevision: 'forged', + }); + expect( + await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + }); +}); + +it('does not return expired originals even before the TTL sweeper runs', async () => { + await tenant('tenant-a', async () => { + const messageId = uuid(); + const conversationId = uuid(); + await methods.saveMessage( + { userId: 'owner', expiredAt: new Date(0), isTemporary: true }, + { + messageId, + conversationId, + text: '[EMAIL_1]', + isCreatedByUser: true, + }, + { privateText: { envelope: 'v1:ciphertext', revision: 'turn' } }, + ); + expect( + await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + }); +}); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 8b85d82d61d..658bbb98eeb 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -11,6 +11,7 @@ import type { AppConfig, IConversation, IMessage } from '~/types'; import { createChatExpirationDate, createTempChatExpirationDate } from '~/utils/tempChatRetention'; import { activeExpirationFilter, createFallbackRetentionDate } from '~/utils/retention'; import { tenantSafeBulkWrite } from '~/utils/tenantBulkWrite'; +import { tenantStorage } from '~/config/tenantContext'; import logger from '~/config/winston'; /** Simple UUID v4 regex to replace zod validation */ @@ -462,6 +463,7 @@ export const CLIENT_MESSAGE_SELECT: string = [ '-conversationSignature', '-summary', '-summaryTokenCount', + '-privateText', '-contextMeta', '-langfuseSampled', '-langfuseDestinationIds', @@ -645,7 +647,25 @@ function toSettledAt(value: unknown): Date | undefined { return undefined; } +export interface PrivateTextWrite { + readonly envelope: string; + readonly revision: string; +} + +export interface PrivateTextRead { + readonly messageId: string; + readonly text: string; + readonly privacyRevision: string; + readonly privateText: string; +} + export interface MessageMethods { + getPrivateMessageTexts(input: { + userId: string; + tenantId?: string; + conversationId: string; + messageIds: readonly string[]; + }): Promise; saveMessage( ctx: { userId: string; @@ -657,7 +677,7 @@ export interface MessageMethods { newMessageId?: string; contextMeta?: IMessage['contextMeta'] | null; }, - metadata?: { context?: string }, + metadata?: { context?: string; privateText?: PrivateTextWrite }, ): Promise; /** * Reads the references a trace viewer needs for one of the user's @@ -881,7 +901,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa /** `null` unsets a previously stored value; omission leaves it in place. */ contextMeta?: IMessage['contextMeta'] | null; }, - metadata?: { context?: string }, + metadata?: { context?: string; privateText?: PrivateTextWrite }, ) { if (!userId) { throw new Error('User not authenticated'); @@ -902,6 +922,15 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa user: userId, messageId: params.newMessageId || params.messageId, }; + delete update.privateText; + delete update.privacyRevision; + if (metadata?.privateText != null) { + if (params.isCreatedByUser !== true || typeof params.text !== 'string') { + throw new Error('Private text requires a user message.'); + } + update.privateText = metadata.privateText.envelope; + update.privacyRevision = metadata.privateText.revision; + } delete update.isTemporary; delete update.expiredAt; let retentionOnInsert: { expiredAt: Date; isTemporary: false } | undefined; @@ -1966,6 +1995,8 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa try { const Message = mongoose.models.Message as Model; const { messageId, ...update } = message; + delete update.privateText; + delete update.privacyRevision; const submittedPaths = normalizeUserSubmittedPaths(update.userSubmittedPaths); const submittedMessageFields = normalizeUserSubmittedMessageFieldPaths( update.userSubmittedMessageFieldPaths, @@ -3655,7 +3686,36 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa return Message.meiliSearch(query, searchOptions, hydrate); } + async function getPrivateMessageTexts(input: { + userId: string; + tenantId?: string; + conversationId: string; + messageIds: readonly string[]; + }): Promise { + if (!input.userId || !UUID_REGEX.test(input.conversationId) || input.messageIds.length > 50) { + throw new Error('Invalid private message read.'); + } + const activeTenant = tenantStorage.getStore()?.tenantId; + if (activeTenant != null && activeTenant !== input.tenantId) { + return []; + } + const Message = mongoose.models.Message as Model; + return Message.find({ + user: input.userId, + ...traceTenantScope(input.tenantId), + conversationId: input.conversationId, + messageId: { $in: input.messageIds }, + isCreatedByUser: true, + privateText: { $exists: true }, + $or: [{ expiredAt: null }, { expiredAt: { $gt: new Date() } }], + }) + .select('messageId text privacyRevision +privateText -_id') + .limit(50) + .lean(); + } + return { + getPrivateMessageTexts, saveMessage, bulkSaveMessages, recordMessage, diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index cb88fc01b53..88a5010d9a1 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -59,6 +59,8 @@ const messageSchema: Schema = new Schema( summary: { type: String, }, + privateText: { type: String, select: false }, + privacyRevision: { type: String }, isCreatedByUser: { type: Boolean, required: true, diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index d2628b61104..76d1f54fce6 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -59,6 +59,9 @@ export interface IMessage extends Document { sender?: string; text?: string; summary?: string; + /** Authenticated ciphertext, available only through the owner-view read. */ + privateText?: string; + privacyRevision?: string; isCreatedByUser: boolean; /** True when the complete stored row came from outside the model. */ isUserSubmitted?: boolean; From 5634f3fa082f141a32ccabd53db81cc071ab3d5a Mon Sep 17 00:00:00 2001 From: Lia Date: Thu, 24 Sep 2026 18:00:20 +0000 Subject: [PATCH 02/18] =?UTF-8?q?=F0=9F=A7=B5=20fix:=20Keep=20Ordinary=20C?= =?UTF-8?q?hat=20Views=20Outside=20Private=20Text=20Provider?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../__tests__/messages-content-edit.spec.js | 1 + .../routes/__tests__/messages-delete.spec.js | 1 + .../__tests__/messages-feedback.spec.js | 1 + .../messages-get-real-validation.spec.js | 1 + .../routes/__tests__/messages-get.spec.js | 1 + .../messages-subagent-thread.spec.js | 1 + .../__tests__/chat.contentFilter.spec.js | 1 + .../Chat/Messages/PrivateText.spec.tsx | 11 +++++++++ .../components/Chat/Messages/PrivateText.tsx | 23 +++++++++++++------ 9 files changed, 34 insertions(+), 7 deletions(-) diff --git a/api/server/routes/__tests__/messages-content-edit.spec.js b/api/server/routes/__tests__/messages-content-edit.spec.js index 9ed509aa45c..ce0ffbb710a 100644 --- a/api/server/routes/__tests__/messages-content-edit.spec.js +++ b/api/server/routes/__tests__/messages-content-edit.spec.js @@ -9,6 +9,7 @@ jest.mock('@librechat/agents', () => ({ jest.mock('@librechat/api', () => ({ unescapeLaTeX: jest.fn((value) => value), countTokens: jest.fn().mockResolvedValue(2), + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (_req, _res, next) => next()), sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), diff --git a/api/server/routes/__tests__/messages-delete.spec.js b/api/server/routes/__tests__/messages-delete.spec.js index a0c27b62b9a..c8174256196 100644 --- a/api/server/routes/__tests__/messages-delete.spec.js +++ b/api/server/routes/__tests__/messages-delete.spec.js @@ -9,6 +9,7 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (req, res, next) => next()), inspectContent: jest.fn(() => null), extractFeedbackContent: jest.fn(() => []), diff --git a/api/server/routes/__tests__/messages-feedback.spec.js b/api/server/routes/__tests__/messages-feedback.spec.js index 46d53a3ce06..b3862ebc283 100644 --- a/api/server/routes/__tests__/messages-feedback.spec.js +++ b/api/server/routes/__tests__/messages-feedback.spec.js @@ -8,6 +8,7 @@ jest.mock('@librechat/agents', () => ({ jest.mock('@librechat/api', () => ({ unescapeLaTeX: jest.fn((value) => value), countTokens: jest.fn().mockResolvedValue(10), + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (req, res, next) => next()), sendFeedbackScore: jest.fn().mockResolvedValue(undefined), traceIdForMessage: jest.fn((messageId) => `trace-${messageId}`), diff --git a/api/server/routes/__tests__/messages-get-real-validation.spec.js b/api/server/routes/__tests__/messages-get-real-validation.spec.js index 46daa50fc00..a7d7e4a4d95 100644 --- a/api/server/routes/__tests__/messages-get-real-validation.spec.js +++ b/api/server/routes/__tests__/messages-get-real-validation.spec.js @@ -21,6 +21,7 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (req, res, next) => next()), inspectContent: jest.fn(() => null), extractFeedbackContent: jest.fn(() => []), diff --git a/api/server/routes/__tests__/messages-get.spec.js b/api/server/routes/__tests__/messages-get.spec.js index b72c0973f48..f376a9f0dce 100644 --- a/api/server/routes/__tests__/messages-get.spec.js +++ b/api/server/routes/__tests__/messages-get.spec.js @@ -54,6 +54,7 @@ jest.mock('@librechat/api', () => { /** The real helper, without loading the rest of the package this suite mocks around. */ withoutTraceRefs: jest.requireActual('../../../../packages/api/src/langfuse/trace.ts') .withoutTraceRefs, + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (req, res, next) => next()), inspectContent, extractChatContent, diff --git a/api/server/routes/__tests__/messages-subagent-thread.spec.js b/api/server/routes/__tests__/messages-subagent-thread.spec.js index 3c7ee69e759..574a1435e3c 100644 --- a/api/server/routes/__tests__/messages-subagent-thread.spec.js +++ b/api/server/routes/__tests__/messages-subagent-thread.spec.js @@ -7,6 +7,7 @@ jest.mock('@librechat/agents', () => ({ sleep: jest.fn() })); jest.mock('@librechat/api', () => ({ withoutTraceRefs: jest.fn((message) => message), + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), createContentFilter: jest.fn(() => (_req, _res, next) => next()), unescapeLaTeX: jest.fn((value) => value), countTokens: jest.fn().mockResolvedValue(1), diff --git a/api/server/routes/agents/__tests__/chat.contentFilter.spec.js b/api/server/routes/agents/__tests__/chat.contentFilter.spec.js index b6439ee18dc..d77e8744d59 100644 --- a/api/server/routes/agents/__tests__/chat.contentFilter.spec.js +++ b/api/server/routes/agents/__tests__/chat.contentFilter.spec.js @@ -14,6 +14,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ createMessageFilterPii: mockCreateMessageFilterPii, + createPrivateTextIngress: jest.fn(() => (_req, _res, next) => next()), generateCheckAccess: jest.fn(() => (_req, _res, next) => next()), skipAgentCheck: jest.fn(), applyResumeContext: jest.fn(), diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index c8a3be3de5d..6b681a91c6d 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -58,6 +58,17 @@ it('renders originals without mutating canonical model/export input, and reloads expect(load).toHaveBeenCalledTimes(2); }); +it('does not fetch originals for an ordinary transcript', () => { + const plain = { ...canonical, privacyRevision: undefined }; + render( + + + , + ); + expect(screen.getByTestId('ordinary-transcript')).toBeInTheDocument(); + expect(load).not.toHaveBeenCalled(); +}); + it('renders only filtered text without an owner provider, as on external viewers', () => { render(); expect(screen.getByText(canonical.text)).toBeInTheDocument(); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index e10d3876320..6a912c4dd21 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -18,17 +18,26 @@ interface OwnerTextState { const empty: OwnerTextState = { scope: '', messages: new Map(), loading: false }; const OwnerTextContext = createContext(empty); -export function OwnerTextProvider({ - messages, - conversationId, - isSubmitting, - children, -}: { +interface OwnerTextProviderProps { messages: readonly TMessage[] | null; conversationId?: string; isSubmitting: boolean; children: ReactNode; -}) { +} + +export function OwnerTextProvider(props: OwnerTextProviderProps) { + if (!props.messages?.some((message) => message.isCreatedByUser && message.privacyRevision)) { + return <>{props.children}; + } + return ; +} + +function ActiveOwnerTextProvider({ + messages, + conversationId, + isSubmitting, + children, +}: OwnerTextProviderProps) { const { user } = useAuthContext(); const selection = useMemo( () => From 28ed3d1ea83c8fc29d498f43c21f8b5aa9894ac8 Mon Sep 17 00:00:00 2001 From: Lia Date: Thu, 24 Sep 2026 22:34:52 +0000 Subject: [PATCH 03/18] =?UTF-8?q?=F0=9F=A7=B7=20fix:=20Preserve=20Owner=20?= =?UTF-8?q?Text=20Across=20Copy=20and=20Edit=20Paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/utils/import/fork.js | 2 + api/server/utils/import/fork.spec.js | 21 +++ api/server/utils/import/importers.js | 2 + api/server/utils/import/importers.spec.js | 22 ++++ .../Chat/Messages/Content/EditMessage.tsx | 2 + .../Chat/Messages/Content/MessageContent.tsx | 2 +- .../Content/__tests__/EditMessage.spec.tsx | 22 ++++ .../Chat/Messages/PrivateText.spec.tsx | 44 ++++++- .../components/Chat/Messages/PrivateText.tsx | 123 ++++++++++++------ e2e/specs/mock/private-text.spec.ts | 10 +- .../src/methods/message.privateText.spec.ts | 119 +++++++++++++++++ .../data-schemas/src/methods/message.spec.ts | 4 +- packages/data-schemas/src/methods/message.ts | 70 ++++++++-- 13 files changed, 385 insertions(+), 58 deletions(-) diff --git a/api/server/utils/import/fork.js b/api/server/utils/import/fork.js index b777eed80da..6db2f8113e9 100644 --- a/api/server/utils/import/fork.js +++ b/api/server/utils/import/fork.js @@ -29,6 +29,8 @@ function cloneMessagesWithTimestamps( parentMessageId, createdAt, }; + delete clonedMessage.privateText; + delete clonedMessage.privacyRevision; if (detachSubagentRuntime) { delete clonedMessage.subagentTask; delete clonedMessage.subagentTranscript; diff --git a/api/server/utils/import/fork.spec.js b/api/server/utils/import/fork.spec.js index d53bb8775c1..97cc3d54c29 100644 --- a/api/server/utils/import/fork.spec.js +++ b/api/server/utils/import/fork.spec.js @@ -1453,6 +1453,27 @@ describe('splitAtTargetLevel', () => { }); describe('cloneMessagesWithTimestamps', () => { + test('does not carry private owner metadata into a new message identity', () => { + const importBatchBuilder = createImportBatchBuilder('owner'); + importBatchBuilder.startConversation(); + cloneMessagesWithTimestamps( + [ + { + messageId: 'source', + parentMessageId: Constants.NO_PARENT, + text: '[EMAIL_1]', + isCreatedByUser: true, + privateText: 'v1:ciphertext', + privacyRevision: 'source-revision', + }, + ], + importBatchBuilder, + ); + const cloned = importBatchBuilder.messages[0]; + expect(cloned.text).toBe('[EMAIL_1]'); + expect(cloned).not.toHaveProperty('privateText'); + expect(cloned).not.toHaveProperty('privacyRevision'); + }); test('should preserve user-submitted provenance without marking untouched model output', () => { const messagesToClone = [ { diff --git a/api/server/utils/import/importers.js b/api/server/utils/import/importers.js index d496749e5cd..ee211c67888 100644 --- a/api/server/utils/import/importers.js +++ b/api/server/utils/import/importers.js @@ -68,6 +68,8 @@ function sanitizeImportedMessage(message) { /** Server-private run state and trace sampling records never come from an import. */ const importable = withoutTraceRefs({ ...message }); delete importable.contextMeta; + delete importable.privateText; + delete importable.privacyRevision; return { ...importable, isUserSubmitted: true, diff --git a/api/server/utils/import/importers.spec.js b/api/server/utils/import/importers.spec.js index b9d945af94c..fed74ac4aba 100644 --- a/api/server/utils/import/importers.spec.js +++ b/api/server/utils/import/importers.spec.js @@ -38,6 +38,28 @@ afterEach(() => { jest.clearAllMocks(); }); +it('strips server-private owner metadata from untrusted LibreChat imports', async () => { + const jsonData = { + conversationId: 'imported', + title: 'Imported', + messages: [ + { + messageId: 'source', + parentMessageId: Constants.NO_PARENT, + text: '[EMAIL_1]', + isCreatedByUser: true, + privateText: 'v1:forged', + privacyRevision: 'forged', + }, + ], + }; + const importBatchBuilder = new ImportBatchBuilder('owner'); + await getImporter(jsonData)(jsonData, 'owner', () => importBatchBuilder); + expect(importBatchBuilder.messages[0].text).toBe('[EMAIL_1]'); + expect(importBatchBuilder.messages[0]).not.toHaveProperty('privateText'); + expect(importBatchBuilder.messages[0]).not.toHaveProperty('privacyRevision'); +}); + describe('importChatGptConvo', () => { it('should import conversation correctly', async () => { const expectedNumberOfMessages = 19; diff --git a/client/src/components/Chat/Messages/Content/EditMessage.tsx b/client/src/components/Chat/Messages/Content/EditMessage.tsx index 0920950cd9b..e05f08efdb2 100644 --- a/client/src/components/Chat/Messages/Content/EditMessage.tsx +++ b/client/src/components/Chat/Messages/Content/EditMessage.tsx @@ -162,6 +162,7 @@ const EditMessage = ({ ); if (!isInMessages) { message.text = data.text; + message.privacyRevision = undefined; } else { setMessages( messages.map((msg) => @@ -169,6 +170,7 @@ const EditMessage = ({ ? { ...msg, text: data.text, + privacyRevision: undefined, } : msg, ), diff --git a/client/src/components/Chat/Messages/Content/MessageContent.tsx b/client/src/components/Chat/Messages/Content/MessageContent.tsx index f7f75459779..a9e222b09dc 100644 --- a/client/src/components/Chat/Messages/Content/MessageContent.tsx +++ b/client/src/components/Chat/Messages/Content/MessageContent.tsx @@ -97,7 +97,7 @@ export const ErrorMessage = ({ ); }; -const DisplayMessage = ({ text, isCreatedByUser, message, showCursor }: TDisplayProps) => { +export const DisplayMessage = ({ text, isCreatedByUser, message, showCursor }: TDisplayProps) => { const { isSubmitting = false, isLatestMessage = false } = useMessageContext(); const enableUserMsgMarkdown = useRecoilValue(store.enableUserMsgMarkdown); const collapseLongUserMessages = useRecoilValue(store.collapseLongUserMessages); diff --git a/client/src/components/Chat/Messages/Content/__tests__/EditMessage.spec.tsx b/client/src/components/Chat/Messages/Content/__tests__/EditMessage.spec.tsx index 6cb0126658c..7a893906569 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/EditMessage.spec.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/EditMessage.spec.tsx @@ -129,6 +129,28 @@ describe('EditMessage', () => { expect(enterEdit).toHaveBeenCalledWith(true); }); + it('clears the private revision when a protected message is saved', async () => { + const protectedMessage = { ...message, privacyRevision: 'previous-revision' }; + mockGetMessages.mockReturnValue([protectedMessage]); + const user = userEvent.setup(); + renderEditor({ editedMessage: protectedMessage }); + + await user.clear(screen.getByTestId('message-text-editor')); + await user.type(screen.getByTestId('message-text-editor'), 'Clean edited message'); + await user.click(screen.getByRole('button', { name: 'com_ui_save' })); + + await waitFor(() => + expect(mockSetMessages).toHaveBeenCalledWith([ + expect.objectContaining({ + messageId: message.messageId, + text: 'Clean edited message', + privacyRevision: undefined, + }), + ]), + ); + expect(protectedMessage.privacyRevision).toBe('previous-revision'); + }); + it('writes the save onto the thread as it stands when the request resolves', async () => { const user = userEvent.setup(); const streamedAnswer = { diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index 6b681a91c6d..175a81c7263 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -10,6 +10,13 @@ jest.mock('~/hooks/AuthContext', () => ({ })); jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); jest.mock('librechat-data-provider', () => ({ dataService: { getOwnerMessageTexts: jest.fn() } })); +jest.mock('./Content/MessageContent', () => ({ + DisplayMessage: ({ text, message }: { text: string; message: TMessage }) => ( +
+ {text} +
+ ), +})); const canonical = Object.freeze({ messageId: 'message', @@ -28,13 +35,15 @@ const original = { function View({ conversationId = 'conversation', messages = [canonical], + displayIndex = 0, }: { conversationId?: string; messages?: TMessage[]; + displayIndex?: number; }) { return ( - + 
{JSON.stringify(messages)}
 ); @@ -50,6 +59,10 @@ it('renders originals without mutating canonical model/export input, and reloads const first = render(); expect(await screen.findByText('alice@example.com')).toBeInTheDocument(); expect(screen.getByTestId('canonical')).not.toHaveTextContent('alice@example.com'); + expect(screen.getByTestId('standard-user-renderer')).toHaveAttribute( + 'data-canonical', + canonical.text, + ); expect(canonical.text).toBe('[EMAIL_1_turn]'); expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_hidden'); first.unmount(); @@ -132,6 +145,35 @@ it('batches selected private rows and never loads ordinary messages', async () = expect(load.mock.calls.flatMap(([, ids]) => ids)).not.toContain('plain'); }); +it('loads batches concurrently, publishes completed batches, and only fetches new revisions', async () => { + let finishFirst!: (value: { messages: (typeof original)[] }) => void; + load.mockImplementationOnce( + () => + new Promise((resolve) => { + finishFirst = resolve; + }), + ); + load.mockResolvedValue({ messages: [{ ...original, messageId: 'message-9' }] }); + const messages: TMessage[] = Array.from({ length: 51 }, (_, index) => ({ + ...canonical, + messageId: `message-${index}`, + })); + const view = render(); + await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + await act(async () => { + finishFirst({ + messages: load.mock.calls[0][1].map((id: string) => ({ ...original, messageId: id })), + }); + }); + view.rerender( + , + ); + await waitFor(() => expect(load).toHaveBeenCalledTimes(3)); + expect(load.mock.calls[2][1]).toEqual(['message-51']); + expect(screen.getByText(original.text)).toBeInTheDocument(); +}); + it('invalidates an already rendered original when the canonical message changes', async () => { load.mockResolvedValue({ messages: [original] }); const view = render(); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index 6a912c4dd21..c47880e55dd 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -1,9 +1,21 @@ -import { createContext, useContext, useEffect, useMemo, useState } from 'react'; +import { + createContext, + lazy, + Suspense, + useContext, + useEffect, + useMemo, + useRef, + useState, +} from 'react'; import { dataService } from 'librechat-data-provider'; import type { TMessage } from 'librechat-data-provider'; import type { ReactNode } from 'react'; import { useAuthContext } from '~/hooks/AuthContext'; import { useLocalize } from '~/hooks'; +const DisplayMessage = lazy(async () => ({ + default: (await import('./Content/MessageContent')).DisplayMessage, +})); interface Original { canonicalText: string; @@ -32,12 +44,7 @@ export function OwnerTextProvider(props: OwnerTextProviderProps) { return ; } -function ActiveOwnerTextProvider({ - messages, - conversationId, - isSubmitting, - children, -}: OwnerTextProviderProps) { +function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTextProviderProps) { const { user } = useAuthContext(); const selection = useMemo( () => @@ -51,57 +58,85 @@ function ActiveOwnerTextProvider({ ); const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); const [state, setState] = useState(empty); + const cached = useRef<{ scope: string; messages: Map }>({ + scope: '', + messages: new Map(), + }); useEffect(() => { let cancelled = false; const selected = JSON.parse(selection) as Array<[string, string, string]>; if (!user?.id || !conversationId || selected.length === 0) { + cached.current = { scope: '', messages: new Map() }; setState(empty); return; } - setState({ scope, messages: new Map(), loading: true }); + const ownerScope = JSON.stringify([user.id, user.tenantId, conversationId]); + if (cached.current.scope !== ownerScope) { + cached.current = { scope: ownerScope, messages: new Map() }; + } + const originals = new Map(); + const pending: Array<[string, string, string]> = []; + for (const [id, revision, text] of selected) { + const prior = cached.current.messages.get(id); + if (prior?.revision === revision && prior.canonicalText === text) { + originals.set(id, prior); + } else { + pending.push([id, revision, text]); + } + } + // Do not retain originals from removed or edited messages. + cached.current.messages = originals; + setState({ scope, messages: new Map(originals), loading: pending.length > 0 }); + if (pending.length === 0) { + return; + } + let next = 0; const load = async () => { - const originals = new Map(); - try { - for (let index = 0; index < selected.length; index += 50) { - if (cancelled) { - return; - } - const batch = selected.slice(index, index + 50); - const result = await dataService.getOwnerMessageTexts( - conversationId, - batch.map(([id]) => id), - ); - for (const message of result.messages) { - if ( - batch.some( - ([id, revision, text]) => - id === message.messageId && - revision === message.revision && - text === message.canonicalText, - ) - ) { - originals.set(message.messageId, { - revision: message.revision, - text: message.text, - canonicalText: message.canonicalText, - }); + const workers = Array.from( + { length: Math.min(3, Math.ceil(pending.length / 50)) }, + async () => { + while (next < pending.length) { + const start = next; + next += 50; + const batch = pending.slice(start, start + 50); + const expected = new Map(batch.map(([id, revision, text]) => [id, { revision, text }])); + try { + const result = await dataService.getOwnerMessageTexts( + conversationId, + batch.map(([id]) => id), + ); + if (cancelled) { + return; + } + for (const message of result.messages) { + const match = expected.get(message.messageId); + if (match?.revision === message.revision && match.text === message.canonicalText) { + const original = { + revision: message.revision, + text: message.text, + canonicalText: message.canonicalText, + }; + originals.set(message.messageId, original); + cached.current.messages.set(message.messageId, original); + } + } + setState({ scope, messages: new Map(originals), loading: true }); + } catch { + // A failed batch does not discard successfully decrypted siblings. } } - } - if (!cancelled) { - setState({ scope, messages: originals, loading: false }); - } - } catch { - if (!cancelled) { - setState({ scope, messages: new Map(), loading: false }); - } + }, + ); + await Promise.all(workers); + if (!cancelled) { + setState({ scope, messages: new Map(originals), loading: false }); } }; void load(); return () => { cancelled = true; }; - }, [scope, selection, conversationId, user?.id, user?.tenantId, isSubmitting]); + }, [scope, selection, conversationId, user?.id, user?.tenantId]); const visible = state.scope === scope ? state : empty; return {children}; } @@ -119,7 +154,9 @@ export function PrivateText({ message }: { message: TMessage }) { : undefined; return (
-
{text ?? message.text}
+ + +

{localize('com_ui_private_text_hidden')} {text == null && ( diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index 214d79665d9..bbbfd4080e1 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -45,7 +45,12 @@ test('owner sees original after reload while provider, sharing, and canonical re await expect( messagesView(page).getByText('E2E private model input verified', { exact: true }), ).toBeVisible(); - await expect(messagesView(page).getByText(original, { exact: true })).toBeVisible(); + const ownerText = messagesView(page).getByText(original, { exact: true }); + await expect(ownerText).toBeVisible(); + const standardContainer = ownerText.locator( + 'xpath=ancestor::div[contains(@class,"text-message")][1]', + ); + await expect(standardContainer).toHaveAttribute('dir', 'auto'); await expect( messagesView(page).getByText('Private details hidden from the model', { exact: true }), ).toBeVisible(); @@ -74,6 +79,9 @@ test('owner sees original after reload while provider, sharing, and canonical re await page.reload(); await expect(messagesView(page).getByText(original, { exact: true })).toBeVisible(); + await expect( + ownerText.locator('xpath=ancestor::div[contains(@class,"text-message")][1]'), + ).toHaveAttribute('dir', 'auto'); for (const theme of ['light', 'dark']) { await page.evaluate( (dark) => document.documentElement.classList.toggle('dark', dark), diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index 57f6ce60538..16d4d98d733 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -134,6 +134,25 @@ it('does not accept sidecar writes from message parameters or generic edits', as }); }); +it('invalidates an encrypted original when an untyped canonical edit sets text to null', async () => { + await tenant('tenant-a', async () => { + const conversationId = uuid(); + const messageId = uuid(); + await methods.saveMessage( + { userId: 'owner' }, + { conversationId, messageId, text: '[EMAIL_1]', isCreatedByUser: true }, + { privateText: { envelope: 'v1:owner', revision: 'original' } }, + ); + await methods.updateMessage('owner', { messageId, text: null as unknown as string }); + const stored = await mongoose.models.Message.findOne({ messageId }) + .select('+privateText') + .lean(); + expect(stored).toMatchObject({ text: null }); + expect(stored).not.toHaveProperty('privateText'); + expect(stored).not.toHaveProperty('privacyRevision'); + }); +}); + it('does not return expired originals even before the TTL sweeper runs', async () => { await tenant('tenant-a', async () => { const messageId = uuid(); @@ -158,3 +177,103 @@ it('does not return expired originals even before the TTL sweeper runs', async ( ).toEqual([]); }); }); + +it('strips private fields from bulk copies, including overwrites of protected rows', async () => { + await tenant('tenant-a', async () => { + const conversationId = uuid(); + const messageId = uuid(); + await methods.saveMessage( + { userId: 'owner' }, + { conversationId, messageId, text: '[EMAIL_1]', isCreatedByUser: true }, + { privateText: { envelope: 'v1:owner', revision: 'original' } }, + ); + await methods.bulkSaveMessages([ + { + user: 'owner', + conversationId, + messageId, + text: 'clean copied', + isCreatedByUser: true, + privateText: 'forged', + privacyRevision: 'forged', + }, + { + user: 'owner', + conversationId, + messageId: uuid(), + text: 'fresh copied', + isCreatedByUser: true, + privateText: 'forged', + privacyRevision: 'forged', + }, + ]); + const rows = await mongoose.models.Message.find({ conversationId }) + .select('+privateText') + .lean(); + expect(rows).toHaveLength(2); + for (const row of rows) { + expect(row).not.toHaveProperty('privateText'); + expect(row).not.toHaveProperty('privacyRevision'); + } + }); +}); + +it.each([ + 'updateMessage', + 'updateMessageWithProvenance', + 'updateMessageText', + 'saveMessage', + 'recordMessage', +])('removes stored private text atomically when %s rewrites text', async (method) => { + await tenant('tenant-a', async () => { + const conversationId = uuid(); + const messageId = uuid(); + await methods.saveMessage( + { userId: 'owner' }, + { conversationId, messageId, text: '[EMAIL_1]', isCreatedByUser: true }, + { privateText: { envelope: 'v1:owner', revision: 'original' } }, + ); + if (method === 'updateMessage' || method === 'updateMessageWithProvenance') { + const updated = await methods.updateMessage('owner', { + messageId, + text: 'clean edited', + ...(method === 'updateMessageWithProvenance' && { userSubmittedPaths: ['/text'] }), + }); + expect(updated).not.toHaveProperty('privacyRevision'); + } else if (method === 'updateMessageText') { + await methods.updateMessageText('owner', { messageId, text: 'clean edited' }); + } else if (method === 'recordMessage') { + await methods.recordMessage({ + user: 'owner', + conversationId, + messageId, + text: 'clean edited', + isCreatedByUser: true, + }); + } else { + await methods.saveMessage( + { userId: 'owner' }, + { + conversationId, + messageId, + text: 'clean edited', + isCreatedByUser: true, + }, + ); + } + const stored = await mongoose.models.Message.findOne({ conversationId, messageId }) + .select('+privateText') + .lean(); + expect(stored).toMatchObject({ text: 'clean edited' }); + expect(stored).not.toHaveProperty('privateText'); + expect(stored).not.toHaveProperty('privacyRevision'); + expect( + await methods.getPrivateMessageTexts({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageIds: [messageId], + }), + ).toEqual([]); + }); +}); diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 8460b85dcfe..29c2112a8eb 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -263,7 +263,7 @@ describe('Message Operations', () => { expect(findOneAndUpdate).toHaveBeenCalledTimes(1); expect(findOneAndUpdate).toHaveBeenCalledWith( expect.anything(), - expect.objectContaining({ $unset: { contextMeta: 1 } }), + expect.objectContaining({ $unset: expect.objectContaining({ contextMeta: 1 }) }), expect.anything(), ); expect(updateOne).not.toHaveBeenCalled(); @@ -297,7 +297,7 @@ describe('Message Operations', () => { expect(findOneAndUpdate).toHaveBeenCalledTimes(1); expect(findOneAndUpdate).toHaveBeenCalledWith( expect.anything(), - expect.objectContaining({ $unset: { contextMeta: 1 } }), + expect.objectContaining({ $unset: expect.objectContaining({ contextMeta: 1 }) }), expect.anything(), ); } finally { diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 658bbb98eeb..e224927a322 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -280,12 +280,14 @@ function buildMessageSaveUpdate( options: { stampModelOutputOnInsert: boolean; unsetContextMeta: boolean; + unsetPrivateText?: boolean; retentionOnInsert?: { expiredAt: Date; isTemporary: false }; }, ): UpdateQuery { if ( !options.stampModelOutputOnInsert && !options.unsetContextMeta && + !options.unsetPrivateText && options.retentionOnInsert == null ) { return update; @@ -298,7 +300,12 @@ function buildMessageSaveUpdate( ...options.retentionOnInsert, }, }), - ...(options.unsetContextMeta && { $unset: { contextMeta: 1 } }), + ...((options.unsetContextMeta || options.unsetPrivateText) && { + $unset: { + ...(options.unsetContextMeta && { contextMeta: 1 }), + ...(options.unsetPrivateText && { privateText: 1, privacyRevision: 1 }), + }, + }), }; } @@ -312,6 +319,7 @@ async function findOneAndMergeMessageProvenance( upsert: boolean; stampModelOutputOnInsert?: boolean; unsetContextMeta?: boolean; + unsetPrivateText?: boolean; retentionOnInsert?: { expiredAt: Date; isTemporary: false }; }, ) { @@ -357,7 +365,12 @@ async function findOneAndMergeMessageProvenance( $set: { ...safeUpdate, ...provenance }, ...(current == null && options.retentionOnInsert != null && { $setOnInsert: options.retentionOnInsert }), - ...(options.unsetContextMeta && { $unset: { contextMeta: 1 } }), + ...((options.unsetContextMeta || options.unsetPrivateText) && { + $unset: { + ...(options.unsetContextMeta && { contextMeta: 1 }), + ...(options.unsetPrivateText && { privateText: 1, privacyRevision: 1 }), + }, + }), }, { upsert: options.upsert && current == null, new: true }, ); @@ -1019,13 +1032,24 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa update, userSubmittedPaths, userSubmittedMessageFieldPaths, - { upsert: true, stampModelOutputOnInsert, unsetContextMeta, retentionOnInsert }, + { + upsert: true, + stampModelOutputOnInsert, + unsetContextMeta, + unsetPrivateText: + metadata?.privateText == null && + Object.prototype.hasOwnProperty.call(update, 'text'), + retentionOnInsert, + }, ) : await Message.findOneAndUpdate( { messageId: params.messageId, user: userId }, buildMessageSaveUpdate(update, { stampModelOutputOnInsert, unsetContextMeta, + unsetPrivateText: + metadata?.privateText == null && + Object.prototype.hasOwnProperty.call(update, 'text'), retentionOnInsert, }), { upsert: true, new: true }, @@ -1114,6 +1138,8 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa const Message = mongoose.models.Message as Model; const bulkOps = messages.map((message) => { const normalizedMessage = { ...message }; + delete normalizedMessage.privateText; + delete normalizedMessage.privacyRevision; const provenance = capNormalizedProvenance( normalizeUserSubmittedPaths(message.userSubmittedPaths), normalizeUserSubmittedMessageFieldPaths(message.userSubmittedMessageFieldPaths), @@ -1135,7 +1161,10 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa return { updateOne: { filter: { messageId: message.messageId }, - update: normalizedMessage, + update: { + $set: normalizedMessage, + $unset: { privateText: 1, privacyRevision: 1 }, + }, timestamps: !overrideTimestamp, upsert: true, }, @@ -1178,6 +1207,8 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa userSubmittedMessageFieldPaths: _userSubmittedMessageFieldPaths, ...safeRest } = rest; + delete safeRest.privateText; + delete safeRest.privacyRevision; const message = { user, endpoint, @@ -1193,11 +1224,18 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa }), ...(provenance.promoteWholeMessage && { isUserSubmitted: true }), }; - const update = + const unsetPrivateText = Object.prototype.hasOwnProperty.call(safeRest, 'text'); + const stampModelOutputOnInsert = rest.isCreatedByUser === false && rest.isUserSubmitted === undefined && - !provenance.promoteWholeMessage - ? { $set: message, $setOnInsert: { isUserSubmitted: false } } + !provenance.promoteWholeMessage; + const update = + stampModelOutputOnInsert || unsetPrivateText + ? { + $set: message, + ...(stampModelOutputOnInsert && { $setOnInsert: { isUserSubmitted: false } }), + ...(unsetPrivateText && { $unset: { privateText: 1, privacyRevision: 1 } }), + } : message; return await Message.findOneAndUpdate({ user, messageId }, update, { @@ -1219,7 +1257,10 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa ) { try { const Message = mongoose.models.Message as Model; - await Message.updateOne({ messageId, user: userId }, { text }); + await Message.updateOne( + { messageId, user: userId }, + { $set: { text }, $unset: { privateText: 1, privacyRevision: 1 } }, + ); } catch (err) { logger.error('Error updating message text:', err); throw err; @@ -2011,9 +2052,18 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa update, submittedPaths, submittedMessageFields, - { upsert: false }, + { + upsert: false, + unsetPrivateText: Object.prototype.hasOwnProperty.call(update, 'text'), + }, ) - : await Message.findOneAndUpdate({ messageId, user: userId }, update, { new: true }); + : await Message.findOneAndUpdate( + { messageId, user: userId }, + Object.prototype.hasOwnProperty.call(update, 'text') + ? { $set: update, $unset: { privateText: 1, privacyRevision: 1 } } + : update, + { new: true }, + ); if (!updatedMessage) { throw new Error('Message not found or user not authorized.'); From 27bd2242a1f4b9ff09296214136e4e9cbdc35a88 Mon Sep 17 00:00:00 2001 From: Lia Date: Sun, 27 Sep 2026 12:34:23 +0000 Subject: [PATCH 04/18] =?UTF-8?q?=F0=9F=94=8F=20fix:=20Protect=20PII=20on?= =?UTF-8?q?=20Denial=20and=20Abort=20Paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/middleware/denyRequest.js | 19 ++- .../denyRequest.privateText.spec.js | 141 ++++++++++++++++++ .../routes/agents/__tests__/abort.spec.js | 48 ++++++ .../__tests__/chat.contentFilter.spec.js | 1 - .../__tests__/idempotencyLimiter.spec.js | 65 +++++++- api/server/routes/agents/chat.js | 8 - api/server/routes/agents/index.js | 29 +++- .../Chat/Messages/PrivateText.spec.tsx | 16 ++ .../components/Chat/Messages/PrivateText.tsx | 16 +- client/src/hooks/Chat/useChatFunctions.ts | 9 +- client/src/locales/en/translation.json | 1 + .../src/protection/private/submission.spec.ts | 81 ++++++++++ .../api/src/protection/private/submission.ts | 107 ++++++++++++- .../src/methods/message.privateText.spec.ts | 30 ++++ packages/data-schemas/src/methods/message.ts | 44 ++++++ 15 files changed, 589 insertions(+), 26 deletions(-) create mode 100644 api/server/middleware/denyRequest.privateText.spec.js diff --git a/api/server/middleware/denyRequest.js b/api/server/middleware/denyRequest.js index 521603a6d9c..8f4adf7b31b 100644 --- a/api/server/middleware/denyRequest.js +++ b/api/server/middleware/denyRequest.js @@ -1,5 +1,10 @@ const crypto = require('crypto'); -const { sendEvent } = require('@librechat/api'); +const { + sendEvent, + rejectUnprotectedDeniedMessage, + stampPrivateTextMessage, + savePrivateTextMessage, +} = require('@librechat/api'); const { getResponseSender, Constants } = require('librechat-data-provider'); const { sendError } = require('~/server/middleware/error'); const { saveMessage } = require('~/models'); @@ -21,6 +26,10 @@ const { saveMessage } = require('~/models'); * @throws {Error} Throws an error if there's an issue saving the message or sending the error. */ const denyRequest = async (req, res, errorMessage) => { + if (rejectUnprotectedDeniedMessage(req, res)) { + return; + } + let responseText = errorMessage; if (typeof errorMessage === 'object') { responseText = JSON.stringify(errorMessage); @@ -29,20 +38,22 @@ const denyRequest = async (req, res, errorMessage) => { const { messageId, conversationId: _convoId, parentMessageId, text } = req.body; const conversationId = _convoId ?? crypto.randomUUID(); - const userMessage = { + const userMessage = stampPrivateTextMessage(req, { sender: 'User', messageId: messageId ?? crypto.randomUUID(), parentMessageId, conversationId, isCreatedByUser: true, text, - }; + }); sendEvent(res, { message: userMessage, created: true }); const shouldSaveMessage = _convoId && parentMessageId && parentMessageId !== Constants.NO_PARENT; if (shouldSaveMessage) { - await saveMessage( + await savePrivateTextMessage( + saveMessage, + req, { userId: req?.user?.id, isTemporary: req?.resolvedConversation?.isTemporary ?? req?.body?.isTemporary, diff --git a/api/server/middleware/denyRequest.privateText.spec.js b/api/server/middleware/denyRequest.privateText.spec.js new file mode 100644 index 00000000000..fce6f63f598 --- /dev/null +++ b/api/server/middleware/denyRequest.privateText.spec.js @@ -0,0 +1,141 @@ +const { v4: uuidv4 } = require('uuid'); +const mockSaveMessage = jest.fn(); +const mockSendEvent = jest.fn(); +const mockSendError = jest.fn().mockResolvedValue(undefined); + +jest.mock('@librechat/api', () => ({ + ...jest.requireActual('@librechat/api'), + sendEvent: (...args) => mockSendEvent(...args), +})); +jest.mock('~/models', () => ({ + saveMessage: (...args) => mockSaveMessage(...args), +})); +jest.mock('~/server/middleware/error', () => ({ + sendError: (...args) => mockSendError(...args), +})); + +const { createPrivateTextIngress } = require('@librechat/api'); +const denyRequest = require('./denyRequest'); +const original = 'alice@example.com'; + +beforeEach(() => { + jest.clearAllMocks(); + mockSaveMessage.mockImplementation(async (_ctx, message, metadata) => ({ + ...message, + ...(metadata?.privateText && { privacyRevision: metadata.privateText.revision }), + })); +}); + +it('encrypts denied PII for an existing conversation and sends only the filtered view', async () => { + const req = { + method: 'POST', + path: '/', + originalUrl: '/api/agents/chat', + user: { id: 'owner', tenantId: 'tenant-a' }, + body: { + text: `Email ${original}`, + messageId: uuidv4(), + conversationId: uuidv4(), + parentMessageId: uuidv4(), + clientRequestId: uuidv4(), + }, + }; + const res = {}; + const next = jest.fn(); + createPrivateTextIngress({ + getFilters: () => ({ + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }), + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, res, next); + expect(next).toHaveBeenCalledTimes(1); + await denyRequest(req, res, { type: 'message_limit' }); + + expect(mockSendEvent).toHaveBeenCalledTimes(1); + const event = mockSendEvent.mock.calls[0][1]; + const [ctx, saved, metadata] = mockSaveMessage.mock.calls[0]; + expect(mockSaveMessage).toHaveBeenCalledTimes(1); + expect(ctx.userId).toBe('owner'); + expect(saved.text).toMatch(/^Email \[EMAIL_1_[a-f0-9]{32}\]$/); + expect(saved.privacyRevision).toBe(event.message.privacyRevision); + expect(metadata.privateText).toMatchObject({ revision: event.message.privacyRevision }); + expect(metadata.privateText.envelope).toMatch(/^v1:/); + expect(event.message.text).toBe(saved.text); + expect(JSON.stringify({ event, ctx, saved, metadata })).not.toContain(original); + expect(mockSendError).toHaveBeenCalledWith( + req, + res, + expect.objectContaining({ shouldSaveMessage: true }), + ); +}); + +it('fails closed before event or storage when denial hits untransformed private text', async () => { + const req = { + path: '/', + user: { id: 'owner', tenantId: 'tenant-a' }, + config: { + filters: { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }, + }, + body: { + text: `Email ${original}`, + files: [{}], + messageId: uuidv4(), + conversationId: uuidv4(), + parentMessageId: uuidv4(), + clientRequestId: uuidv4(), + }, + }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const next = jest.fn(); + createPrivateTextIngress({ + getFilters: () => req.config.filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, res, next); + expect(next).toHaveBeenCalledTimes(1); + await denyRequest(req, res, { type: 'message_limit' }); + expect(res.status).toHaveBeenCalledWith(400); + expect(JSON.stringify(res.json.mock.calls)).not.toContain(original); + expect(mockSendEvent).not.toHaveBeenCalled(); + expect(mockSaveMessage).not.toHaveBeenCalled(); + expect(mockSendError).not.toHaveBeenCalled(); +}); + +it('retains the existing denial behavior when no PII transformer ran', async () => { + const req = { + user: { id: 'owner' }, + body: { + text: 'ordinary message', + messageId: uuidv4(), + conversationId: uuidv4(), + parentMessageId: uuidv4(), + }, + }; + await denyRequest(req, {}, { type: 'message_limit' }); + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'owner' }), + expect.objectContaining({ text: 'ordinary message' }), + expect.not.objectContaining({ privateText: expect.anything() }), + ); +}); diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index e31bdc95c7d..ce1732bee43 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -24,6 +24,7 @@ const mockGenerationJobManager = { }; const mockSaveMessage = jest.fn(); +const mockHasPersistedPrivateText = jest.fn(); const mockRecordScheduleOutcome = jest.fn(); const mockBeginScheduledStop = jest.fn(); @@ -47,6 +48,7 @@ jest.mock('@librechat/api', () => ({ jest.mock('~/models', () => ({ saveMessage: (...args) => mockSaveMessage(...args), + hasPersistedPrivateText: (...args) => mockHasPersistedPrivateText(...args), })); jest.mock('~/server/services/Schedules', () => ({ @@ -99,6 +101,8 @@ describe('Agent Abort Endpoint', () => { mockGenerationJobManager.getActiveJobIdsForUser.mockReset(); mockSaveMessage.mockReset(); mockSaveMessage.mockImplementation(async (_context, message) => message); + mockHasPersistedPrivateText.mockReset(); + mockHasPersistedPrivateText.mockResolvedValue(true); mockRecordScheduleOutcome.mockReset(); mockRecordScheduleOutcome.mockResolvedValue(true); mockBeginScheduledStop.mockReset(); @@ -361,6 +365,50 @@ describe('Agent Abort Endpoint', () => { }); describe('Partial Response Saving', () => { + it('does not overwrite a persisted protected user sidecar while stopping a run', async () => { + const conversationId = 'test-stream-123'; + const userMessageId = 'protected-user-msg'; + const privacyRevision = 'protected-revision'; + const text = 'Email [EMAIL_1_protected]'; + const abortResult = { + success: true, + jobData: { + userMessage: { messageId: userMessageId, privacyRevision, text }, + responseMessageId: 'protected-response', + conversationId, + endpoint: 'agents', + }, + content: [{ type: 'text', text: 'Partial answer' }], + text: 'Partial answer', + }; + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'test-user-123' }, + }); + mockGenerationJobManager.abortJob.mockImplementation(async (_streamId, options) => { + await options.beforePublish(abortResult); + return abortResult; + }); + + const response = await request(app).post('/api/agents/chat/abort').send({ conversationId }); + + expect(response.status).toBe(200); + expect(mockHasPersistedPrivateText).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'test-user-123', + messageId: userMessageId, + conversationId, + privacyRevision, + text, + }), + ); + expect(mockSaveMessage).toHaveBeenCalledTimes(1); + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ messageId: 'protected-response', isCreatedByUser: false }), + expect.anything(), + ); + }); + it('should save partial response when both userMessage and responseMessageId exist', async () => { const jobStreamId = 'test-stream-123'; const userMessageId = 'user-msg-123'; diff --git a/api/server/routes/agents/__tests__/chat.contentFilter.spec.js b/api/server/routes/agents/__tests__/chat.contentFilter.spec.js index d77e8744d59..b6439ee18dc 100644 --- a/api/server/routes/agents/__tests__/chat.contentFilter.spec.js +++ b/api/server/routes/agents/__tests__/chat.contentFilter.spec.js @@ -14,7 +14,6 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ createMessageFilterPii: mockCreateMessageFilterPii, - createPrivateTextIngress: jest.fn(() => (_req, _res, next) => next()), generateCheckAccess: jest.fn(() => (_req, _res, next) => next()), skipAgentCheck: jest.fn(), applyResumeContext: jest.fn(), diff --git a/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js index d911f4d0f25..d842b7fce27 100644 --- a/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js +++ b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js @@ -8,6 +8,19 @@ const mockRetryLimiter = jest.fn((_req, _res, next) => next()); const mockRetryProbeLimiter = jest.fn((_req, _res, next) => next()); const mockExemptAgentTrigger = jest.fn(() => false); const mockExemptSchedule = jest.fn(() => false); +const mockIngress = jest.fn((req, _res, next) => { + if (req.config?.filters?.messages?.pii?.action === 'redact') { + req.body.text = '[EMAIL_1]'; + } + next(); +}); +const mockCheckBan = jest.fn((_req, _res, next) => next()); +const mockConfigMiddleware = jest.fn((req, _res, next) => { + if (req.headers['x-test-private'] === 'yes') { + req.config = { filters: { messages: { pii: { action: 'redact' } } } }; + } + next(); +}); jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -35,18 +48,27 @@ jest.mock('@librechat/api', () => ({ exemptAgentTriggerFromIpLimiter: (...args) => mockExemptAgentTrigger(...args), exemptFromUserLimiter: (...args) => mockExemptSchedule(...args), createMessageFilterPii: jest.fn(() => (_req, _res, next) => next()), + createPrivateTextIngress: jest.fn( + () => + (...args) => + mockIngress(...args), + ), + isPrivateTextChatSubmission: (req) => + req.method === 'POST' && + req.originalUrl === '/agents/chat' && + typeof req.body?.text === 'string', })); jest.mock('~/server/middleware', () => ({ uaParser: (_req, _res, next) => next(), - checkBan: (_req, _res, next) => next(), + checkBan: (...args) => mockCheckBan(...args), requireJwtAuth: (req, _res, next) => { req.user = { id: 'user-1' }; next(); }, moderateText: (_req, _res, next) => next(), messageIpLimiter: (...args) => mockIpLimiter(...args), - configMiddleware: (_req, _res, next) => next(), + configMiddleware: (...args) => mockConfigMiddleware(...args), messageUserLimiter: (...args) => mockUserLimiter(...args), })); @@ -103,6 +125,45 @@ describe('start-generation idempotency before message limiters', () => { mockExemptSchedule.mockReturnValue(false); }); + it('filters before a ban denial, IP limit, and user limit without charging config twice', async () => { + mockHasGenerationClaim.mockResolvedValue(false); + const payload = { text: 'alice@example.com', clientRequestId: 'request-privacy' }; + mockCheckBan.mockImplementationOnce((req, res) => { + expect(req.body.text).toBe('[EMAIL_1]'); + res.status(403).json({ banned: true }); + }); + const banned = await request(app) + .post('/agents/chat') + .set('X-Test-Private', 'yes') + .send(payload); + expect(banned.status).toBe(403); + expect(mockIpLimiter).not.toHaveBeenCalled(); + + mockIpLimiter.mockImplementationOnce((req, res) => { + expect(req.body.text).toBe('[EMAIL_1]'); + res.status(429).json({ limited: 'ip' }); + }); + const ipLimited = await request(app) + .post('/agents/chat') + .set('X-Test-Private', 'yes') + .send(payload); + expect(ipLimited.status).toBe(429); + expect(mockUserLimiter).not.toHaveBeenCalled(); + + mockIpLimiter.mockImplementationOnce((_req, _res, next) => next()); + mockUserLimiter.mockImplementationOnce((req, res) => { + expect(req.body.text).toBe('[EMAIL_1]'); + res.status(429).json({ limited: 'user' }); + }); + const userLimited = await request(app) + .post('/agents/chat') + .set('X-Test-Private', 'yes') + .send(payload); + expect(userLimited.status).toBe(429); + expect(mockIngress).toHaveBeenCalledTimes(3); + expect(mockConfigMiddleware).toHaveBeenCalledTimes(3); + }); + it('keeps a confirmed retry behind the shared IP limiter', async () => { mockHasGenerationClaim.mockResolvedValue(true); mockIpLimiter.mockImplementationOnce((_req, _res, next) => next()); diff --git a/api/server/routes/agents/chat.js b/api/server/routes/agents/chat.js index d25d5e353b2..16ecf54fd04 100644 --- a/api/server/routes/agents/chat.js +++ b/api/server/routes/agents/chat.js @@ -2,7 +2,6 @@ const express = require('express'); const { logger } = require('@librechat/data-schemas'); const { createMessageFilterPii, - createPrivateTextIngress, reportLocatorTraversalFailure, generateCheckAccess, skipAgentCheck, @@ -72,13 +71,6 @@ const restoreResumeContext = async (req, res, next) => { }; router.use(restoreResumeContext); -router.use( - createPrivateTextIngress({ - getFilters: (req) => req.config?.filters, - getLegacyPii: (req) => req.config?.messageFilter?.pii, - getKey: () => process.env.CREDS_KEY ?? '', - }), -); router.use( createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 52a4c34eeb1..5c84b63a192 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -15,6 +15,9 @@ const { attachAskUserQuestionAnswers, attachAskUserQuestionArgs, createMessageFilterPii, + createPrivateTextIngress, + isPrivateTextChatSubmission, + saveAbortedUserMessage, isAgentTriggerRequest, exemptAgentTriggerFromIpLimiter, captureScheduleFireContext, @@ -49,7 +52,7 @@ const { getServerGenerationProtocol, negotiateExistingGenerationProtocol, } = require('~/server/controllers/agents/protocol'); -const { getFiles, saveMessage } = require('~/models'); +const { getFiles, saveMessage, hasPersistedPrivateText } = require('~/models'); const { recordScheduleOutcome, beginScheduledStop, @@ -146,6 +149,18 @@ router.use((req, _res, next) => { captureScheduleFireContext(req); next(); }); +// Run config and the privacy boundary before ban/limiter denials, which may persist +// the submitted user turn. Other chat routes retain their existing config path. +const privateTextIngress = createPrivateTextIngress({ + getFilters: (req) => req.config?.filters, + getLegacyPii: (req) => req.config?.messageFilter?.pii, + getKey: () => process.env.CREDS_KEY ?? '', +}); +router.use( + '/chat', + unless((req) => !isPrivateTextChatSubmission(req), configMiddleware), + unless((req) => !isPrivateTextChatSubmission(req), privateTextIngress), +); router.use(checkBan); router.use(uaParser); @@ -829,9 +844,13 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { * write and checkpoint cleanup so every independently useful * operation gets a chance to succeed. */ try { - const persistedRequest = await saveMessage(messageContext, requestMessage, { - context: 'api/server/routes/agents/index.js - abort user prerequisite', - }); + const persistedRequest = await saveAbortedUserMessage( + { saveMessage, hasPersistedPrivateText }, + messageContext, + requestMessage, + { context: 'api/server/routes/agents/index.js - abort user prerequisite' }, + req.user?.tenantId, + ); if (!persistedRequest) { throw new Error('Abort user prerequisite was not persisted'); } @@ -1153,7 +1172,7 @@ router.use('/', v1); const chatRouter = express.Router(); const useMessageIpLimiter = isEnabled(LIMIT_MESSAGE_IP); const useMessageUserLimiter = isEnabled(LIMIT_MESSAGE_USER); -chatRouter.use(configMiddleware); +chatRouter.use(unless(isPrivateTextChatSubmission, configMiddleware)); if (useMessageIpLimiter || useMessageUserLimiter) { chatRouter.use( unless( diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index 175a81c7263..76561d3ab26 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -104,6 +104,22 @@ it('shows loading then safe unavailable text when decryption or authorization fa expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'); }); +it('offers a safe retry after a transient owner-text request failure', async () => { + load.mockRejectedValueOnce(new Error('temporary outage')); + load.mockResolvedValueOnce({ messages: [original] }); + render(); + + const retry = await screen.findByRole('button', { name: 'com_ui_private_text_retry' }); + expect(screen.getByText(canonical.text)).toBeInTheDocument(); + expect(screen.getByTestId('canonical')).not.toHaveTextContent(original.text); + await act(async () => { + retry.click(); + }); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(2); + expect(screen.queryByRole('button', { name: 'com_ui_private_text_retry' })).toBeNull(); +}); + it('rejects stale revisions instead of restoring a previous original', async () => { load.mockResolvedValue({ messages: [{ ...original, revision: 'old-revision' }] }); render(); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index c47880e55dd..08f1b13dd64 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -26,6 +26,7 @@ interface OwnerTextState { scope: string; messages: ReadonlyMap; loading: boolean; + retry?: () => void; } const empty: OwnerTextState = { scope: '', messages: new Map(), loading: false }; const OwnerTextContext = createContext(empty); @@ -58,6 +59,7 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe ); const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); const [state, setState] = useState(empty); + const [retryAttempt, setRetryAttempt] = useState(0); const cached = useRef<{ scope: string; messages: Map }>({ scope: '', messages: new Map(), @@ -86,7 +88,8 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe } // Do not retain originals from removed or edited messages. cached.current.messages = originals; - setState({ scope, messages: new Map(originals), loading: pending.length > 0 }); + const retry = () => setRetryAttempt((attempt) => attempt + 1); + setState({ scope, messages: new Map(originals), loading: pending.length > 0, retry }); if (pending.length === 0) { return; } @@ -120,7 +123,7 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe cached.current.messages.set(message.messageId, original); } } - setState({ scope, messages: new Map(originals), loading: true }); + setState({ scope, messages: new Map(originals), loading: true, retry }); } catch { // A failed batch does not discard successfully decrypted siblings. } @@ -129,14 +132,14 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe ); await Promise.all(workers); if (!cancelled) { - setState({ scope, messages: new Map(originals), loading: false }); + setState({ scope, messages: new Map(originals), loading: false, retry }); } }; void load(); return () => { cancelled = true; }; - }, [scope, selection, conversationId, user?.id, user?.tenantId]); + }, [scope, selection, conversationId, user?.id, user?.tenantId, retryAttempt]); const visible = state.scope === scope ? state : empty; return {children}; } @@ -169,6 +172,11 @@ export function PrivateText({ message }: { message: TMessage }) { )}

+ {text == null && !state.loading && state.retry != null && ( + + )}
); } diff --git a/client/src/hooks/Chat/useChatFunctions.ts b/client/src/hooks/Chat/useChatFunctions.ts index e102ec97249..3dc1e86ab34 100644 --- a/client/src/hooks/Chat/useChatFunctions.ts +++ b/client/src/hooks/Chat/useChatFunctions.ts @@ -806,7 +806,14 @@ export default function useChatFunctions({ ); setSubmissionStart(Date.now()); setSubmission(submission); - logger.dir('message_stream', submission, { depth: null }); + logger.dir('message_stream', { + conversationId, + messageId: currentMsg.messageId, + parentMessageId: currentMsg.parentMessageId, + isEdited: isEditOrContinue, + isRegenerate: regenerateShaped, + isContinued, + }); }; const regenerate = ( diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index a99d7586e78..b78997b8f0c 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1,6 +1,7 @@ { "com_ui_private_text_hidden": "Private details hidden from the model", "com_ui_private_text_loading": "Loading your original message…", + "com_ui_private_text_retry": "Retry loading original", "com_ui_private_text_unavailable": "Original unavailable. Showing the filtered message.", "com_ui_background_tasks": "Background tasks", "com_ui_background_tasks_cancel_disabled": "Stopping background tools is not enabled on this server", diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 44bda5f0eae..9ef1d3331b9 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -6,6 +6,8 @@ import { savePrivateTextMessage, stampPrivateTextMessage, requirePrivateTextPersistence, + saveAbortedUserMessage, + isPrivateTextChatSubmission, privateTextBinding, } from './submission'; import { createPrivateTextCipher } from './crypto'; @@ -52,6 +54,85 @@ function submit(overrides: object = {}, encryptionKey = key) { } describe('private text submission boundary', () => { + it('limits early processing to interactive POSTs, excluding controls and queued work', () => { + for (const path of ['/api/agents/chat', '/api/agents/chat/safe-ephemeral']) { + expect( + isPrivateTextChatSubmission({ + method: 'POST', + originalUrl: path, + body: { text: original }, + } as Request), + ).toBe(true); + } + for (const path of [ + '/api/agents/chat/abort', + '/api/agents/chat/resume', + '/api/agents/chat/queued-turns', + '/api/agents/chat/steer', + '/api/agents/chat/status/one', + ]) { + expect( + isPrivateTextChatSubmission({ + method: 'POST', + originalUrl: path, + body: { text: original }, + } as Request), + ).toBe(false); + } + expect( + isPrivateTextChatSubmission({ + method: 'GET', + originalUrl: '/api/agents/chat', + body: { text: original }, + } as Request), + ).toBe(false); + }); + + it('preserves the exact encrypted user row on Stop and fails closed when missing', async () => { + const saveMessage: MessageMethods['saveMessage'] = jest.fn( + async (_ctx, message) => message as IMessage, + ); + const hasPersistedPrivateText = jest.fn(async () => true); + const store = { saveMessage, hasPersistedPrivateText }; + const { message } = submit(); + expect( + await saveAbortedUserMessage( + store, + { userId: 'owner' }, + message, + { context: 'Stop' }, + 'tenant-a', + ), + ).toBe(true); + expect(hasPersistedPrivateText).toHaveBeenCalledWith({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId: message.conversationId, + messageId: message.messageId, + text: message.text, + privacyRevision: message.privacyRevision, + }); + expect(saveMessage).not.toHaveBeenCalled(); + hasPersistedPrivateText.mockResolvedValueOnce(false); + await expect( + saveAbortedUserMessage(store, { userId: 'owner' }, message, { context: 'Stop' }, 'tenant-a'), + ).rejects.toThrow('private value'); + expect(saveMessage).not.toHaveBeenCalled(); + expect( + await saveAbortedUserMessage( + store, + { userId: 'owner' }, + { + ...message, + privacyRevision: undefined, + }, + { context: 'ordinary Stop' }, + 'tenant-a', + ), + ).toBe(true); + expect(saveMessage).toHaveBeenCalledTimes(1); + }); + it('replaces request text before consumers and exposes no original in metadata or serialization', () => { const { req, message, next } = submit(); expect(next).toHaveBeenCalledTimes(1); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 98c0bb8fdda..1ad0fc7ca9a 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -1,6 +1,6 @@ import type { FiltersConfig, MessageFilterPiiConfig } from 'librechat-data-provider'; +import type { RequestHandler, Request, Response } from 'express'; import type { MessageMethods } from '@librechat/data-schemas'; -import type { RequestHandler, Request } from 'express'; import type { PrivateTextCipher } from './crypto'; import { ContentFilterError } from '../../middleware/contentFilter'; import { createPiiTextTransformer } from '../transform'; @@ -26,6 +26,33 @@ interface Capture { const captures = new WeakMap(); +const CONTROL_ROUTES = new Set([ + 'abort', + 'steer', + 'queued-turns', + 'stream', + 'status', + 'active', + 'resume', +]); + +/** Only actual interactive chat POSTs have an owner-view text sidecar in this slice. */ +export function isPrivateTextChatSubmission(req: Request): boolean { + if (req.method !== 'POST' || typeof req.body?.text !== 'string') { + return false; + } + const path = req.originalUrl?.split('?', 1)[0]?.replace(/\/$/, ''); + const base = '/api/agents/chat'; + if (path === base) { + return true; + } + if (path == null || !path.startsWith(`${base}/`)) { + return false; + } + const child = path.slice(base.length + 1); + return child.length > 0 && !child.includes('/') && !CONTROL_ROUTES.has(child.toLowerCase()); +} + function unavailable(): ContentFilterError { return new ContentFilterError({ detectorId: 'pii-pattern', @@ -39,6 +66,51 @@ function unavailable(): ContentFilterError { }); } +type PrivateTextRequest = Request & { + config?: { filters?: FiltersConfig; messageFilter?: { pii?: MessageFilterPiiConfig } }; +}; + +/** A denial may persist a turn before the regular message filter runs. Never store unfiltered PII. */ +export function rejectUnprotectedDeniedMessage(req: PrivateTextRequest, res: Response): boolean { + const filters = req.config?.filters; + const legacyPii = req.config?.messageFilter?.pii; + const text = req.body?.text; + if ( + typeof text !== 'string' || + (filters?.messages?.pii?.action !== 'redact' && legacyPii == null) || + captures.get(req)?.text === text + ) { + return false; + } + try { + const finding = inspectContent( + [ + { + id: 'chat.text', + path: '/text', + text, + source: 'message', + field: 'text', + format: 'plain', + treatment: 'replaceable', + provenance: 'user', + }, + ], + { filters, legacyPii }, + ); + if (finding == null) { + return false; + } + } catch { + // A broken policy must not turn a denial into unfiltered storage. + } + res.status(400).json({ + error: 'content_filter_block', + message: 'Private details could not be protected. Nothing was sent to the model.', + }); + return true; +} + export function privateTextBinding( userId: string, tenantId: string, @@ -236,3 +308,36 @@ export async function requirePrivateTextPersistence( throw unavailable(); } } + +/** + * A Stop request has a different Express request from the original turn and cannot + * retrieve the original plaintext. A protected turn already passed the pre-model + * persistence barrier: verify that exact owner row instead of resaving its text + * and unsetting its sidecar. If that prerequisite disappeared, fail closed. + */ +export async function saveAbortedUserMessage( + store: Pick, + ctx: Parameters[0], + message: Parameters[1], + metadata: Parameters[2], + tenantId?: string, +): Promise { + if (typeof message.privacyRevision !== 'string' || message.privacyRevision.length === 0) { + return (await store.saveMessage(ctx, message, metadata)) != null; + } + if (!message.messageId || !message.conversationId || typeof message.text !== 'string') { + throw unavailable(); + } + const persisted = await store.hasPersistedPrivateText({ + userId: ctx.userId, + tenantId, + conversationId: message.conversationId, + messageId: message.messageId, + text: message.text, + privacyRevision: message.privacyRevision, + }); + if (!persisted) { + throw unavailable(); + } + return true; +} diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index 16d4d98d733..4d7e3a13137 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -61,6 +61,26 @@ it('stores both views atomically and excludes ciphertext from ordinary and clien privacyRevision: 'turn', text: '[EMAIL_1_turn]', }); + const protectedRow = { + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageId, + text: '[EMAIL_1_turn]', + privacyRevision: 'turn', + }; + expect(await methods.hasPersistedPrivateText(protectedRow)).toBe(true); + for (const mismatch of [ + { userId: 'other' }, + { tenantId: 'tenant-b' }, + { conversationId: uuid() }, + { messageId: uuid() }, + { text: 'changed text' }, + { privacyRevision: 'wrong' }, + ]) { + expect(await methods.hasPersistedPrivateText({ ...protectedRow, ...mismatch })).toBe(false); + } + expect( await methods.getPrivateMessageTexts({ userId: 'other', @@ -275,5 +295,15 @@ it.each([ messageIds: [messageId], }), ).toEqual([]); + expect( + await methods.hasPersistedPrivateText({ + userId: 'owner', + tenantId: 'tenant-a', + conversationId, + messageId, + text: '[EMAIL_1]', + privacyRevision: 'turn', + }), + ).toBe(false); }); }); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 783915f0f0e..b0395ef9155 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -673,6 +673,14 @@ export interface PrivateTextRead { } export interface MessageMethods { + hasPersistedPrivateText(input: { + userId: string; + tenantId?: string; + conversationId: string; + messageId: string; + privacyRevision: string; + text: string; + }): Promise; getPrivateMessageTexts(input: { userId: string; tenantId?: string; @@ -3835,6 +3843,41 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa return Message.meiliSearch(query, searchOptions, hydrate); } + async function hasPersistedPrivateText(input: { + userId: string; + tenantId?: string; + conversationId: string; + messageId: string; + privacyRevision: string; + text: string; + }): Promise { + if ( + !input.userId || + !input.messageId || + !input.privacyRevision || + !UUID_REGEX.test(input.conversationId) + ) { + return false; + } + const activeTenant = tenantStorage.getStore()?.tenantId; + if (activeTenant != null && activeTenant !== input.tenantId) { + return false; + } + const Message = mongoose.models.Message as Model; + return ( + (await Message.exists({ + user: input.userId, + ...traceTenantScope(input.tenantId), + conversationId: input.conversationId, + messageId: input.messageId, + text: input.text, + privacyRevision: input.privacyRevision, + privateText: { $exists: true }, + $or: [{ expiredAt: null }, { expiredAt: { $gt: new Date() } }], + })) != null + ); + } + async function getPrivateMessageTexts(input: { userId: string; tenantId?: string; @@ -3864,6 +3907,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } return { + hasPersistedPrivateText, getPrivateMessageTexts, saveMessage, bulkSaveMessages, From a036847b15ee623efb0d2a04cb9a4c954c9d9c98 Mon Sep 17 00:00:00 2001 From: Lia Date: Mon, 28 Sep 2026 01:33:02 +0000 Subject: [PATCH 05/18] fix: preserve protected text through denial, stop, and owner reload --- api/server/controllers/agents/request.js | 13 ++- .../denyRequest.privateText.spec.js | 73 ++++++++++-- .../routes/agents/__tests__/abort.spec.js | 56 ++++++++++ .../__tests__/idempotencyLimiter.spec.js | 30 +++++ api/server/routes/agents/index.js | 38 ++++--- api/test/server/middleware/checkBan.test.js | 13 +++ .../Chat/Messages/PrivateText.spec.tsx | 47 +++++++- .../components/Chat/Messages/PrivateText.tsx | 69 +++++++++--- .../src/protection/private/submission.spec.ts | 61 +++++++++- .../api/src/protection/private/submission.ts | 104 ++++++++++++++++-- .../api/src/stream/GenerationJobManager.ts | 2 + .../api/src/stream/__tests__/startup.spec.ts | 48 ++++++++ .../api/src/stream/interfaces/IJobStore.ts | 1 + packages/api/src/types/events.ts | 1 + packages/data-provider/src/types/agents.ts | 2 + .../src/methods/message.privateText.spec.ts | 47 ++++++++ packages/data-schemas/src/methods/message.ts | 25 ++++- 17 files changed, 572 insertions(+), 58 deletions(-) diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index bdbff3a216d..be22de124b3 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -53,6 +53,7 @@ const { getFailedTurnTraceFields, resolveFailedTurnContent, savePrivateTextMessage, + stampPreliminaryPrivateTextMessage, } = require('@librechat/api'); const { disposeClient } = require('~/server/cleanup'); const { @@ -1593,10 +1594,13 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit const responseModel = getAgentResponseModel(req, endpointOption); const preliminaryUserMessage = isCompaction ? projectCompactionAnchor({ messageId: parentMessageId, conversationId }) - : getPreliminaryUserMessage( - { ...req.body, messageId: preallocatedUserMessageId }, - conversationId, - req._agentEventTriggerProjection, + : stampPreliminaryPrivateTextMessage( + req, + getPreliminaryUserMessage( + { ...req.body, messageId: preallocatedUserMessageId }, + conversationId, + req._agentEventTriggerProjection, + ), ); const job = await GenerationJobManager.createJob(streamId, userId, conversationId, { startupTelemetry, @@ -2307,6 +2311,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit conversationId: userMsg.conversationId, text: userMsg.text, quotes: userMsg.quotes, + privacyRevision: userMsg.privacyRevision, // Persist the turn's uploaded files here (authoritative job metadata) so a // HITL resume sources them from the job, not the user DB row — which the // approval prompt can race (the row save may still be in flight when a fast diff --git a/api/server/middleware/denyRequest.privateText.spec.js b/api/server/middleware/denyRequest.privateText.spec.js index fce6f63f598..9e33d36ec6d 100644 --- a/api/server/middleware/denyRequest.privateText.spec.js +++ b/api/server/middleware/denyRequest.privateText.spec.js @@ -31,6 +31,20 @@ it('encrypts denied PII for an existing conversation and sends only the filtered method: 'POST', path: '/', originalUrl: '/api/agents/chat', + config: { + filters: { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }, + }, user: { id: 'owner', tenantId: 'tenant-a' }, body: { text: `Email ${original}`, @@ -43,18 +57,7 @@ it('encrypts denied PII for an existing conversation and sends only the filtered const res = {}; const next = jest.fn(); createPrivateTextIngress({ - getFilters: () => ({ - messages: { - pii: { - action: 'redact', - fields: ['text'], - starterPatterns: [], - customPatterns: [ - { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, - ], - }, - }, - }), + getFilters: () => req.config.filters, getLegacyPii: () => undefined, getKey: () => 'ab'.repeat(32), })(req, res, next); @@ -122,6 +125,52 @@ it('fails closed before event or storage when denial hits untransformed private expect(mockSendError).not.toHaveBeenCalled(); }); +it.each(['redact', 'block'])( + 'rejects a queued-turn denial before events or storage with %s policy', + async (action) => { + const req = { + method: 'POST', + originalUrl: '/api/agents/chat/queued-turns', + user: { id: 'owner', tenantId: 'tenant-a' }, + config: { + filters: { + messages: { + pii: { + action, + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }, + }, + body: { text: `Email ${original}`, conversationId: uuidv4(), parentMessageId: uuidv4() }, + }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + await denyRequest(req, res, { type: 'ban' }); + expect(res.status).toHaveBeenCalledWith(400); + expect(JSON.stringify(res.json.mock.calls)).not.toContain(original); + expect(mockSendEvent).not.toHaveBeenCalled(); + expect(mockSaveMessage).not.toHaveBeenCalled(); + }, +); + +it('fails closed on a denied Agent submission when its policy was not loaded', async () => { + const req = { + method: 'POST', + originalUrl: '/api/agents/chat/queued-turns', + user: { id: 'owner' }, + body: { text: original, conversationId: uuidv4(), parentMessageId: uuidv4() }, + }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + await denyRequest(req, res, { type: 'ban' }); + expect(res.status).toHaveBeenCalledWith(400); + expect(mockSendEvent).not.toHaveBeenCalled(); + expect(mockSaveMessage).not.toHaveBeenCalled(); +}); + it('retains the existing denial behavior when no PII transformer ran', async () => { const req = { user: { id: 'owner' }, diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index ce1732bee43..07c5d6564d2 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -25,6 +25,7 @@ const mockGenerationJobManager = { const mockSaveMessage = jest.fn(); const mockHasPersistedPrivateText = jest.fn(); +const mockGetPrivateMessageTexts = jest.fn(); const mockRecordScheduleOutcome = jest.fn(); const mockBeginScheduledStop = jest.fn(); @@ -49,6 +50,7 @@ jest.mock('@librechat/api', () => ({ jest.mock('~/models', () => ({ saveMessage: (...args) => mockSaveMessage(...args), hasPersistedPrivateText: (...args) => mockHasPersistedPrivateText(...args), + getPrivateMessageTexts: (...args) => mockGetPrivateMessageTexts(...args), })); jest.mock('~/server/services/Schedules', () => ({ @@ -103,6 +105,8 @@ describe('Agent Abort Endpoint', () => { mockSaveMessage.mockImplementation(async (_context, message) => message); mockHasPersistedPrivateText.mockReset(); mockHasPersistedPrivateText.mockResolvedValue(true); + mockGetPrivateMessageTexts.mockReset(); + mockGetPrivateMessageTexts.mockResolvedValue([]); mockRecordScheduleOutcome.mockReset(); mockRecordScheduleOutcome.mockResolvedValue(true); mockBeginScheduledStop.mockReset(); @@ -409,6 +413,58 @@ describe('Agent Abort Endpoint', () => { ); }); + it('recovers an older revisionless job from its exact protected row without unsetting it', async () => { + const conversationId = 'test-stream-revisionless'; + const userMessageId = 'protected-user'; + const text = `Email [EMAIL_1_${'a'.repeat(32)}]`; + const abortResult = { + success: true, + jobData: { + createdEventEmitted: true, + userMessage: { messageId: userMessageId, text }, + responseMessageId: 'protected-response', + conversationId, + endpoint: 'agents', + }, + finalEvent: { requestMessage: { messageId: userMessageId, text } }, + content: [], + text: '', + }; + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { userId: 'test-user-123' }, + }); + mockGenerationJobManager.abortJob.mockImplementation(async (_id, options) => { + await options.beforePublish(abortResult); + return abortResult; + }); + mockGetPrivateMessageTexts.mockResolvedValueOnce([ + { + messageId: userMessageId, + text, + privacyRevision: 'recovered-revision', + privateText: 'v1:protected', + }, + ]); + + const response = await request(app).post('/api/agents/chat/abort').send({ conversationId }); + expect(response.status).toBe(200); + expect(mockGetPrivateMessageTexts).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'test-user-123', + conversationId, + messageIds: [userMessageId], + }), + ); + expect(abortResult.finalEvent.requestMessage.privacyRevision).toBe('recovered-revision'); + expect(mockHasPersistedPrivateText).not.toHaveBeenCalled(); + expect(mockSaveMessage).toHaveBeenCalledTimes(1); + expect(mockSaveMessage).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ messageId: 'protected-response' }), + expect.anything(), + ); + }); + it('should save partial response when both userMessage and responseMessageId exist', async () => { const jobStreamId = 'test-stream-123'; const userMessageId = 'user-msg-123'; diff --git a/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js index d842b7fce27..62b1120837c 100644 --- a/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js +++ b/api/server/routes/agents/__tests__/idempotencyLimiter.spec.js @@ -53,6 +53,7 @@ jest.mock('@librechat/api', () => ({ (...args) => mockIngress(...args), ), + isPreDenialTextSubmission: (req) => req.method === 'POST' && typeof req.body?.text === 'string', isPrivateTextChatSubmission: (req) => req.method === 'POST' && req.originalUrl === '/agents/chat' && @@ -164,6 +165,35 @@ describe('start-generation idempotency before message limiters', () => { expect(mockConfigMiddleware).toHaveBeenCalledTimes(3); }); + it.each(['/agents/chat/queued-turns', '/agents/chat/queued-turns/v2', '/agents/chat/steer'])( + 'loads policy once ahead of a banned text submission to %s without transforming it', + async (path) => { + mockCheckBan.mockImplementationOnce((req, res) => { + expect(req.config?.filters?.messages?.pii?.action).toBe('redact'); + expect(req.body.text).toBe('alice@example.com'); + res.status(403).json({ banned: true }); + }); + const response = await request(app) + .post(path) + .set('X-Test-Private', 'yes') + .send({ text: 'alice@example.com' }); + expect(response.status).toBe(403); + expect(mockConfigMiddleware).toHaveBeenCalledTimes(1); + expect(mockIngress).not.toHaveBeenCalled(); + }, + ); + + it('does not reload pre-denial config on an admitted queued submission', async () => { + mockIpLimiter.mockImplementationOnce((_req, _res, next) => next()); + mockUserLimiter.mockImplementationOnce((_req, _res, next) => next()); + const response = await request(app) + .post('/agents/chat/queued-turns') + .set('X-Test-Private', 'yes') + .send({ text: 'clean queued turn' }); + expect(response.status).toBe(202); + expect(mockConfigMiddleware).toHaveBeenCalledTimes(1); + }); + it('keeps a confirmed retry behind the shared IP limiter', async () => { mockHasGenerationClaim.mockResolvedValue(true); mockIpLimiter.mockImplementationOnce((_req, _res, next) => next()); diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 5c84b63a192..bec4e4a6c8f 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -17,6 +17,7 @@ const { createMessageFilterPii, createPrivateTextIngress, isPrivateTextChatSubmission, + isPreDenialTextSubmission, saveAbortedUserMessage, isAgentTriggerRequest, exemptAgentTriggerFromIpLimiter, @@ -52,7 +53,12 @@ const { getServerGenerationProtocol, negotiateExistingGenerationProtocol, } = require('~/server/controllers/agents/protocol'); -const { getFiles, saveMessage, hasPersistedPrivateText } = require('~/models'); +const { + getFiles, + saveMessage, + hasPersistedPrivateText, + getPrivateMessageTexts, +} = require('~/models'); const { recordScheduleOutcome, beginScheduledStop, @@ -149,16 +155,17 @@ router.use((req, _res, next) => { captureScheduleFireContext(req); next(); }); -// Run config and the privacy boundary before ban/limiter denials, which may persist -// the submitted user turn. Other chat routes retain their existing config path. +// Denials may persist submitted text on chat control routes too. Load policy before +// any such denial; transformation remains limited to fresh interactive turns. const privateTextIngress = createPrivateTextIngress({ getFilters: (req) => req.config?.filters, getLegacyPii: (req) => req.config?.messageFilter?.pii, getKey: () => process.env.CREDS_KEY ?? '', }); +const chatConfigMiddleware = unless((req) => req.config != null, configMiddleware); router.use( '/chat', - unless((req) => !isPrivateTextChatSubmission(req), configMiddleware), + unless((req) => !isPreDenialTextSubmission(req), configMiddleware), unless((req) => !isPrivateTextChatSubmission(req), privateTextIngress), ); router.use(checkBan); @@ -582,7 +589,7 @@ router.get('/chat/status/:conversationId', async (req, res) => { * @access Private * @description Mounted before chatRouter to bypass buildEndpointOption middleware */ -router.post('/chat/abort', configMiddleware, async (req, res, next) => { +router.post('/chat/abort', chatConfigMiddleware, async (req, res, next) => { logger.debug(`[AgentStream] ========== ABORT ENDPOINT HIT ==========`); logger.debug(`[AgentStream] Method: ${req.method}, Path: ${req.path}`); @@ -845,11 +852,12 @@ router.post('/chat/abort', configMiddleware, async (req, res, next) => { * operation gets a chance to succeed. */ try { const persistedRequest = await saveAbortedUserMessage( - { saveMessage, hasPersistedPrivateText }, + { saveMessage, hasPersistedPrivateText, getPrivateMessageTexts }, messageContext, requestMessage, { context: 'api/server/routes/agents/index.js - abort user prerequisite' }, req.user?.tenantId, + pendingAbortResult.finalEvent, ); if (!persistedRequest) { throw new Error('Abort user prerequisite was not persisted'); @@ -1072,7 +1080,7 @@ if (isEnabled(LIMIT_MESSAGE_USER)) { } router.post( '/chat/steer', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, @@ -1093,7 +1101,7 @@ router.post( */ router.post( '/chat/steer/deliver', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, @@ -1113,7 +1121,7 @@ router.post( */ router.post( '/chat/steer/cancel', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, SteerController.SteerCancelController, ); @@ -1126,14 +1134,14 @@ router.post( */ router.post( '/chat/steer/arm', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, SteerController.SteerArmController, ); router.post( '/chat/queued-turns', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, @@ -1146,7 +1154,7 @@ router.post( ); router.post( '/chat/queued-turns/v2', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, createMessageFilterPii({ onTraversalFailure: reportLocatorTraversalFailure, @@ -1159,10 +1167,10 @@ router.post( ); /** Synchronizing durable queue state is read-only and polled while work is * pending. It must not consume the model-submission admission budget. */ -router.get('/chat/queued-turns', configMiddleware, AgentQueuedTurnListController); +router.get('/chat/queued-turns', chatConfigMiddleware, AgentQueuedTurnListController); router.delete( '/chat/queued-turns/:queuedTurnId', - configMiddleware, + chatConfigMiddleware, ...steerLimiters, AgentQueuedTurnCancelController, ); @@ -1172,7 +1180,7 @@ router.use('/', v1); const chatRouter = express.Router(); const useMessageIpLimiter = isEnabled(LIMIT_MESSAGE_IP); const useMessageUserLimiter = isEnabled(LIMIT_MESSAGE_USER); -chatRouter.use(unless(isPrivateTextChatSubmission, configMiddleware)); +chatRouter.use(chatConfigMiddleware); if (useMessageIpLimiter || useMessageUserLimiter) { chatRouter.use( unless( diff --git a/api/test/server/middleware/checkBan.test.js b/api/test/server/middleware/checkBan.test.js index 39775c389a4..aeb2f4e3162 100644 --- a/api/test/server/middleware/checkBan.test.js +++ b/api/test/server/middleware/checkBan.test.js @@ -219,6 +219,19 @@ describe('checkBan middleware', () => { }, ); + it('treats a banned queued-turn POST as an early text-saving denial', async () => { + mockBanCacheGet.mockResolvedValueOnce({ expiresAt: Date.now() + 60000 }); + const req = createReq({ + method: 'POST', + baseUrl: '/api/agents', + originalUrl: '/api/agents/chat/queued-turns', + body: { text: 'submitted text' }, + }); + const res = createRes(); + await checkBan(req, res, jest.fn()); + expect(denyRequest).toHaveBeenCalledWith(req, res, { type: ViolationTypes.BAN }); + }); + it('preserves SSE denial for a banned browser interactive chat request', async () => { mockBanCacheGet.mockResolvedValueOnce({ expiresAt: Date.now() + 60000 }); const req = createReq({ diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index 76561d3ab26..b5f6a23bb84 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -36,13 +36,19 @@ function View({ conversationId = 'conversation', messages = [canonical], displayIndex = 0, + isSubmitting = false, }: { conversationId?: string; messages?: TMessage[]; displayIndex?: number; + isSubmitting?: boolean; }) { return ( - + 
{JSON.stringify(messages)}
 @@ -104,6 +110,34 @@ it('shows loading then safe unavailable text when decryption or authorization fa expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'); }); +it('retries a provisional empty owner read once on turn completion, without repeatedly polling', async () => { + load.mockResolvedValueOnce({ messages: [] }); + load.mockResolvedValueOnce({ messages: [original] }); + const view = render(); + await waitFor(() => + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), + ); + expect(load).toHaveBeenCalledTimes(1); + view.rerender(); + expect(load).toHaveBeenCalledTimes(1); + view.rerender(); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(2); + view.rerender(); + expect(load).toHaveBeenCalledTimes(2); +}); + +it('does not automatically retry an old missing owner row on unrelated submission transitions', async () => { + load.mockResolvedValue({ messages: [] }); + const view = render(); + await waitFor(() => + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), + ); + view.rerender(); + view.rerender(); + expect(load).toHaveBeenCalledTimes(1); +}); + it('offers a safe retry after a transient owner-text request failure', async () => { load.mockRejectedValueOnce(new Error('temporary outage')); load.mockResolvedValueOnce({ messages: [original] }); @@ -120,6 +154,17 @@ it('offers a safe retry after a transient owner-text request failure', async () expect(screen.queryByRole('button', { name: 'com_ui_private_text_retry' })).toBeNull(); }); +it('does not cache a failed decryption as if it contained a usable original', async () => { + load.mockResolvedValueOnce({ messages: [{ ...original, text: undefined }] }); + load.mockResolvedValueOnce({ messages: [original] }); + render(); + const retry = await screen.findByRole('button', { name: 'com_ui_private_text_retry' }); + expect(screen.getByText(canonical.text)).toBeInTheDocument(); + await act(async () => retry.click()); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(2); +}); + it('rejects stale revisions instead of restoring a previous original', async () => { load.mockResolvedValue({ messages: [{ ...original, revision: 'old-revision' }] }); render(); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index 08f1b13dd64..f1d4873266e 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -22,6 +22,11 @@ interface Original { revision: string; text?: string; } +interface MissingOriginal { + canonicalText: string; + revision: string; + retryOnCompletion: boolean; +} interface OwnerTextState { scope: string; messages: ReadonlyMap; @@ -45,7 +50,12 @@ export function OwnerTextProvider(props: OwnerTextProviderProps) { return ; } -function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTextProviderProps) { +function ActiveOwnerTextProvider({ + messages, + conversationId, + isSubmitting, + children, +}: OwnerTextProviderProps) { const { user } = useAuthContext(); const selection = useMemo( () => @@ -60,34 +70,47 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); const [state, setState] = useState(empty); const [retryAttempt, setRetryAttempt] = useState(0); - const cached = useRef<{ scope: string; messages: Map }>({ - scope: '', - messages: new Map(), - }); + const cached = useRef<{ + scope: string; + messages: Map; + missing: Map; + retryAttempt: number; + }>({ scope: '', messages: new Map(), missing: new Map(), retryAttempt: 0 }); useEffect(() => { let cancelled = false; const selected = JSON.parse(selection) as Array<[string, string, string]>; if (!user?.id || !conversationId || selected.length === 0) { - cached.current = { scope: '', messages: new Map() }; + cached.current = { scope: '', messages: new Map(), missing: new Map(), retryAttempt }; setState(empty); return; } const ownerScope = JSON.stringify([user.id, user.tenantId, conversationId]); if (cached.current.scope !== ownerScope) { - cached.current = { scope: ownerScope, messages: new Map() }; + cached.current = { scope: ownerScope, messages: new Map(), missing: new Map(), retryAttempt }; } + const forceRetry = cached.current.retryAttempt !== retryAttempt; + cached.current.retryAttempt = retryAttempt; const originals = new Map(); + const missing = new Map(); const pending: Array<[string, string, string]> = []; for (const [id, revision, text] of selected) { const prior = cached.current.messages.get(id); - if (prior?.revision === revision && prior.canonicalText === text) { + if (prior?.revision === revision && prior.canonicalText === text && prior.text != null) { originals.set(id, prior); - } else { - pending.push([id, revision, text]); + continue; } + const missed = cached.current.missing.get(id); + if (missed?.revision === revision && missed.canonicalText === text) { + if (!forceRetry && (isSubmitting || !missed.retryOnCompletion)) { + missing.set(id, missed); + continue; + } + } + pending.push([id, revision, text]); } - // Do not retain originals from removed or edited messages. + // Do not retain originals or failed reads from removed or edited messages. cached.current.messages = originals; + cached.current.missing = missing; const retry = () => setRetryAttempt((attempt) => attempt + 1); setState({ scope, messages: new Map(originals), loading: pending.length > 0, retry }); if (pending.length === 0) { @@ -113,7 +136,11 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe } for (const message of result.messages) { const match = expected.get(message.messageId); - if (match?.revision === message.revision && match.text === message.canonicalText) { + if ( + match?.revision === message.revision && + match.text === message.canonicalText && + typeof message.text === 'string' + ) { const original = { revision: message.revision, text: message.text, @@ -123,9 +150,23 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe cached.current.messages.set(message.messageId, original); } } + for (const [id, revision, text] of batch) { + if (!originals.has(id)) { + missing.set(id, { + revision, + canonicalText: text, + retryOnCompletion: isSubmitting, + }); + } + } setState({ scope, messages: new Map(originals), loading: true, retry }); } catch { - // A failed batch does not discard successfully decrypted siblings. + if (cancelled) { + return; + } + for (const [id, revision, text] of batch) { + missing.set(id, { revision, canonicalText: text, retryOnCompletion: isSubmitting }); + } } } }, @@ -139,7 +180,7 @@ function ActiveOwnerTextProvider({ messages, conversationId, children }: OwnerTe return () => { cancelled = true; }; - }, [scope, selection, conversationId, user?.id, user?.tenantId, retryAttempt]); + }, [scope, selection, conversationId, user?.id, user?.tenantId, isSubmitting, retryAttempt]); const visible = state.scope === scope ? state : empty; return {children}; } diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 9ef1d3331b9..99cdb07ff84 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -5,8 +5,10 @@ import { createPrivateTextIngress, savePrivateTextMessage, stampPrivateTextMessage, + stampPreliminaryPrivateTextMessage, requirePrivateTextPersistence, saveAbortedUserMessage, + isPreDenialTextSubmission, isPrivateTextChatSubmission, privateTextBinding, } from './submission'; @@ -79,6 +81,12 @@ describe('private text submission boundary', () => { } as Request), ).toBe(false); } + expect( + isPreDenialTextSubmission({ + method: 'POST', + body: { text: original }, + } as Request), + ).toBe(true); expect( isPrivateTextChatSubmission({ method: 'GET', @@ -93,7 +101,8 @@ describe('private text submission boundary', () => { async (_ctx, message) => message as IMessage, ); const hasPersistedPrivateText = jest.fn(async () => true); - const store = { saveMessage, hasPersistedPrivateText }; + const getPrivateMessageTexts = jest.fn(async (): Promise => []); + const store = { saveMessage, hasPersistedPrivateText, getPrivateMessageTexts }; const { message } = submit(); expect( await saveAbortedUserMessage( @@ -124,6 +133,7 @@ describe('private text submission boundary', () => { { userId: 'owner' }, { ...message, + text: 'ordinary turn without private values', privacyRevision: undefined, }, { context: 'ordinary Stop' }, @@ -133,6 +143,55 @@ describe('private text submission boundary', () => { expect(saveMessage).toHaveBeenCalledTimes(1); }); + it('checks old revisionless protected jobs against storage instead of overwriting them', async () => { + const { message } = submit(); + const saveMessage = jest.fn(async () => message as IMessage); + const getPrivateMessageTexts = jest.fn(async () => [ + { + messageId: message.messageId!, + text: message.text!, + privacyRevision: message.privacyRevision!, + privateText: 'v1:encrypted', + }, + ]); + const store = { + saveMessage, + hasPersistedPrivateText: jest.fn(async () => true), + getPrivateMessageTexts, + }; + const finalEvent = { requestMessage: { messageId: message.messageId, privacyRevision: '' } }; + const older = { ...message, privacyRevision: undefined }; + expect( + await saveAbortedUserMessage( + store, + { userId: 'owner' }, + older, + undefined, + 'tenant-a', + finalEvent, + ), + ).toBe(true); + expect(saveMessage).not.toHaveBeenCalled(); + expect(finalEvent.requestMessage.privacyRevision).toBe(message.privacyRevision); + getPrivateMessageTexts.mockResolvedValueOnce([]); + await expect( + saveAbortedUserMessage(store, { userId: 'owner' }, older, undefined, 'tenant-a'), + ).rejects.toThrow('private value'); + expect(saveMessage).not.toHaveBeenCalled(); + }); + + it('stamps a protected preliminary job message before the created event', () => { + const { req, message } = submit(); + const preliminary = stampPreliminaryPrivateTextMessage(req, { + messageId: message.messageId, + conversationId: message.conversationId, + text: message.text, + }); + expect(preliminary?.privacyRevision).toBe(message.privacyRevision); + expect(JSON.stringify(preliminary)).not.toContain(original); + expect(stampPreliminaryPrivateTextMessage(req, null)).toBeNull(); + }); + it('replaces request text before consumers and exposes no original in metadata or serialization', () => { const { req, message, next } = submit(); expect(next).toHaveBeenCalledTimes(1); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 1ad0fc7ca9a..2059245393b 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -25,6 +25,7 @@ interface Capture { } const captures = new WeakMap(); +const PRIVATE_PLACEHOLDER = /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]/; const CONTROL_ROUTES = new Set([ 'abort', @@ -36,6 +37,11 @@ const CONTROL_ROUTES = new Set([ 'resume', ]); +/** Every submitted text path can be denied before its route-specific filter runs. */ +export function isPreDenialTextSubmission(req: Request): boolean { + return req.method === 'POST' && typeof req.body?.text === 'string'; +} + /** Only actual interactive chat POSTs have an owner-view text sidecar in this slice. */ export function isPrivateTextChatSubmission(req: Request): boolean { if (req.method !== 'POST' || typeof req.body?.text !== 'string') { @@ -75,9 +81,20 @@ export function rejectUnprotectedDeniedMessage(req: PrivateTextRequest, res: Res const filters = req.config?.filters; const legacyPii = req.config?.messageFilter?.pii; const text = req.body?.text; + if ( + typeof text === 'string' && + req.config == null && + req.originalUrl?.split('?', 1)[0]?.startsWith('/api/agents/chat') + ) { + res.status(400).json({ + error: 'content_filter_block', + message: 'Private details could not be protected. Nothing was sent to the model.', + }); + return true; + } if ( typeof text !== 'string' || - (filters?.messages?.pii?.action !== 'redact' && legacyPii == null) || + (filters?.messages?.pii == null && legacyPii == null) || captures.get(req)?.text === text ) { return false; @@ -240,6 +257,21 @@ export function stampPrivateTextMessage( return message; } +/** The preliminary job record precedes the created event and may be read by Stop. */ +export function stampPreliminaryPrivateTextMessage( + req: object | undefined, + message: T | null, +): (T & { privacyRevision?: string }) | null { + if (message == null) { + return null; + } + const capture = req == null ? undefined : captures.get(req); + if (capture != null && capture.text === message.text) { + return { ...message, privacyRevision: capture.revision }; + } + return message; +} + /** Encrypts against final server-resolved message identity, then commits both views in one write. */ export async function savePrivateTextMessage( save: MessageMethods['saveMessage'], @@ -310,20 +342,74 @@ export async function requirePrivateTextPersistence( } /** - * A Stop request has a different Express request from the original turn and cannot - * retrieve the original plaintext. A protected turn already passed the pre-model - * persistence barrier: verify that exact owner row instead of resaving its text - * and unsetting its sidecar. If that prerequisite disappeared, fail closed. + * Stop has a different request and cannot retrieve the original plaintext. Verify + * protected rows instead of rewriting their canonical text. Older job records may + * omit the revision, so their prerequisite is an insert-only write. */ export async function saveAbortedUserMessage( - store: Pick, + store: Pick, ctx: Parameters[0], message: Parameters[1], metadata: Parameters[2], tenantId?: string, + finalEvent?: { requestMessage?: { privacyRevision?: string } | null }, ): Promise { - if (typeof message.privacyRevision !== 'string' || message.privacyRevision.length === 0) { - return (await store.saveMessage(ctx, message, metadata)) != null; + const revision = message.privacyRevision; + if (typeof revision !== 'string' || revision.length === 0) { + if (!message.messageId || !message.conversationId) { + throw unavailable(); + } + if (typeof message.text !== 'string' || !PRIVATE_PLACEHOLDER.test(message.text)) { + const saved = await store.saveMessage(ctx, message, { ...metadata, insertOnly: true }); + if ( + saved == null || + saved.messageId !== message.messageId || + saved.conversationId !== message.conversationId || + saved.text !== message.text + ) { + throw unavailable(); + } + if (typeof saved.privacyRevision === 'string' && saved.privacyRevision.length > 0) { + if (typeof saved.text !== 'string') { + throw unavailable(); + } + const exists = await store.hasPersistedPrivateText({ + userId: ctx.userId, + tenantId, + conversationId: message.conversationId, + messageId: message.messageId, + text: saved.text, + privacyRevision: saved.privacyRevision, + }); + if (!exists) { + throw unavailable(); + } + if (finalEvent?.requestMessage != null) { + finalEvent.requestMessage.privacyRevision = saved.privacyRevision; + } + } + return true; + } + const rows = await store.getPrivateMessageTexts({ + userId: ctx.userId, + tenantId, + conversationId: message.conversationId, + messageIds: [message.messageId], + }); + const row = rows.find( + (candidate) => + candidate.messageId === message.messageId && + candidate.text === message.text && + candidate.privacyRevision && + candidate.privateText, + ); + if (row == null) { + throw unavailable(); + } + if (finalEvent?.requestMessage != null) { + finalEvent.requestMessage.privacyRevision = row.privacyRevision; + } + return true; } if (!message.messageId || !message.conversationId || typeof message.text !== 'string') { throw unavailable(); @@ -334,7 +420,7 @@ export async function saveAbortedUserMessage( conversationId: message.conversationId, messageId: message.messageId, text: message.text, - privacyRevision: message.privacyRevision, + privacyRevision: revision, }); if (!persisted) { throw unavailable(); diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index a54ee838beb..ffe87a27699 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -4797,6 +4797,7 @@ class GenerationJobManagerClass { conversationId: jobData.conversationId, text: jobData.userMessage.text ?? '', quotes: jobData.userMessage.quotes, + privacyRevision: jobData.userMessage.privacyRevision, isCreatedByUser: true, } : null, @@ -7916,6 +7917,7 @@ class GenerationJobManagerClass { conversationId: message.conversationId, text: message.text, quotes: message.quotes, + privacyRevision: message.privacyRevision, // Persist the turn's uploaded files so a HITL resume sources them from the job // (this authoritative writer), not a user DB row whose save can still be racing // the approval prompt. diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index f85cb84d531..6f71729f773 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -637,6 +637,54 @@ describe('GenerationJobManager startup telemetry', () => { await manager.destroy(); }); + it('carries only canonical protected text and revision from created through abort', async () => { + const manager = createManager(); + const streamId = 'stream-protected-stop'; + const text = 'Email [EMAIL_1_0123456789abcdef0123456789abcdef]'; + const revision = '0123456789abcdef0123456789abcdef'; + const job = await manager.createJob(streamId, 'owner', streamId, { + initialMetadata: { + responseMessageId: 'response-1', + userMessage: { messageId: 'user-1', conversationId: streamId, text, privacyRevision: revision }, + }, + }); + const created: ServerSentEvent = { + created: true, + streamId, + message: { + messageId: 'user-1', + conversationId: streamId, + sender: 'User', + isCreatedByUser: true, + text, + privacyRevision: revision, + }, + }; + try { + await manager.emitChunk(streamId, created, { expectedCreatedAt: job.createdAt }); + expect((await manager.getJob(streamId))?.metadata.userMessage).toMatchObject({ + messageId: 'user-1', + text, + privacyRevision: revision, + }); + const result = await manager.abortJob(streamId, { + expectedCreatedAt: job.createdAt, + beforePublish: async (pending) => { + expect(pending.jobData?.userMessage).toMatchObject({ text, privacyRevision: revision }); + expect(pending.finalEvent).toMatchObject({ + requestMessage: { text, privacyRevision: revision }, + }); + }, + }); + expect(result.finalEvent).toMatchObject({ + requestMessage: { messageId: 'user-1', text, privacyRevision: revision }, + }); + expect(JSON.stringify(result)).not.toContain('alice@example.com'); + } finally { + await manager.destroy(); + } + }); + it('ends an active startup when the manager shuts down', async () => { const manager = createManager(); const telemetry = createTelemetry(); diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index 2e7d471b360..5f1681979c4 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -153,6 +153,7 @@ export interface SerializableJobData { parentMessageId?: string; conversationId?: string; text?: string; + privacyRevision?: string; /** Quoted excerpts referenced on this turn, carried so resumable/aborted * reconstructions of the user message keep their `MessageQuotes`. */ quotes?: string[]; diff --git a/packages/api/src/types/events.ts b/packages/api/src/types/events.ts index 90194b9e51f..c32b11b1481 100644 --- a/packages/api/src/types/events.ts +++ b/packages/api/src/types/events.ts @@ -14,6 +14,7 @@ export type CreatedEvent = { parentMessageId?: string; conversationId?: string; text?: string; + privacyRevision?: string; sender: string; isCreatedByUser: boolean; /** Quoted excerpts referenced on this turn, carried through resumable job diff --git a/packages/data-provider/src/types/agents.ts b/packages/data-provider/src/types/agents.ts index 98f9591cb6e..c14614a7c8e 100644 --- a/packages/data-provider/src/types/agents.ts +++ b/packages/data-provider/src/types/agents.ts @@ -282,6 +282,8 @@ export namespace Agents { /** User message metadata for rebuilding submission on reconnect */ export interface UserMessageMeta { + /** Canonical, nonsecret revision used to verify owner-only private text. */ + privacyRevision?: string; messageId: string; parentMessageId?: string; conversationId?: string; diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index 4d7e3a13137..ddc7af4e550 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -121,6 +121,53 @@ it('stores both views atomically and excludes ciphertext from ordinary and clien }); }); +it('inserts a missing Stop prerequisite once without overwriting a stored or concurrent protected row', async () => { + await tenant('tenant-a', async () => { + const conversationId = uuid(); + const messageId = uuid(); + const context = { userId: 'owner', expiredAt: new Date(Date.now() + 60_000) }; + const user = { + messageId, + conversationId, + isCreatedByUser: true, + text: '[EMAIL_1_turn]', + }; + const inserted = await methods.saveMessage(context, user, { insertOnly: true }); + expect(inserted?.text).toBe(user.text); + expect(inserted?.expiredAt).toEqual(context.expiredAt); + await methods.saveMessage(context, user, { + privateText: { envelope: 'v1:owner', revision: 'turn' }, + }); + const retry = await methods.saveMessage(context, user, { insertOnly: true }); + expect(retry?.privacyRevision).toBe('turn'); + expect(retry).not.toHaveProperty('privateText'); + const stored = await mongoose.models.Message.findOne({ messageId }) + .select('+privateText') + .lean(); + expect(stored).toMatchObject({ + text: user.text, + privateText: 'v1:owner', + privacyRevision: 'turn', + }); + + const concurrentId = uuid(); + await Promise.all([ + methods.saveMessage(context, { ...user, messageId: concurrentId }, { insertOnly: true }), + methods.saveMessage( + context, + { ...user, messageId: concurrentId }, + { + privateText: { envelope: 'v1:concurrent', revision: 'turn' }, + }, + ), + ]); + const concurrent = await mongoose.models.Message.findOne({ messageId: concurrentId }) + .select('+privateText') + .lean(); + expect(concurrent).toMatchObject({ privateText: 'v1:concurrent', privacyRevision: 'turn' }); + }); +}); + it('does not accept sidecar writes from message parameters or generic edits', async () => { await tenant('tenant-a', async () => { const messageId = uuid(); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index b0395ef9155..6e71c117244 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -698,7 +698,7 @@ export interface MessageMethods { newMessageId?: string; contextMeta?: IMessage['contextMeta'] | null; }, - metadata?: { context?: string; privateText?: PrivateTextWrite }, + metadata?: { context?: string; privateText?: PrivateTextWrite; insertOnly?: boolean }, ): Promise; /** * Reads the references a trace viewer needs for one of the user's @@ -964,7 +964,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa /** `null` unsets a previously stored value; omission leaves it in place. */ contextMeta?: IMessage['contextMeta'] | null; }, - metadata?: { context?: string; privateText?: PrivateTextWrite }, + metadata?: { context?: string; privateText?: PrivateTextWrite; insertOnly?: boolean }, ) { if (!userId) { throw new Error('User not authenticated'); @@ -1075,6 +1075,27 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa params.isCreatedByUser === false && params.isUserSubmitted === undefined; const hasProvenance = userSubmittedPaths.length > 0 || userSubmittedMessageFieldPaths.length > 0; + if (metadata?.insertOnly === true) { + if (metadata.privateText != null) { + throw new Error('A private message cannot use the ordinary insert-only writer.'); + } + const now = new Date(); + const existingOrInserted = await Message.findOneAndUpdate( + { messageId: params.messageId, user: userId }, + { + $setOnInsert: { + ...update, + ...(userSubmittedPaths.length > 0 && { userSubmittedPaths }), + ...(userSubmittedMessageFieldPaths.length > 0 && { userSubmittedMessageFieldPaths }), + ...retentionOnInsert, + createdAt: now, + updatedAt: now, + }, + }, + { upsert: true, new: true, timestamps: false }, + ); + return existingOrInserted?.toObject(); + } const message = hasProvenance ? await findOneAndMergeMessageProvenance( Message, From 0cf0686ec6c426cbafe52842c89e38eaf8b61831 Mon Sep 17 00:00:00 2001 From: Lia Date: Mon, 28 Sep 2026 01:46:07 +0000 Subject: [PATCH 06/18] test: cover preliminary private metadata in disconnect fixture --- .../agents/__tests__/request.partialDisconnect.spec.js | 1 + packages/api/src/stream/__tests__/startup.spec.ts | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 01bdb323379..a36aa3cb70f 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -48,6 +48,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ savePrivateTextMessage: (save, _req, ...args) => save(...args), + stampPreliminaryPrivateTextMessage: (_req, message) => message, getAgentErrorMetadata: (...args) => jest.requireActual('@librechat/api').getAgentErrorMetadata(...args), sendEvent: jest.fn(), diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index 6f71729f773..90197a343e5 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -645,7 +645,12 @@ describe('GenerationJobManager startup telemetry', () => { const job = await manager.createJob(streamId, 'owner', streamId, { initialMetadata: { responseMessageId: 'response-1', - userMessage: { messageId: 'user-1', conversationId: streamId, text, privacyRevision: revision }, + userMessage: { + messageId: 'user-1', + conversationId: streamId, + text, + privacyRevision: revision, + }, }, }); const created: ServerSentEvent = { From 7696f792b40d2a643dfe7bb7fd337c7f80af2630 Mon Sep 17 00:00:00 2001 From: Lia Date: Mon, 28 Sep 2026 02:01:52 +0000 Subject: [PATCH 07/18] test: carry preliminary metadata in resume controller fixture --- .../controllers/agents/__tests__/request.resumeMetadata.spec.js | 1 + 1 file changed, 1 insertion(+) diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index ea72a8aedf4..658c543a4f6 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -253,6 +253,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ savePrivateTextMessage: (save, _req, ...args) => save(...args), + stampPreliminaryPrivateTextMessage: (_req, message) => message, getSteerRecoveryFailure: jest.requireActual( '../../../../../packages/api/src/stream/SteerRecovery', ).getSteerRecoveryFailure, From 22734756b250eb9aad664dd2dcf1a64658163f71 Mon Sep 17 00:00:00 2001 From: Lia Date: Mon, 28 Sep 2026 02:10:25 +0000 Subject: [PATCH 08/18] test: preserve Stop prerequisite identity in tenant fixture --- api/server/routes/agents/__tests__/streamTenant.spec.js | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/api/server/routes/agents/__tests__/streamTenant.spec.js b/api/server/routes/agents/__tests__/streamTenant.spec.js index e23338c937f..94c7710ced9 100644 --- a/api/server/routes/agents/__tests__/streamTenant.spec.js +++ b/api/server/routes/agents/__tests__/streamTenant.spec.js @@ -89,7 +89,10 @@ describe('SSE stream tenant isolation', () => { checkpointIds: ['checkpoint-a'], }); mockDeleteAgentCheckpoint.mockResolvedValue(undefined); - mockSaveMessage.mockResolvedValue({ persisted: true }); + mockSaveMessage.mockImplementation(async (_context, message) => ({ + ...message, + persisted: true, + })); mockGenerationJobManager.getActiveJobIdsForUser.mockResolvedValue([]); mockGenerationJobManager.steering.claim.mockResolvedValue([]); mockGenerationJobManager.steering.claimDetailed.mockResolvedValue({ @@ -982,7 +985,7 @@ describe('SSE stream tenant isolation', () => { 1, expect.any(Object), expect.objectContaining({ messageId: 'user-1', isCreatedByUser: true }), - expect.any(Object), + expect.objectContaining({ insertOnly: true }), ); expect(mockSaveMessage).toHaveBeenNthCalledWith( 2, From ce8cd951dcfeb6501b759b0d8040742685418aee Mon Sep 17 00:00:00 2001 From: Lia Date: Mon, 28 Sep 2026 18:00:38 +0000 Subject: [PATCH 09/18] fix: keep protected owner text out of search and exports --- api/server/middleware/denyRequest.js | 27 +-- .../denyRequest.privateText.spec.js | 50 +++++ .../routes/__tests__/messages-get.spec.js | 7 + api/server/routes/agents/chat.js | 2 + api/server/routes/messages.js | 10 +- .../Chat/Messages/PrivateText.spec.tsx | 36 ++++ .../components/Chat/Messages/PrivateText.tsx | 22 ++- .../__tests__/useExportConversation.spec.tsx | 179 ++++++++++++++++++ .../Conversations/useExportConversation.ts | 83 ++++++-- client/src/locales/en/translation.json | 2 + client/src/utils/messages.spec.ts | 34 +++- client/src/utils/messages.ts | 10 + e2e/specs/mock/private-text.spec.ts | 139 ++++++++++++++ .../api/src/middleware/messageFilterPii.ts | 14 +- .../src/protection/private/submission.spec.ts | 57 ++++++ .../api/src/protection/private/submission.ts | 16 +- packages/api/src/protection/private/view.ts | 8 + .../src/models/plugins/mongoMeili.spec.ts | 25 +++ .../src/models/plugins/mongoMeili.ts | 2 +- 19 files changed, 678 insertions(+), 45 deletions(-) create mode 100644 client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx diff --git a/api/server/middleware/denyRequest.js b/api/server/middleware/denyRequest.js index 8f4adf7b31b..2d45372244e 100644 --- a/api/server/middleware/denyRequest.js +++ b/api/server/middleware/denyRequest.js @@ -37,18 +37,22 @@ const denyRequest = async (req, res, errorMessage) => { const { messageId, conversationId: _convoId, parentMessageId, text } = req.body; const conversationId = _convoId ?? crypto.randomUUID(); + const shouldSaveMessage = Boolean( + _convoId && parentMessageId && parentMessageId !== Constants.NO_PARENT, + ); - const userMessage = stampPrivateTextMessage(req, { - sender: 'User', - messageId: messageId ?? crypto.randomUUID(), - parentMessageId, - conversationId, - isCreatedByUser: true, - text, - }); - sendEvent(res, { message: userMessage, created: true }); - - const shouldSaveMessage = _convoId && parentMessageId && parentMessageId !== Constants.NO_PARENT; + const userMessage = stampPrivateTextMessage( + req, + { + sender: 'User', + messageId: messageId ?? crypto.randomUUID(), + parentMessageId, + conversationId, + isCreatedByUser: true, + text, + }, + shouldSaveMessage, + ); if (shouldSaveMessage) { await savePrivateTextMessage( @@ -65,6 +69,7 @@ const denyRequest = async (req, res, errorMessage) => { ); } + sendEvent(res, { message: userMessage, created: true }); return await sendError(req, res, { sender: getResponseSender(req.body), messageId: crypto.randomUUID(), diff --git a/api/server/middleware/denyRequest.privateText.spec.js b/api/server/middleware/denyRequest.privateText.spec.js index 9e33d36ec6d..41084da62a0 100644 --- a/api/server/middleware/denyRequest.privateText.spec.js +++ b/api/server/middleware/denyRequest.privateText.spec.js @@ -68,6 +68,9 @@ it('encrypts denied PII for an existing conversation and sends only the filtered const event = mockSendEvent.mock.calls[0][1]; const [ctx, saved, metadata] = mockSaveMessage.mock.calls[0]; expect(mockSaveMessage).toHaveBeenCalledTimes(1); + expect(mockSaveMessage.mock.invocationCallOrder[0]).toBeLessThan( + mockSendEvent.mock.invocationCallOrder[0], + ); expect(ctx.userId).toBe('owner'); expect(saved.text).toMatch(/^Email \[EMAIL_1_[a-f0-9]{32}\]$/); expect(saved.privacyRevision).toBe(event.message.privacyRevision); @@ -82,6 +85,53 @@ it('encrypts denied PII for an existing conversation and sends only the filtered ); }); +it('does not advertise an original on a denied first turn that is never persisted', async () => { + const req = { + method: 'POST', + path: '/', + originalUrl: '/api/agents/chat', + config: { + filters: { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }, + }, + user: { id: 'owner', tenantId: 'tenant-a' }, + body: { text: `Email ${original}`, clientRequestId: uuidv4() }, + }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const next = jest.fn(); + createPrivateTextIngress({ + getFilters: () => req.config.filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, res, next); + expect(next).toHaveBeenCalledTimes(1); + await denyRequest(req, res, { type: 'ban' }); + + expect(mockSaveMessage).not.toHaveBeenCalled(); + expect(mockSendEvent).toHaveBeenCalledTimes(1); + expect(mockSendEvent.mock.calls[0][1].message).toMatchObject({ + text: expect.stringMatching(/^Email \[EMAIL_1_[a-f0-9]{32}\]$/), + isCreatedByUser: true, + }); + expect(mockSendEvent.mock.calls[0][1].message).not.toHaveProperty('privacyRevision'); + expect(JSON.stringify(mockSendEvent.mock.calls)).not.toContain(original); + expect(mockSendError).toHaveBeenCalledWith( + req, + res, + expect.objectContaining({ shouldSaveMessage: false }), + ); +}); + it('fails closed before event or storage when denial hits untransformed private text', async () => { const req = { path: '/', diff --git a/api/server/routes/__tests__/messages-get.spec.js b/api/server/routes/__tests__/messages-get.spec.js index f376a9f0dce..bac78f10468 100644 --- a/api/server/routes/__tests__/messages-get.spec.js +++ b/api/server/routes/__tests__/messages-get.spec.js @@ -55,6 +55,9 @@ jest.mock('@librechat/api', () => { withoutTraceRefs: jest.requireActual('../../../../packages/api/src/langfuse/trace.ts') .withoutTraceRefs, createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), + stripPrivateMessageFields: jest.requireActual( + '../../../../packages/api/src/protection/private/view', + ).stripPrivateMessageFields, createContentFilter: jest.fn(() => (req, res, next) => next()), inspectContent, extractChatContent, @@ -641,6 +644,8 @@ describe('message route conversation ownership filters', () => { messageId: 'hit-1', conversationId: 'convo-1', text: 'needle in a haystack', + privateText: 'v1:encrypted-original', + privacyRevision: 'public-revision', contextMeta: { calibrationRatio: 1.2, encoding: 'claude', @@ -662,6 +667,8 @@ describe('message route conversation ownership filters', () => { expect(response.body.messages).toHaveLength(1); expect(response.body.messages[0]).toMatchObject({ messageId: 'hit-1', title: 'Found' }); expect(response.body.messages[0]).not.toHaveProperty('contextMeta'); + expect(response.body.messages[0]).not.toHaveProperty('privateText'); + expect(response.body.messages[0].privacyRevision).toBe('public-revision'); }); it('returns indistinguishable not-found responses for child and missing query reads', async () => { diff --git a/api/server/routes/agents/chat.js b/api/server/routes/agents/chat.js index 16ecf54fd04..b9c9770089a 100644 --- a/api/server/routes/agents/chat.js +++ b/api/server/routes/agents/chat.js @@ -2,6 +2,7 @@ const express = require('express'); const { logger } = require('@librechat/data-schemas'); const { createMessageFilterPii, + getPreinspectedPrivateText, reportLocatorTraversalFailure, generateCheckAccess, skipAgentCheck, @@ -76,6 +77,7 @@ router.use( onTraversalFailure: reportLocatorTraversalFailure, getConfig: (req) => req.config?.messageFilter?.pii, getFilters: (req) => req.config?.filters, + getPreinspectedText: getPreinspectedPrivateText, getFiles, }), ); diff --git a/api/server/routes/messages.js b/api/server/routes/messages.js index fb56d950363..6797595251b 100644 --- a/api/server/routes/messages.js +++ b/api/server/routes/messages.js @@ -28,6 +28,7 @@ const { isContentFilterError, withoutTraceRefs, createPrivateTextView, + stripPrivateMessageFields, } = require('@librechat/api'); const subagentThreadTaskStore = require('~/server/services/Endpoints/agents/subagentThreadStore'); const { findAllArtifacts, replaceArtifactContent } = require('~/server/services/Artifacts/update'); @@ -199,9 +200,8 @@ router.get('/', async (req, res) => { for (const message of cleanedMessages) { const convo = result.convoMap[message.conversationId]; const dbMessage = dbMessageMap[message.messageId]; - /** Search hydrates every schema field; server-private state never leaves. */ - const publicHit = { ...message }; - delete publicHit.contextMeta; + /** Search may hydrate server-private fields; only a public projection leaves. */ + const publicHit = stripPrivateMessageFields(message); activeMessages.push({ ...publicHit, @@ -243,9 +243,7 @@ router.get('/', async (req, res) => { * @returns {TMessage} */ function toClientMessage(message) { - const clientMessage = { ...message }; - delete clientMessage.contextMeta; - return clientMessage; + return stripPrivateMessageFields(message); } router.post('/branch', configMiddleware, async (req, res) => { diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index b5f6a23bb84..76754777c99 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -77,6 +77,42 @@ it('renders originals without mutating canonical model/export input, and reloads expect(load).toHaveBeenCalledTimes(2); }); +it('loads the acknowledged first-turn owner text before navigating away from /new', async () => { + const firstTurn = { + ...canonical, + conversationId: '11111111-1111-4111-8111-111111111111', + }; + load.mockResolvedValue({ messages: [{ ...original, canonicalText: firstTurn.text }] }); + const view = render(); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledWith(firstTurn.conversationId, ['message']); + expect(load).not.toHaveBeenCalledWith('new', expect.anything()); + + view.rerender( + , + ); + expect(screen.getByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(1); +}); + +it('never mixes protected rows from stale conversations into a first-turn owner read', async () => { + const old = { + ...canonical, + messageId: 'old-message', + conversationId: '11111111-1111-4111-8111-111111111111', + }; + const fresh = { + ...canonical, + messageId: 'fresh-message', + conversationId: '22222222-2222-4222-8222-222222222222', + }; + load.mockResolvedValue({ messages: [{ ...original, messageId: 'fresh-message' }] }); + render(); + expect(await screen.findByText(original.text)).toBeInTheDocument(); + expect(load).toHaveBeenCalledTimes(1); + expect(load).toHaveBeenCalledWith(fresh.conversationId, ['fresh-message']); +}); + it('does not fetch originals for an ordinary transcript', () => { const plain = { ...canonical, privacyRevision: undefined }; render( diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index f1d4873266e..112f9f87d3e 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -44,10 +44,19 @@ interface OwnerTextProviderProps { } export function OwnerTextProvider(props: OwnerTextProviderProps) { - if (!props.messages?.some((message) => message.isCreatedByUser && message.privacyRevision)) { + const protectedMessage = props.messages?.findLast( + (message) => message.isCreatedByUser && message.privacyRevision && message.conversationId, + ); + if (protectedMessage == null) { return <>{props.children}; } - return ; + const conversationId = + props.conversationId === 'new' && + typeof protectedMessage.conversationId === 'string' && + protectedMessage.conversationId !== 'new' + ? protectedMessage.conversationId + : props.conversationId; + return ; } function ActiveOwnerTextProvider({ @@ -61,11 +70,16 @@ function ActiveOwnerTextProvider({ () => JSON.stringify( (messages ?? []) - .filter((message) => message.isCreatedByUser && message.privacyRevision) + .filter( + (message) => + message.isCreatedByUser && + message.privacyRevision && + message.conversationId === conversationId, + ) .map((message) => [message.messageId, message.privacyRevision, message.text]) .sort(), ), - [messages], + [messages, conversationId], ); const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); const [state, setState] = useState(empty); diff --git a/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx b/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx new file mode 100644 index 00000000000..af3dfca614e --- /dev/null +++ b/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx @@ -0,0 +1,179 @@ +import { Provider } from 'jotai'; +import download from 'downloadjs'; +import exportFromJSON from 'export-from-json'; +import { act, renderHook } from '@testing-library/react'; +import { MemoryRouter, Route, Routes } from 'react-router-dom'; +import { Constants, QueryKeys } from 'librechat-data-provider'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import type { TConversation, TMessage } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import useExportConversation from '../useExportConversation'; + +const mockGetMessages = jest.fn(); +const mockShowToast = jest.fn(); +const mockCaptureScreenshot = jest.fn(); + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + getMessagesByConvoId: (...args: unknown[]) => mockGetMessages(...args), + }, + }; +}); +jest.mock('@librechat/client', () => ({ + ...jest.requireActual('@librechat/client'), + useToastContext: () => ({ showToast: mockShowToast }), +})); +jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); +jest.mock('~/hooks/ScreenshotContext', () => ({ + useScreenshot: () => ({ captureScreenshot: mockCaptureScreenshot }), + ScreenshotLimitError: class ScreenshotLimitError extends Error {}, +})); +jest.mock('downloadjs', () => jest.fn()); +jest.mock('export-from-json', () => + Object.assign(jest.fn(), { types: { csv: 'csv', txt: 'txt' } }), +); + +const conversationId = '11111111-1111-4111-8111-111111111111'; +const conversation = { conversationId, title: 'Protected chat' } as TConversation; +const submitted = { + conversationId, + messageId: 'user-1', + parentMessageId: Constants.NO_PARENT, + isCreatedByUser: true, + clientTimestamp: '2026-09-28T15:00:00', + text: 'Email alice@example.com', +} as TMessage; +const canonical = { + ...submitted, + text: 'Email [EMAIL_1_revision]', + privacyRevision: 'revision', +} as TMessage; + +function setup(type: string, pending: TMessage[] = [submitted]) { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + queryClient.setQueryData([QueryKeys.messages, conversationId], pending); + const hook = renderHook( + () => + useExportConversation({ + conversation, + filename: 'protected', + type, + includeOptions: false, + exportBranches: false, + recursive: false, + }), + { + wrapper: function Wrapper({ children }: { children: ReactNode }) { + return ( + + + + + + + + + + ); + }, + }, + ); + return { ...hook, queryClient }; +} + +async function readBlob(blob: Blob): Promise { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(String(reader.result)); + reader.onerror = () => reject(reader.error); + reader.readAsText(blob); + }); +} + +beforeEach(() => { + jest.clearAllMocks(); +}); + +it('exports acknowledged canonical text rather than a pending private cache value', async () => { + mockGetMessages.mockResolvedValueOnce([canonical]); + const { result } = setup('json'); + await act(async () => { + await result.current.exportConversation(); + }); + expect(mockGetMessages).toHaveBeenCalledWith(conversationId); + const saved = (download as jest.Mock).mock.calls[0]?.[0] as Blob; + expect(saved).toBeInstanceOf(Blob); + const fileText = await readBlob(saved); + expect(fileText).toContain(canonical.text); + expect(fileText).not.toContain('alice@example.com'); + expect(mockShowToast).not.toHaveBeenCalled(); +}); + +it.each(['csv', 'markdown', 'text'])( + 'uses the canonical server response for %s exports', + async (type) => { + mockGetMessages.mockResolvedValueOnce([canonical]); + const { result } = setup(type); + await act(async () => { + await result.current.exportConversation(); + }); + expect(mockGetMessages).toHaveBeenCalledWith(conversationId); + expect(exportFromJSON).toHaveBeenCalledTimes(1); + const payload = JSON.stringify(jest.mocked(exportFromJSON).mock.calls[0][0]); + expect(payload).toContain(canonical.text); + expect(payload).not.toContain('alice@example.com'); + }, +); + +it('does not fall back to the unfiltered cache when the canonical read fails', async () => { + mockGetMessages.mockRejectedValueOnce(new Error('Temporary outage')); + const { result } = setup('json'); + await act(async () => { + await result.current.exportConversation(); + }); + expect(download).not.toHaveBeenCalled(); + expect(mockShowToast).toHaveBeenCalledWith( + expect.objectContaining({ message: 'com_nav_export_unavailable' }), + ); +}); + +it('does not download a screenshot when a protected row arrives during capture', async () => { + const clean = { ...canonical, privacyRevision: undefined, createdAt: '2026-09-28T15:00:01' }; + let finish!: (result: Blob) => void; + mockCaptureScreenshot.mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result, queryClient } = setup('screenshot', [clean]); + const exportAction = result.current.exportConversation(); + expect(mockCaptureScreenshot).toHaveBeenCalledTimes(1); + act(() => queryClient.setQueryData([QueryKeys.messages, conversationId], [canonical])); + finish(new Blob(['captured original'], { type: 'image/png' })); + await act(async () => { + await exportAction; + }); + expect(download).not.toHaveBeenCalled(); + expect(mockShowToast).toHaveBeenCalledWith( + expect.objectContaining({ message: 'com_nav_export_screenshot_private_text' }), + ); +}); + +it.each([ + { kind: 'unsent', message: submitted }, + { kind: 'protected', message: canonical }, +])('does not screenshot $kind text', async ({ message }) => { + const { result } = setup('screenshot', [message]); + await act(async () => { + await result.current.exportConversation(); + }); + expect(mockCaptureScreenshot).not.toHaveBeenCalled(); + expect(download).not.toHaveBeenCalled(); + expect(mockShowToast).toHaveBeenCalledWith( + expect.objectContaining({ message: 'com_nav_export_screenshot_private_text' }), + ); +}); diff --git a/client/src/hooks/Conversations/useExportConversation.ts b/client/src/hooks/Conversations/useExportConversation.ts index 2ea3d31a163..80900ebda53 100644 --- a/client/src/hooks/Conversations/useExportConversation.ts +++ b/client/src/hooks/Conversations/useExportConversation.ts @@ -4,10 +4,11 @@ import { useParams } from 'react-router-dom'; import exportFromJSON from 'export-from-json'; import { useToastContext } from '@librechat/client'; import { useQueryClient } from '@tanstack/react-query'; -import { buildTree, QueryKeys } from 'librechat-data-provider'; +import { buildTree, Constants, dataService, QueryKeys } from 'librechat-data-provider'; import type { TConversation, TMessage, TPreset } from 'librechat-data-provider'; import { ScreenshotLimitError, useScreenshot } from '~/hooks/ScreenshotContext'; import useBuildMessageTree from '~/hooks/Messages/useBuildMessageTree'; +import { isUnacknowledgedUserMessage } from '~/utils/messages'; import { NotificationSeverity } from '~/common'; import { formatMessageText } from './format'; import { cleanupPreset } from '~/utils'; @@ -42,15 +43,48 @@ export default function useExportConversation({ const { conversationId: paramId } = useParams(); - const getMessageTree = useCallback(() => { + const getCachedMessages = useCallback(() => { const queryParam = paramId === 'new' ? paramId : (conversation?.conversationId ?? paramId ?? ''); - const messages = queryClient.getQueryData([QueryKeys.messages, queryParam]) ?? []; + return queryClient.getQueryData([QueryKeys.messages, queryParam]) ?? []; + }, [paramId, conversation?.conversationId, queryClient]); + + const getMessageTree = useCallback(async () => { + const conversationId = conversation?.conversationId; + if ( + !conversationId || + conversationId === Constants.NEW_CONVO || + conversationId === Constants.PENDING_CONVO || + conversationId === Constants.SEARCH + ) { + throw new Error('Conversation is unavailable for export'); + } + const messages = await dataService.getMessagesByConvoId(conversationId); const dataTree = buildTree({ messages }); return dataTree?.length === 0 ? null : (dataTree ?? null); - }, [paramId, conversation?.conversationId, queryClient]); + }, [conversation?.conversationId]); + + const screenshotWouldExposePrivateText = () => + getCachedMessages().some( + (message) => message.privacyRevision != null || isUnacknowledgedUserMessage(message), + ); + + const refuseUnsafeScreenshot = () => { + if (!screenshotWouldExposePrivateText()) { + return false; + } + showToast({ + message: localize('com_nav_export_screenshot_private_text'), + severity: NotificationSeverity.ERROR, + showIcon: true, + }); + return true; + }; const exportScreenshot = async () => { + if (refuseUnsafeScreenshot()) { + return; + } let data: Blob; try { data = await captureScreenshot(); @@ -67,6 +101,9 @@ export default function useExportConversation({ }); return; } + if (refuseUnsafeScreenshot()) { + return; + } download(data, `${filename}.png`, 'image/png'); }; @@ -76,7 +113,7 @@ export default function useExportConversation({ const messages = await buildMessageTree({ messageId: conversation?.conversationId, message: null, - messages: getMessageTree(), + messages: await getMessageTree(), branches: Boolean(exportBranches), recursive: false, }); @@ -154,7 +191,7 @@ export default function useExportConversation({ const messages = await buildMessageTree({ messageId: conversation?.conversationId, message: null, - messages: getMessageTree(), + messages: await getMessageTree(), branches: false, recursive: false, }); @@ -210,7 +247,7 @@ export default function useExportConversation({ const messages = await buildMessageTree({ messageId: conversation?.conversationId, message: null, - messages: getMessageTree(), + messages: await getMessageTree(), branches: false, recursive: false, }); @@ -262,7 +299,7 @@ export default function useExportConversation({ const messages = await buildMessageTree({ messageId: conversation?.conversationId, message: null, - messages: getMessageTree(), + messages: await getMessageTree(), branches: Boolean(exportBranches), recursive: Boolean(recursive), }); @@ -279,17 +316,25 @@ export default function useExportConversation({ download(blob, `${filename}.json`, 'application/json'); }; - const exportConversation = () => { - if (type === 'json') { - exportJSON(); - } else if (type == 'text') { - exportText(); - } else if (type == 'markdown') { - exportMarkdown(); - } else if (type == 'csv') { - exportCSV(); - } else if (type == 'screenshot') { - exportScreenshot(); + const exportConversation = async () => { + try { + if (type === 'json') { + await exportJSON(); + } else if (type == 'text') { + await exportText(); + } else if (type == 'markdown') { + await exportMarkdown(); + } else if (type == 'csv') { + await exportCSV(); + } else if (type == 'screenshot') { + await exportScreenshot(); + } + } catch { + showToast({ + message: localize('com_nav_export_unavailable'), + severity: NotificationSeverity.ERROR, + showIcon: true, + }); } }; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index b78997b8f0c..11a256d368f 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -735,8 +735,10 @@ "com_nav_export_recursive": "Recursive", "com_nav_export_recursive_or_sequential": "Recursive or sequential?", "com_nav_export_screenshot_error": "Couldn't capture the screenshot. Try a different export type.", + "com_nav_export_screenshot_private_text": "Screenshots can't be exported while a chat contains protected or unsent text. Export as text or JSON instead.", "com_nav_export_screenshot_too_large": "This conversation is too large to export as a screenshot. Try a different export type.", "com_nav_export_type": "Type", + "com_nav_export_unavailable": "Couldn't load this conversation for export. Try again.", "com_nav_external": "External", "com_nav_font_size": "Message Font Size", "com_nav_font_size_base": "Medium", diff --git a/client/src/utils/messages.spec.ts b/client/src/utils/messages.spec.ts index b06f3e5a485..33f17b5d38a 100644 --- a/client/src/utils/messages.spec.ts +++ b/client/src/utils/messages.spec.ts @@ -1,6 +1,10 @@ import { ContentTypes } from 'librechat-data-provider'; import type { TMessage, TMessageContentParts } from 'librechat-data-provider'; -import { preserveStreamedContentIdentity, stripStreamedIndexStamps } from './messages'; +import { + preserveStreamedContentIdentity, + isUnacknowledgedUserMessage, + stripStreamedIndexStamps, +} from './messages'; const text = (value: string, extra: Record = {}): TMessageContentParts => ({ type: ContentTypes.TEXT, text: value, ...extra }) as TMessageContentParts; @@ -20,6 +24,34 @@ const label = (value: string, extra: Record = {}): TMessageCont const streamedIndexes = (content: TMessage['content']): Array => (content ?? []).map((part) => part?.streamedIndex); +describe('private screenshot admission', () => { + it('identifies unsent user text without classifying acknowledged messages as pending', () => { + const optimistic = { + messageId: 'pending', + isCreatedByUser: true, + clientTimestamp: '2026-09-28T15:00:00', + text: 'alice@example.com', + } as TMessage; + const assistant = { + messageId: 'previous-answer', + isCreatedByUser: false, + text: 'Safe', + } as TMessage; + const canonical = { + ...optimistic, + privacyRevision: 'accepted-revision', + text: '[EMAIL_1_accepted-revision]', + }; + const persisted = { ...optimistic, createdAt: '2026-09-28T15:00:01' }; + + expect(isUnacknowledgedUserMessage(optimistic)).toBe(true); + expect(isUnacknowledgedUserMessage(canonical)).toBe(false); + expect(isUnacknowledgedUserMessage(persisted)).toBe(false); + expect(isUnacknowledgedUserMessage(assistant)).toBe(false); + expect(optimistic.text).toBe('alice@example.com'); + }); +}); + describe('preserveStreamedContentIdentity', () => { it('stamps every part shifted by compacted holes with its streamed index', () => { const streamed = [ diff --git a/client/src/utils/messages.ts b/client/src/utils/messages.ts index 2c59cee4a4b..6ef5efc344f 100644 --- a/client/src/utils/messages.ts +++ b/client/src/utils/messages.ts @@ -21,6 +21,16 @@ import type { LocalizeFunction, TMessageProps } from '~/common'; export const TEXT_KEY_DIVIDER = '|||'; export const STREAM_START_FAILED_METADATA_KEY = 'streamStartFailed'; +/** A locally submitted user row is not canonical until the server acknowledges it. */ +export function isUnacknowledgedUserMessage(message: TMessage): boolean { + return ( + message.isCreatedByUser === true && + message.clientTimestamp != null && + message.createdAt == null && + !message.privacyRevision + ); +} + type SiblingIndexLookup = (parentMessageId: string | null | undefined) => number; export type BranchSiblingIndex = { diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index bbbfd4080e1..94dc7ef37cd 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -1,3 +1,4 @@ +import { readFile } from 'node:fs/promises'; import { expect, test } from '@playwright/test'; import type { TMessage } from 'librechat-data-provider'; import { withMongo } from './db'; @@ -10,6 +11,7 @@ import { import { MOCK_ENDPOINTS, selectMockEndpoint, + sendMessage, sendMessageAndWaitForCompletion, messagesView, fetchJson, @@ -108,6 +110,89 @@ test('owner sees original after reload while provider, sharing, and canonical re expect(shared.text).not.toContain('privateText'); expect(shared.text).toContain('EMAIL_1_'); + const openExport = async () => { + await page.getByRole('button', { name: 'Export/Share' }).click(); + await page.getByRole('menuitem', { name: 'Export' }).click(); + return page.getByRole('dialog', { name: 'Export conversation' }); + }; + const selectType = async (label: string) => { + const dialog = page.getByRole('dialog', { name: 'Export conversation' }); + await dialog.getByTestId('dropdown-menu').click(); + await page.getByRole('option', { name: label }).click(); + return dialog; + }; + let exportDialog = await openExport(); + exportDialog = await selectType('json (.json)'); + const [canonicalDownload] = await Promise.all([ + page.waitForEvent('download'), + exportDialog.getByRole('button', { name: 'Export', exact: true }).click(), + ]); + const exported = await readFile(await canonicalDownload.path(), 'utf8'); + expect(exported).not.toContain('alice@example.com'); + expect(exported).toContain('EMAIL_1_'); + + exportDialog = await selectType('screenshot (.png)'); + await exportDialog.getByRole('button', { name: 'Export', exact: true }).click(); + await expect( + page + .getByText("Screenshots can't be exported while a chat contains protected or unsent text.") + .first(), + ).toBeVisible(); + await page.keyboard.press('Escape'); + await expect(exportDialog).toBeHidden(); + + const pending = 'Unacknowledged private export: alice@example.com'; + const delayedPath = '**/api/agents/chat/**'; + let releasePending!: () => void; + const pendingGate = new Promise((resolve) => { + releasePending = resolve; + }); + await page.route(delayedPath, async (route) => { + if (route.request().method() !== 'POST' || route.request().postDataJSON()?.text !== pending) { + await route.continue(); + return; + } + await pendingGate; + await route.continue(); + }); + const pendingRequest = (requestToCheck: { + method(): string; + url(): string; + postDataJSON(): { text?: string }; + }) => + requestToCheck.method() === 'POST' && + new URL(requestToCheck.url()).pathname.startsWith('/api/agents/chat/') && + requestToCheck.postDataJSON()?.text === pending; + const requestSeen = page.waitForRequest(pendingRequest); + const pendingResponse = page.waitForResponse((result) => pendingRequest(result.request())); + try { + const input = page.getByRole('textbox', { name: 'Message input' }); + await input.fill(pending); + await input.press('Enter'); + await requestSeen; + exportDialog = await openExport(); + exportDialog = await selectType('json (.json)'); + const [pendingDownload] = await Promise.all([ + page.waitForEvent('download'), + exportDialog.getByRole('button', { name: 'Export', exact: true }).click(), + ]); + const pendingExport = await readFile(await pendingDownload.path(), 'utf8'); + expect(pendingExport).not.toContain(pending); + expect(pendingExport).not.toContain('alice@example.com'); + expect(pendingExport).toContain('EMAIL_1_'); + } finally { + releasePending(); + await pendingResponse.catch(() => undefined); + await page.unroute(delayedPath); + } + expect((await pendingResponse).ok()).toBe(true); + await expect + .poll( + async () => + (await fetchJson(page, `/api/messages/${conversationId}`, token)).length, + ) + .toBeGreaterThanOrEqual(4); + const unauthorized = await request.post(`/api/messages/${conversationId}/owner-text`, { data: { messageIds: [user.messageId] }, }); @@ -127,3 +212,57 @@ test('owner sees original after reload while provider, sharing, and canonical re } } }); + +test('the first protected owner view loads from its server ID while generation is streaming', async ({ + page, + request, +}) => { + test.setTimeout(90000); + const token = await loginAdmin(request); + await setRuntimeFilters(request, token, { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }); + let conversationId: string | undefined; + const text = 'E2E_SLOW_REPLY:owner-stream alice@example.com'; + try { + await page.goto('/c/new'); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const ownerRead = page.waitForResponse( + (response) => + response.request().method() === 'POST' && + /\/api\/messages\/[^/]+\/owner-text$/.test(new URL(response.url()).pathname), + ); + const started = await sendMessage(page, text); + conversationId = (await started.json()).conversationId as string; + expect(conversationId).toMatch(/^[0-9a-f-]{36}$/); + const owner = await ownerRead; + expect(owner.status()).toBe(200); + expect(new URL(owner.url()).pathname).toBe(`/api/messages/${conversationId}/owner-text`); + expect((await owner.json()).messages[0].text).toBe(text); + await expect( + messagesView(page).getByText('Private details hidden from the model'), + ).toBeVisible(); + await expect(messagesView(page).getByText(text, { exact: true })).toBeVisible(); + await expect(page.getByTestId('stop-generation-button')).toBeVisible(); + await expect(page.getByTestId('stop-generation-button')).toBeHidden({ timeout: 60000 }); + } finally { + await restoreRuntimeFilters(request, token); + if (conversationId) { + await requestResult(request, { + path: '/api/convos', + token, + method: 'DELETE', + data: { arg: { conversationId } }, + }); + } + } +}); diff --git a/packages/api/src/middleware/messageFilterPii.ts b/packages/api/src/middleware/messageFilterPii.ts index a9cd123c052..4d2626e0da0 100644 --- a/packages/api/src/middleware/messageFilterPii.ts +++ b/packages/api/src/middleware/messageFilterPii.ts @@ -147,6 +147,7 @@ export interface CreateMessageFilterPiiOptions { getConfig: (req: ServerRequest) => MessageFilterPiiConfig | undefined; getFilters?: (req: ServerRequest) => FiltersConfig | undefined; getFiles?: GetCanonicalFilesForInspection; + getPreinspectedText?: (req: ServerRequest) => string | undefined; } export function createMessageFilterPii(options: CreateMessageFilterPiiOptions): RequestHandler { @@ -241,7 +242,18 @@ export function createMessageFilterPii(options: CreateMessageFilterPiiOptions): if (filters != null && !collect(() => extractStoredMessageContent(req.body))) { return; } - const finding = inspectContent(fragments, { filters, legacyPii }); + const preinspectedText = options.getPreinspectedText?.(req); + const inspectable = + preinspectedText == null + ? fragments + : fragments.filter( + (fragment) => + fragment.source !== 'message' || + fragment.field !== 'text' || + fragment.path !== '/text' || + fragment.text !== preinspectedText, + ); + const finding = inspectContent(inspectable, { filters, legacyPii }); if (finding != null) { if (finding.detectorId !== 'legacy-pattern') { res.status(400).json(contentFilterBlockResponse(finding)); diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 99cdb07ff84..991a570d804 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -10,8 +10,10 @@ import { saveAbortedUserMessage, isPreDenialTextSubmission, isPrivateTextChatSubmission, + getPreinspectedPrivateText, privateTextBinding, } from './submission'; +import { createMessageFilterPii } from '../../middleware/messageFilterPii'; import { createPrivateTextCipher } from './crypto'; import { createPrivateTextView } from './view'; @@ -200,6 +202,61 @@ describe('private text submission boundary', () => { expect(JSON.stringify({ req, message })).not.toContain('alice@example.com'); }); + it('skips only the verified, preinspected text in the second PII pass', async () => { + const patterns: FiltersConfig = { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'revision', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, + ], + }, + }, + }; + const req = { + path: '/', + user: { id: 'owner', tenantId: 'tenant-a' }, + body: { + text: original, + clientRequestId: 'hex-rule-1', + input: undefined as string | undefined, + }, + } as unknown as Request; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const ingressNext = jest.fn(); + createPrivateTextIngress({ + getFilters: () => patterns, + getLegacyPii: () => undefined, + getKey: () => key, + })(req, res as unknown as Response, ingressNext); + expect(ingressNext).toHaveBeenCalledTimes(1); + expect(req.body.text).toMatch(/\[EMAIL_1_[a-f0-9]{32}\]/); + expect(getPreinspectedPrivateText(req)).toBe(req.body.text); + + const secondPass = createMessageFilterPii({ + getConfig: () => undefined, + getFilters: () => patterns, + getPreinspectedText: getPreinspectedPrivateText, + }); + const next = jest.fn(); + await secondPass(req, res as unknown as Response, next); + expect(next).toHaveBeenCalledTimes(1); + expect(res.status).not.toHaveBeenCalled(); + + req.body.input = original; + await secondPass(req, res as unknown as Response, next); + expect(res.status).toHaveBeenCalledWith(400); + expect(next).toHaveBeenCalledTimes(1); + req.body.input = undefined; + req.body.text = original; + expect(getPreinspectedPrivateText(req)).toBeUndefined(); + await secondPass(req, res as unknown as Response, next); + expect(res.status).toHaveBeenCalledTimes(2); + }); + it('uses stable retry revisions and distinct namespaces for different turns or originals', () => { const first = submit().message; expect(submit().message).toEqual(first); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 2059245393b..2be510516a8 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -245,18 +245,30 @@ export function createPrivateTextIngress(options: { }; } -/** Only safe metadata joins the user-message projection sent to events and ordinary readers. */ +/** Only a user turn that will persist its sidecar advertises an owner-readable revision. */ export function stampPrivateTextMessage( req: object | undefined, message: T, + willPersist = true, ): T & { privacyRevision?: string } { const capture = req == null ? undefined : captures.get(req); - if (capture != null && message.isCreatedByUser === true && message.text === capture.text) { + if ( + willPersist && + capture != null && + message.isCreatedByUser === true && + message.text === capture.text + ) { message.privacyRevision = capture.revision; } return message; } +/** Only the request that ran ingress may skip its exact, already-inspected submitted text. */ +export function getPreinspectedPrivateText(req: Request): string | undefined { + const capture = captures.get(req); + return capture != null && req.body?.text === capture.text ? capture.text : undefined; +} + /** The preliminary job record precedes the created event and may be read by Stop. */ export function stampPreliminaryPrivateTextMessage( req: object | undefined, diff --git a/packages/api/src/protection/private/view.ts b/packages/api/src/protection/private/view.ts index ebfdef77f0f..76ebd57b0e7 100644 --- a/packages/api/src/protection/private/view.ts +++ b/packages/api/src/protection/private/view.ts @@ -4,6 +4,14 @@ import type { RequestHandler, Request } from 'express'; import { createPrivateTextCipher } from './crypto'; import { privateTextBinding } from './submission'; +/** Defense in depth for search hits and mutation responses that bypass normal Mongo projections. */ +export function stripPrivateMessageFields< + T extends { privateText?: unknown; contextMeta?: unknown }, +>(message: T): Omit { + const { privateText: _privateText, contextMeta: _contextMeta, ...publicMessage } = message; + return publicMessage; +} + const inputSchema = z .object({ messageIds: z.array(z.string().min(1).max(256)).min(1).max(50), diff --git a/packages/data-schemas/src/models/plugins/mongoMeili.spec.ts b/packages/data-schemas/src/models/plugins/mongoMeili.spec.ts index 4c9a7511ca6..9be20e6e0dc 100644 --- a/packages/data-schemas/src/models/plugins/mongoMeili.spec.ts +++ b/packages/data-schemas/src/models/plugins/mongoMeili.spec.ts @@ -83,6 +83,7 @@ const mockAddDocumentsInBatches = jest.fn(); const mockUpdateDocuments = jest.fn(); const mockDeleteDocument = jest.fn(); const mockDeleteDocuments = jest.fn(); +const mockSearch = jest.fn(); const mockGetDocument = jest.fn(); const mockGetDocuments = jest.fn().mockResolvedValue({ results: [] }); const mockWaitForTask = jest.fn().mockResolvedValue({ status: 'succeeded' }); @@ -96,6 +97,7 @@ const mockIndex = jest.fn().mockReturnValue({ deleteDocuments: mockDeleteDocuments, getDocument: mockGetDocument, getDocuments: mockGetDocuments, + search: mockSearch, }); jest.mock('meilisearch', () => { return { @@ -134,6 +136,7 @@ describe('Meilisearch Mongoose plugin', () => { mockDeleteDocuments.mockReset().mockResolvedValue({ taskUid: 1 }); mockGetDocument.mockClear(); mockGetDocuments.mockReset().mockResolvedValue({ results: [] }); + mockSearch.mockReset().mockResolvedValue({ hits: [] }); mockWaitForTask.mockReset().mockResolvedValue({ status: 'succeeded' }); }); @@ -164,6 +167,28 @@ describe('Meilisearch Mongoose plugin', () => { } }); + test('hydrated message search excludes encrypted originals and all schema-hidden fields', async () => { + const Message = createMessageModel(mongoose) as unknown as SchemaWithMeiliMethods; + const messageId = new mongoose.Types.ObjectId().toString(); + const conversationId = new mongoose.Types.ObjectId().toString(); + await Message.collection.insertOne({ + messageId, + conversationId, + user: 'search-owner', + text: '[EMAIL_1_private]', + isCreatedByUser: true, + privacyRevision: 'private', + privateText: 'v1:encrypted-original', + }); + mockSearch.mockResolvedValueOnce({ hits: [{ messageId, text: '[EMAIL_1_private]' }] }); + + const result = await Message.meiliSearch('EMAIL', { filter: 'user = "search-owner"' }, true); + expect(result.hits).toHaveLength(1); + expect(result.hits[0]).toMatchObject({ messageId, privacyRevision: 'private' }); + expect(result.hits[0]).not.toHaveProperty('privateText'); + expect(result.hits[0]).not.toHaveProperty('_meiliIndex'); + }); + test('saving conversation indexes w/ meilisearch', async () => { await createConversationModel(mongoose).create({ conversationId: new mongoose.Types.ObjectId(), diff --git a/packages/data-schemas/src/models/plugins/mongoMeili.ts b/packages/data-schemas/src/models/plugins/mongoMeili.ts index 20e4498a734..f543fc47b62 100644 --- a/packages/data-schemas/src/models/plugins/mongoMeili.ts +++ b/packages/data-schemas/src/models/plugins/mongoMeili.ts @@ -767,7 +767,7 @@ const createMeiliMongooseModel = ({ const projection = Object.keys(this.schema.obj).reduce>( (results, key) => { - if (!key.startsWith('$')) { + if (!key.startsWith('$') && this.schema.path(key)?.options?.select !== false) { results[key] = 1; } return results; From 077e21df48c3c8c6df3f349ea2a0cbbed8c1532d Mon Sep 17 00:00:00 2001 From: Lia Date: Thu, 1 Oct 2026 13:31:47 +0000 Subject: [PATCH 10/18] fix: preserve trusted PII placeholders through provider inspection --- api/server/controllers/agents/client.js | 3 + .../routes/agents/__tests__/abort.spec.js | 3 +- api/server/routes/agents/index.js | 4 +- client/src/components/Chat/ChatView.tsx | 2 +- .../components/Chat/Messages/PrivateText.tsx | 4 +- e2e/specs/mock/private-text.spec.ts | 1 + .../api/src/middleware/modelBoundContent.ts | 35 +++++-- .../src/protection/private/submission.spec.ts | 98 +++++++++++++++++-- .../api/src/protection/private/submission.ts | 51 ++++++++-- .../src/methods/message.privateText.spec.ts | 1 + packages/data-schemas/src/methods/message.ts | 67 +++++++++---- 11 files changed, 221 insertions(+), 48 deletions(-) diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 960e7269ab0..5a0ea32aab9 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -158,6 +158,7 @@ const { reportLocatorTraversalFailure, filterFilesByEndpointRuntimeConfig, createModelBoundChatModelCallback: createModelBoundContentCallback, + getPrivateTextInspectionTokens, createInitialModelBoundAdmissionCallback, hasModelBoundContentProtection, assertResumeRuntimeContentAllowed, @@ -2012,6 +2013,7 @@ class AgentClient extends BaseClient { onTraversalFailure: reportLocatorTraversalFailure, legacyPii, storedMessages: this.modelBoundStoredMessages, + privateTextTokens: getPrivateTextInspectionTokens(this.modelBoundStoredMessages ?? []), }); } @@ -2030,6 +2032,7 @@ class AgentClient extends BaseClient { filters: this.options.req?.config?.filters, legacyPii: this.options.req?.config?.messageFilter?.pii, storedMessages: this.modelBoundStoredMessages, + privateTextTokens: getPrivateTextInspectionTokens(this.modelBoundStoredMessages ?? []), fileIdsBySourceMessageId: fileProjection.fileIdsBySourceMessageId, resolvedFiles: fileProjection.resolvedFiles, sourceFileProjectionOverflowed: fileProjection.overflowed, diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index b9976cb5cfb..b2b5bad2325 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -51,7 +51,8 @@ jest.mock('@librechat/api', () => ({ jest.mock('~/models', () => ({ initializeMessageBudget: jest.fn(), saveMessage: (...args) => mockSaveMessage(...args), - hasPersistedPrivateText: (...args) => mockHasPersistedPrivateText(...args), + getPersistedPrivateTextId: async (...args) => + (await mockHasPersistedPrivateText(...args)) ? 'protected-row-id' : null, getPrivateMessageTexts: (...args) => mockGetPrivateMessageTexts(...args), saveConvo: (...args) => mockSaveConvo(...args), })); diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index 3d455bcf9dd..8c6308fa787 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -58,7 +58,7 @@ const { getFiles, saveMessage, saveConvo, - hasPersistedPrivateText, + getPersistedPrivateTextId, getPrivateMessageTexts, } = require('~/models'); const { @@ -856,7 +856,7 @@ router.post('/chat/abort', chatConfigMiddleware, async (req, res, next) => { let persistedRequestId; try { const persistedRequest = await saveAbortedUserMessage( - { saveMessage, hasPersistedPrivateText, getPrivateMessageTexts }, + { saveMessage, getPersistedPrivateTextId, getPrivateMessageTexts }, messageContext, requestMessage, { context: 'api/server/routes/agents/index.js - abort user prerequisite' }, diff --git a/client/src/components/Chat/ChatView.tsx b/client/src/components/Chat/ChatView.tsx index 314dbde4a73..01f54f17035 100644 --- a/client/src/components/Chat/ChatView.tsx +++ b/client/src/components/Chat/ChatView.tsx @@ -28,10 +28,10 @@ import { QueuedTurnPortalProvider } from './Steering/QueuedTurnPortal'; import ApprovalProvider from './Messages/Content/ApprovalContext'; import ConversationStarters from './Input/ConversationStarters'; import { pendingApprovalActionFamily } from './approval/state'; -import { OwnerTextProvider } from './Messages/PrivateText'; import ProjectBadge from '~/components/Projects/ProjectBadge'; import { composerLiftFamily } from './Input/Composer/state'; import { showComposerTipsAtom } from '~/store/composerTips'; +import { OwnerTextProvider } from './Messages/PrivateText'; import { useGetMessagesByConvoId } from '~/data-provider'; import Footer, { useConfiguredFooter } from './Footer'; import { AskAnswerHostProvider } from './ask/state'; diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index 112f9f87d3e..2a9b3f06475 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -215,7 +215,7 @@ export function PrivateText({ message }: { message: TMessage }) { -

+

{localize('com_ui_private_text_hidden')} {text == null && ( @@ -228,7 +228,7 @@ export function PrivateText({ message }: { message: TMessage }) { )}

{text == null && !state.loading && state.retry != null && ( - )} diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index 94dc7ef37cd..4f8189b2a90 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -33,6 +33,7 @@ test('owner sees original after reload while provider, sharing, and canonical re starterPatterns: [], customPatterns: [ { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, ], }, }, diff --git a/packages/api/src/middleware/modelBoundContent.ts b/packages/api/src/middleware/modelBoundContent.ts index 212acc3ea18..74388e3256f 100644 --- a/packages/api/src/middleware/modelBoundContent.ts +++ b/packages/api/src/middleware/modelBoundContent.ts @@ -264,6 +264,8 @@ export interface ModelBoundProviderContentInput { readonly onTraversalFailure?: LocatorTraversalReporter; readonly filters?: FiltersConfig; readonly legacyPii?: MessageFilterPiiConfig; + /** Exact generated tokens from server-owned canonical user rows, never request metadata. */ + readonly privateTextTokens?: ReadonlySet; readonly providerMessages: readonly ModelBoundProviderMessage[]; readonly storedMessages?: readonly (StoredModelBoundMessage | null | undefined)[]; readonly resolvedFiles?: readonly (ModelBoundCanonicalFile | null | undefined)[]; @@ -725,6 +727,8 @@ export interface ModelBoundContentInput { readonly onTraversalFailure?: LocatorTraversalReporter; readonly filters?: FiltersConfig; readonly legacyPii?: MessageFilterPiiConfig; + /** Exact generated tokens from server-owned canonical user rows, never request metadata. */ + readonly privateTextTokens?: ReadonlySet; /** Fresh API input: every role is caller-submitted. */ readonly submittedMessages?: readonly ModelBoundMessage[]; /** Persisted chat history: user rows plus structured tool fragments are re-inspected. */ @@ -3146,6 +3150,7 @@ function assertIndexedModelBoundProviderContent( onTraversalFailure: input.onTraversalFailure, filters: input.filters, legacyPii: input.legacyPii, + privateTextTokens: input.privateTextTokens, storedMessages: projection.storedMessages, resolvedFiles: projection.resolvedFiles, deferredTraversalErrors: projection.deferredTraversalErrors, @@ -3242,6 +3247,7 @@ export function createModelBoundChatModelCallback( onTraversalFailure: input.onTraversalFailure, filters: input.filters, legacyPii: input.legacyPii, + privateTextTokens: input.privateTextTokens, storedMessages: storedMessageSnapshot.values, resolvedFiles: resolvedFileSnapshot.values, fileIdsBySourceMessageId: sourceFileIdSnapshot.values, @@ -3492,16 +3498,29 @@ function inspectModelBoundContent( const inspectionSession = inspector?.createSession(); const shouldContinueAfterFinding = inspectionSession?.hasAuditRules === true; let finding: ReturnType['inspect']> = null; - const inspectFragments = (fragments: Iterable): void => { - if (finding == null || shouldContinueAfterFinding) { - const nextFinding = inspectionSession?.inspect(fragments) ?? null; - finding ??= nextFinding; + const inspectFragment = (fragment: TextContentFragment): void => { + if (finding != null && !shouldContinueAfterFinding) { + return; } + const inspectableText = + input.privateTextTokens?.size && + (fragment.source === 'message' || fragment.source === 'assembled_context') + ? fragment.text.replace( + /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]/g, + (token) => (input.privateTextTokens!.has(token) ? '' : token), + ) + : fragment.text; + const inspected = + inspectableText === fragment.text ? fragment : { ...fragment, text: inspectableText }; + const nextFinding = inspectionSession?.inspectFragment(inspected) ?? null; + finding ??= nextFinding; }; - const inspectFragment = (fragment: TextContentFragment): void => { - if (finding == null || shouldContinueAfterFinding) { - const nextFinding = inspectionSession?.inspectFragment(fragment) ?? null; - finding ??= nextFinding; + const inspectFragments = (fragments: Iterable): void => { + for (const fragment of fragments) { + inspectFragment(fragment); + if (finding != null && !shouldContinueAfterFinding) { + return; + } } }; const traversalErrors: ContentTraversalLimitError[] = [ diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 991a570d804..06a940ebc6d 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -11,8 +11,13 @@ import { isPreDenialTextSubmission, isPrivateTextChatSubmission, getPreinspectedPrivateText, + getPrivateTextInspectionTokens, privateTextBinding, } from './submission'; +import { + createModelBoundChatModelCallback, + assertModelBoundContent, +} from '../../middleware/modelBoundContent'; import { createMessageFilterPii } from '../../middleware/messageFilterPii'; import { createPrivateTextCipher } from './crypto'; import { createPrivateTextView } from './view'; @@ -102,9 +107,17 @@ describe('private text submission boundary', () => { const saveMessage: MessageMethods['saveMessage'] = jest.fn( async (_ctx, message) => message as IMessage, ); - const hasPersistedPrivateText = jest.fn(async () => true); + const hasPersistedPrivateText = jest.fn( + async (_input: Parameters[0]) => true, + ); const getPrivateMessageTexts = jest.fn(async (): Promise => []); - const store = { saveMessage, hasPersistedPrivateText, getPrivateMessageTexts }; + const store = { + saveMessage, + getPersistedPrivateTextId: async ( + input: Parameters[0], + ) => ((await hasPersistedPrivateText(input)) ? 'protected-row-id' : null), + getPrivateMessageTexts, + }; const { message } = submit(); expect( await saveAbortedUserMessage( @@ -114,7 +127,7 @@ describe('private text submission boundary', () => { { context: 'Stop' }, 'tenant-a', ), - ).toBe(true); + ).toEqual({ _id: 'protected-row-id' }); expect(hasPersistedPrivateText).toHaveBeenCalledWith({ userId: 'owner', tenantId: 'tenant-a', @@ -141,7 +154,7 @@ describe('private text submission boundary', () => { { context: 'ordinary Stop' }, 'tenant-a', ), - ).toBe(true); + ).toEqual({ _id: undefined }); expect(saveMessage).toHaveBeenCalledTimes(1); }); @@ -158,7 +171,7 @@ describe('private text submission boundary', () => { ]); const store = { saveMessage, - hasPersistedPrivateText: jest.fn(async () => true), + getPersistedPrivateTextId: jest.fn(async () => 'protected-row-id'), getPrivateMessageTexts, }; const finalEvent = { requestMessage: { messageId: message.messageId, privacyRevision: '' } }; @@ -172,7 +185,7 @@ describe('private text submission boundary', () => { 'tenant-a', finalEvent, ), - ).toBe(true); + ).toEqual({ _id: undefined }); expect(saveMessage).not.toHaveBeenCalled(); expect(finalEvent.requestMessage.privacyRevision).toBe(message.privacyRevision); getPrivateMessageTexts.mockResolvedValueOnce([]); @@ -257,6 +270,79 @@ describe('private text submission boundary', () => { expect(res.status).toHaveBeenCalledTimes(2); }); + it('keeps trusted placeholders safe through every provider call, including restored history', () => { + const { message } = submit(); + const patterns: FiltersConfig = { + messages: { + pii: { + action: 'redact', + fields: ['text', 'content_part', 'assembled_context'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, + ], + }, + }, + }; + const privateTextTokens = getPrivateTextInspectionTokens([message]); + expect(privateTextTokens.size).toBe(1); + const providerMessage = { role: 'user', content: message.text }; + const storedMessages = [{ ...message, role: 'user' }]; + const callback = createModelBoundChatModelCallback({ + filters: patterns, + storedMessages, + privateTextTokens, + }); + expect(() => callback.handleChatModelStart(undefined, [[providerMessage]])).not.toThrow(); + expect(() => + callback.handleChatModelStart(undefined, [ + [providerMessage, { role: 'user', content: 'Follow up' }], + ]), + ).not.toThrow(); + const restored = createModelBoundChatModelCallback({ + filters: patterns, + storedMessages: JSON.parse(JSON.stringify(storedMessages)) as typeof storedMessages, + privateTextTokens: getPrivateTextInspectionTokens( + JSON.parse(JSON.stringify(storedMessages)) as typeof storedMessages, + ), + }); + expect(() => restored.handleChatModelStart(undefined, [[providerMessage]])).not.toThrow(); + for (const unsafe of [ + original, + `${message.text} ${'f'.repeat(32)}`, + `Email [EMAIL_1_${'f'.repeat(32)}]`, + ]) { + expect(() => + callback.handleChatModelStart(undefined, [[{ role: 'user', content: unsafe }]]), + ).toThrow(); + } + expect(getPrivateTextInspectionTokens([{ ...message, privacyRevision: undefined }]).size).toBe( + 0, + ); + expect( + getPrivateTextInspectionTokens([{ ...message, privacyRevision: 'f'.repeat(32) }]).size, + ).toBe(0); + const untrusted = createModelBoundChatModelCallback({ + filters: patterns, + storedMessages: [], + privateTextTokens: new Set(), + }); + expect(() => untrusted.handleChatModelStart(undefined, [[providerMessage]])).toThrow(); + expect(() => + assertModelBoundContent({ + filters: { + ...patterns, + agentInstructions: { + pii: { customPatterns: [{ id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}' }] }, + }, + }, + privateTextTokens, + agents: [{ instructions: message.text }], + }), + ).toThrow(); + }); + it('uses stable retry revisions and distinct namespaces for different turns or originals', () => { const first = submit().message; expect(submit().message).toEqual(first); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 2be510516a8..7ebf79b2514 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -269,6 +269,40 @@ export function getPreinspectedPrivateText(req: Request): string | undefined { return capture != null && req.body?.text === capture.text ? capture.text : undefined; } +/** Only call with server-owned canonical rows; generic writes strip their private revision. */ +export function getPrivateTextInspectionTokens( + messages: readonly (PrivateTextMessage | null | undefined)[], +): ReadonlySet { + const tokens = new Set(); + // Mirror the existing provider work and maximum transformed-text ceilings. + if (messages.length > 4096) { + throw unavailable(); + } + for (const message of messages) { + if ( + message?.isCreatedByUser !== true || + typeof message.text !== 'string' || + !/^[a-f0-9]{32}$/.test(message.privacyRevision ?? '') + ) { + continue; + } + if (message.text.length > 524288) { + throw unavailable(); + } + for (const match of message.text.matchAll( + /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_([a-f0-9]{32})\]/g, + )) { + if (match[1] === message.privacyRevision) { + tokens.add(match[0]); + if (tokens.size > 4096) { + throw unavailable(); + } + } + } + } + return tokens; +} + /** The preliminary job record precedes the created event and may be read by Stop. */ export function stampPreliminaryPrivateTextMessage( req: object | undefined, @@ -359,13 +393,16 @@ export async function requirePrivateTextPersistence( * omit the revision, so their prerequisite is an insert-only write. */ export async function saveAbortedUserMessage( - store: Pick, + store: Pick< + MessageMethods, + 'saveMessage' | 'getPersistedPrivateTextId' | 'getPrivateMessageTexts' + >, ctx: Parameters[0], message: Parameters[1], metadata: Parameters[2], tenantId?: string, finalEvent?: { requestMessage?: { privacyRevision?: string } | null }, -): Promise { +): Promise<{ _id?: unknown }> { const revision = message.privacyRevision; if (typeof revision !== 'string' || revision.length === 0) { if (!message.messageId || !message.conversationId) { @@ -385,7 +422,7 @@ export async function saveAbortedUserMessage( if (typeof saved.text !== 'string') { throw unavailable(); } - const exists = await store.hasPersistedPrivateText({ + const exists = await store.getPersistedPrivateTextId({ userId: ctx.userId, tenantId, conversationId: message.conversationId, @@ -400,7 +437,7 @@ export async function saveAbortedUserMessage( finalEvent.requestMessage.privacyRevision = saved.privacyRevision; } } - return true; + return { _id: saved._id }; } const rows = await store.getPrivateMessageTexts({ userId: ctx.userId, @@ -421,12 +458,12 @@ export async function saveAbortedUserMessage( if (finalEvent?.requestMessage != null) { finalEvent.requestMessage.privacyRevision = row.privacyRevision; } - return true; + return { _id: row._id }; } if (!message.messageId || !message.conversationId || typeof message.text !== 'string') { throw unavailable(); } - const persisted = await store.hasPersistedPrivateText({ + const persisted = await store.getPersistedPrivateTextId({ userId: ctx.userId, tenantId, conversationId: message.conversationId, @@ -437,5 +474,5 @@ export async function saveAbortedUserMessage( if (!persisted) { throw unavailable(); } - return true; + return { _id: persisted }; } diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index ddc7af4e550..cfc0f2018ad 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -70,6 +70,7 @@ it('stores both views atomically and excludes ciphertext from ordinary and clien privacyRevision: 'turn', }; expect(await methods.hasPersistedPrivateText(protectedRow)).toBe(true); + expect(await methods.getPersistedPrivateTextId(protectedRow)).toBe(String(saved?._id)); for (const mismatch of [ { userId: 'other' }, { tenantId: 'tenant-b' }, diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 00f25eeadfa..eaad2eaacd1 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -702,6 +702,7 @@ export interface PrivateTextWrite { } export interface PrivateTextRead { + readonly _id?: unknown; readonly messageId: string; readonly text: string; readonly privacyRevision: string; @@ -709,6 +710,9 @@ export interface PrivateTextRead { } export interface MessageMethods { + getPersistedPrivateTextId( + input: Parameters[0], + ): Promise; hasPersistedPrivateText(input: { userId: string; tenantId?: string; @@ -1260,8 +1264,12 @@ export function createMessageMethods( { messageId: params.messageId, user: userId }, { ...update, userSubmittedPaths, userSubmittedMessageFieldPaths }, { - upsert: true, stampModelOutputOnInsert, unsetContextMeta, retentionOnInsert, - unsetPrivateText: metadata?.privateText == null && Object.prototype.hasOwnProperty.call(update, 'text'), + upsert: true, + stampModelOutputOnInsert, + unsetContextMeta, + retentionOnInsert, + unsetPrivateText: + metadata?.privateText == null && Object.prototype.hasOwnProperty.call(update, 'text'), }, ); @@ -1374,7 +1382,11 @@ export function createMessageMethods( messageId: message.messageId, ...(message.user != null ? { user: message.user } : {}), }, - update: { $set: normalizedMessage, $inc: { __v: 1 }, $unset: { privateText: 1, privacyRevision: 1 } }, + update: { + $set: normalizedMessage, + $inc: { __v: 1 }, + $unset: { privateText: 1, privacyRevision: 1 }, + }, timestamps: !overrideTimestamp, upsert: true, }, @@ -1531,7 +1543,11 @@ export function createMessageMethods( { messageId, text }: { messageId: string; text: string }, ) { try { - await writeMessage({ messageId, user: userId }, { text }, { upsert: false, unsetPrivateText: true }); + await writeMessage( + { messageId, user: userId }, + { text }, + { upsert: false, unsetPrivateText: true }, + ); } catch (err) { logger.error('Error updating message text:', err); throw err; @@ -2443,8 +2459,11 @@ export function createMessageMethods( ) { try { const { messageId, ...update } = message; + delete update.privateText; + delete update.privacyRevision; const updatedMessage = await writeMessage({ messageId, user: userId }, update, { upsert: false, + unsetPrivateText: Object.prototype.hasOwnProperty.call(update, 'text'), }); if (!updatedMessage) { @@ -4118,39 +4137,44 @@ export function createMessageMethods( return Message.meiliSearch(query, searchOptions, hydrate); } - async function hasPersistedPrivateText(input: { + async function getPersistedPrivateTextId(input: { userId: string; tenantId?: string; conversationId: string; messageId: string; privacyRevision: string; text: string; - }): Promise { + }): Promise { if ( !input.userId || !input.messageId || !input.privacyRevision || !UUID_REGEX.test(input.conversationId) ) { - return false; + return null; } const activeTenant = tenantStorage.getStore()?.tenantId; if (activeTenant != null && activeTenant !== input.tenantId) { - return false; + return null; } const Message = mongoose.models.Message as Model; - return ( - (await Message.exists({ - user: input.userId, - ...traceTenantScope(input.tenantId), - conversationId: input.conversationId, - messageId: input.messageId, - text: input.text, - privacyRevision: input.privacyRevision, - privateText: { $exists: true }, - $or: [{ expiredAt: null }, { expiredAt: { $gt: new Date() } }], - })) != null - ); + const stored = await Message.exists({ + user: input.userId, + ...traceTenantScope(input.tenantId), + conversationId: input.conversationId, + messageId: input.messageId, + text: input.text, + privacyRevision: input.privacyRevision, + privateText: { $exists: true }, + $or: [{ expiredAt: null }, { expiredAt: { $gt: new Date() } }], + }); + return stored == null ? null : String(stored._id); + } + + async function hasPersistedPrivateText( + input: Parameters[0], + ): Promise { + return (await getPersistedPrivateTextId(input)) != null; } async function getPrivateMessageTexts(input: { @@ -4176,13 +4200,14 @@ export function createMessageMethods( privateText: { $exists: true }, $or: [{ expiredAt: null }, { expiredAt: { $gt: new Date() } }], }) - .select('messageId text privacyRevision +privateText -_id') + .select('messageId text privacyRevision +privateText') .limit(50) .lean(); } return { hasPersistedPrivateText, + getPersistedPrivateTextId, getPrivateMessageTexts, saveMessage, bulkSaveMessages, From fad416b74e77882023cb866e8425748591cabccf Mon Sep 17 00:00:00 2001 From: Lia Date: Thu, 1 Oct 2026 13:51:29 +0000 Subject: [PATCH 11/18] fix: coordinate protected turns across Stop and native sharing --- api/app/clients/BaseClient.js | 6 ++ api/app/clients/specs/BaseClient.test.js | 62 +++++++++++++++++++ api/server/controllers/agents/client.js | 1 + api/server/controllers/agents/client.test.js | 41 ++++++++++++ .../__tests__/messages-retention.spec.js | 4 ++ .../routes/__tests__/messages-update.spec.js | 4 ++ packages/api/src/imports.ts | 3 + .../api/src/protection/private/submission.ts | 6 +- .../api/src/shared-links/protection.spec.ts | 38 ++++++++++++ packages/api/src/shared-links/protection.ts | 27 +++++--- packages/data-schemas/src/methods/share.ts | 13 +++- 11 files changed, 193 insertions(+), 12 deletions(-) diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index cc38bdc1757..43068a3520b 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -636,6 +636,12 @@ class BaseClient { }); } + // A created frame enables Stop and owner reads on other replicas. Protected + // turns must have their atomic user/conversation write committed first. + await requirePrivateTextPersistence(this.options.req, () => + this.saveMessageToDatabase(userMessage, saveOptions, user), + ); + if (typeof opts?.onStart === 'function') { const isNewConvo = !requestConvoId && parentMessageId === Constants.NO_PARENT; opts.onStart(userMessage, responseMessageId, isNewConvo); diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index d6e1549adec..10dd5de6e73 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -2,6 +2,7 @@ const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-prov const BaseClientClass = require('../BaseClient'); const { ContentFilterError, + createPrivateTextIngress, resolveTurnDeliveryRouting, buildSteerMedia, Tokenizer, @@ -1734,6 +1735,67 @@ describe('BaseClient', () => { ); }); + test('protected created events wait for the atomic user write and fail closed on write failure', async () => { + const req = { + user: { id: 'owner' }, + path: '/', + body: { text: 'alice@example.com', clientRequestId: 'created-privacy' }, + }; + const res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; + const next = jest.fn(); + createPrivateTextIngress({ + getFilters: () => ({ + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }), + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, res, next); + expect(next).toHaveBeenCalledTimes(1); + const client = Object.create(BaseClientClass.prototype); + client.options = { req }; + client.sender = 'Agent'; + client.resolveStartUserMessage = jest.fn(() => ({ + messageId: 'private-user', + conversationId: 'private-conversation', + text: req.body.text, + isCreatedByUser: true, + })); + client.setMessageOptions = jest.fn(async () => ({ + user: 'owner', + saveOptions: {}, + conversationId: 'private-conversation', + responseMessageId: 'private-response', + parentMessageId: Constants.NO_PARENT, + })); + const committed = deferred(); + let written; + client.saveMessageToDatabase = jest.fn((message) => { + written = message; + return committed.promise; + }); + const onStart = jest.fn(); + const started = client.handleStartMethods(req.body.text, { onStart }); + await Promise.resolve(); + await Promise.resolve(); + expect(client.saveMessageToDatabase).toHaveBeenCalledTimes(1); + expect(onStart).not.toHaveBeenCalled(); + committed.resolve({ message: { ...written } }); + await started; + expect(onStart).toHaveBeenCalledTimes(1); + client.saveMessageToDatabase.mockResolvedValueOnce({}); + await expect(client.handleStartMethods(req.body.text, { onStart })).rejects.toThrow(); + expect(onStart).toHaveBeenCalledTimes(1); + }); + test('onStart is called with the correct arguments', async () => { const onStart = jest.fn(); const opts = { onStart }; diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 5a0ea32aab9..468d0262a10 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -2750,6 +2750,7 @@ class AgentClient extends BaseClient { filters: this.options.req.config?.filters, legacyPii: this.options.req.config?.messageFilter?.pii, submittedMessages: [{ role: 'user', content: latestFormatted.content }], + privateTextTokens: getPrivateTextInspectionTokens([latestOrdered]), }); /** Google rejects an unusable video with a generic `INVALID_ARGUMENT` that names no cause, * so `#sendCompletion` can only attribute one by knowing this turn carried a video. */ diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index a6a3991c421..c03c096cd69 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -7335,6 +7335,47 @@ describe('AgentClient - titleConvo', () => { require('~/models').getUserMemories.mockReset().mockResolvedValue([]); }); + it.each([Providers.GOOGLE, Providers.VERTEXAI])( + 'allows protected text through the late %s urlContext preflight without exempting raw content', + async (provider) => { + const filters = { + messages: { + pii: { + action: 'redact', + fields: ['text', 'content_part'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, + ], + }, + }, + }; + const { client, invokeModel } = createClient({ provider, filters }); + const req = client.options.req; + req.body.text = 'Email alice@example.com'; + req.body.clientRequestId = 'google-private-text'; + require('@librechat/api').createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, { status: jest.fn().mockReturnThis(), json: jest.fn() }, jest.fn()); + client.skipSaveUserMessage = false; + client.saveMessageToDatabase.mockImplementation(async (message) => ({ + message: { ...message }, + })); + await expect( + client.sendMessage(req.body.text, { + conversationId: 'protected-google', + parentMessageId: Constants.NO_PARENT, + user: 'user-123', + }), + ).resolves.toBeDefined(); + expect(invokeModel).toHaveBeenCalledTimes(1); + expect(req.body.text).not.toContain('alice@example.com'); + }, + ); + it.each([Providers.GOOGLE, Providers.VERTEXAI])( 'blocks a late %s fileUri before model invocation under strict content policy', async (provider) => { diff --git a/api/server/routes/__tests__/messages-retention.spec.js b/api/server/routes/__tests__/messages-retention.spec.js index fc796580549..9bbbc0ce9da 100644 --- a/api/server/routes/__tests__/messages-retention.spec.js +++ b/api/server/routes/__tests__/messages-retention.spec.js @@ -7,6 +7,10 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), + stripPrivateMessageFields: jest.requireActual( + '../../../../packages/api/src/protection/private/view', + ).stripPrivateMessageFields, unescapeLaTeX: jest.fn((value) => value), countTokens: jest.fn().mockResolvedValue(10), createContentFilter: jest.fn(() => (_req, _res, next) => next()), diff --git a/api/server/routes/__tests__/messages-update.spec.js b/api/server/routes/__tests__/messages-update.spec.js index a58e3cd4750..5394b49cb41 100644 --- a/api/server/routes/__tests__/messages-update.spec.js +++ b/api/server/routes/__tests__/messages-update.spec.js @@ -7,6 +7,10 @@ jest.mock('@librechat/agents', () => ({ })); jest.mock('@librechat/api', () => ({ + createPrivateTextView: jest.fn(() => (_req, _res, next) => next()), + stripPrivateMessageFields: jest.requireActual( + '../../../../packages/api/src/protection/private/view', + ).stripPrivateMessageFields, unescapeLaTeX: jest.fn((value) => value), countTokens: jest.fn().mockResolvedValue(10), createContentFilter: jest.fn(() => (_req, _res, next) => next()), diff --git a/packages/api/src/imports.ts b/packages/api/src/imports.ts index d216ce5a5ba..4f26093832d 100644 --- a/packages/api/src/imports.ts +++ b/packages/api/src/imports.ts @@ -56,6 +56,8 @@ export interface ConversationImportProtectionContext { readonly getFiles?: GetCanonicalFilesForInspection; readonly trustedLiveFiles?: readonly CanonicalFileInspectionFile[]; readonly legacyPii?: MessageFilterPiiConfig | null; + /** Supplied only by native canonical-copy/share callers, never parsed from an import. */ + readonly privateTextTokens?: ReadonlySet; /** Injectable only at the legacy package boundary and in focused tests. */ readonly assertModelBoundContent?: (input: ModelBoundContentInput) => void; } @@ -161,6 +163,7 @@ async function inspectConversationImportContent( onTraversalFailure: context.onTraversalFailure, filters: activeFilters, legacyPii, + privateTextTokens: context.privateTextTokens, storedMessages: [message], }); } catch (error) { diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 7ebf79b2514..0b62163a3a5 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -276,7 +276,7 @@ export function getPrivateTextInspectionTokens( const tokens = new Set(); // Mirror the existing provider work and maximum transformed-text ceilings. if (messages.length > 4096) { - throw unavailable(); + return tokens; } for (const message of messages) { if ( @@ -287,7 +287,7 @@ export function getPrivateTextInspectionTokens( continue; } if (message.text.length > 524288) { - throw unavailable(); + return new Set(); } for (const match of message.text.matchAll( /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_([a-f0-9]{32})\]/g, @@ -295,7 +295,7 @@ export function getPrivateTextInspectionTokens( if (match[1] === message.privacyRevision) { tokens.add(match[0]); if (tokens.size > 4096) { - throw unavailable(); + return new Set(); } } } diff --git a/packages/api/src/shared-links/protection.spec.ts b/packages/api/src/shared-links/protection.spec.ts index 86cf9eee6e3..f6690f2b351 100644 --- a/packages/api/src/shared-links/protection.spec.ts +++ b/packages/api/src/shared-links/protection.spec.ts @@ -32,6 +32,44 @@ function capturePolicyError( throw new Error('Expected content policy to reject the shared metadata'); } +describe('native protected message sharing', () => { + it('inspects the remainder of canonical messages without matching generated revision tokens', async () => { + const revision = 'a'.repeat(32); + const message = { + isCreatedByUser: true, + text: `Email [EMAIL_1_${revision}]`, + privacyRevision: revision, + }; + const preflight = createShareContentPreflight({ + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, + ], + }, + }, + }); + await expect(preflight!({ title: 'Safe', messages: [message] })).resolves.toBeUndefined(); + const publicMessage = { isCreatedByUser: true, text: message.text }; + await expect( + preflight!({ title: 'Safe', messages: [publicMessage] }, { canonicalMessages: [message] }), + ).resolves.toBeUndefined(); + expect(publicMessage).not.toHaveProperty('privacyRevision'); + await expect( + preflight!({ + title: 'Safe', + messages: [{ ...message, text: `${message.text} ${'f'.repeat(32)}` }], + }), + ).rejects.toThrow(); + await expect( + preflight!({ title: 'Safe', messages: [{ ...message, privacyRevision: undefined }] }), + ).rejects.toThrow(); + }); +}); + describe('shared file metadata protection', () => { const attachmentFilters: FiltersConfig = { messages: { diff --git a/packages/api/src/shared-links/protection.ts b/packages/api/src/shared-links/protection.ts index 5a788a28b1b..a9c2334f44d 100644 --- a/packages/api/src/shared-links/protection.ts +++ b/packages/api/src/shared-links/protection.ts @@ -33,6 +33,7 @@ import { extractFileContent, extractStoredMessageContent, } from '../protection/adapters/submissions'; +import { getPrivateTextInspectionTokens } from '../protection/private/submission'; import { assertModelBoundContent } from '../middleware/modelBoundContent'; import { getUserSubmittedPathState } from '../protection/provenance'; import { assertConversationImportContentAllowed } from '../imports'; @@ -69,6 +70,7 @@ interface SerializedSharedMessagePart { } export interface SerializedSharedMessage { + readonly privacyRevision?: string; readonly isCreatedByUser?: boolean; readonly isUserSubmitted?: boolean; readonly userSubmittedPaths?: readonly string[]; @@ -111,7 +113,16 @@ export interface ShareContentPreflightOptions { readonly sharedFileMetadataFiles?: boolean; } -export type ShareContentPreflight = (input: ShareContentPreflightInput) => Promise; +export type ShareContentPreflight = ( + input: ShareContentPreflightInput, + context?: { + readonly canonicalMessages: readonly { + readonly text?: string; + readonly isCreatedByUser?: boolean; + readonly privacyRevision?: string; + }[]; + }, +) => Promise; function omitUnsharedMessageFiles( messages: readonly ShareContentPreflightMessage[], @@ -142,11 +153,10 @@ export function createShareContentPreflight( return undefined; } - const inspectSharedContent = async ({ - title, - messages, - shareId, - }: ShareContentPreflightInput): Promise => { + const inspectSharedContent = async ( + { title, messages, shareId }: ShareContentPreflightInput, + context?: Parameters[1], + ): Promise => { const inspectSharedFileMetadata = options.sharedFileMetadata === true; const inspectSharedFiles = inspectSharedFileMetadata && options.sharedFileMetadataFiles !== false; @@ -165,6 +175,9 @@ export function createShareContentPreflight( user: options.user, getFiles: options.getFiles, onTraversalFailure: options.onTraversalFailure, + privateTextTokens: getPrivateTextInspectionTokens( + context?.canonicalMessages ?? snapshotMessages, + ), }, ); if (!inspectSharedFileMetadata) { @@ -178,7 +191,7 @@ export function createShareContentPreflight( }); }; - return (input) => aggregateAuditFindings(() => inspectSharedContent(input)); + return (input, context) => aggregateAuditFindings(() => inspectSharedContent(input, context)); } const SERIALIZED_LOCATOR_KEYS = [ diff --git a/packages/data-schemas/src/methods/share.ts b/packages/data-schemas/src/methods/share.ts index 1849be8f139..ff7a27aee34 100644 --- a/packages/data-schemas/src/methods/share.ts +++ b/packages/data-schemas/src/methods/share.ts @@ -75,7 +75,16 @@ export type SharedLinkContentPreflight = ( snapshot: SharedLinkContentSnapshot, ) => void | Promise; -export type SharedMessagesPreflight = (snapshot: t.SharedMessagesResult) => void | Promise; +export type SharedMessagesPreflight = ( + snapshot: t.SharedMessagesResult, + context?: { + readonly canonicalMessages: readonly { + readonly text?: string; + readonly isCreatedByUser?: boolean; + readonly privacyRevision?: string; + }[]; + }, +) => void | Promise; export interface GetSharedMessagesOptions { readonly snapshotFiles?: boolean; @@ -1048,7 +1057,7 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { }; try { - await options?.preflight?.(result); + await options?.preflight?.(result, { canonicalMessages: messagesToShare }); } catch (error) { preflightFailed = true; throw error; From 10499e6f6773f1d249bf0feb8a281442c612d253 Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 02:55:23 +0000 Subject: [PATCH 12/18] fix: Preserve Private Text Admission and Query Lifecycles --- api/app/clients/BaseClient.js | 24 +-- api/app/clients/specs/BaseClient.test.js | 152 +++++++++++---- api/server/controllers/agents/client.js | 6 + .../routes/agents/__tests__/abort.spec.js | 22 +-- .../Chat/Messages/PrivateText.spec.tsx | 104 ++++++++-- .../components/Chat/Messages/PrivateText.tsx | 179 +++--------------- client/src/data-provider/Messages/private.ts | 152 +++++++++++++++ .../src/middleware/modelBoundContent.spec.ts | 37 ++++ .../api/src/middleware/modelBoundContent.ts | 8 +- .../src/protection/private/submission.spec.ts | 34 +++- .../api/src/protection/private/submission.ts | 116 +++++++----- packages/data-provider/src/keys.ts | 1 + 12 files changed, 543 insertions(+), 292 deletions(-) create mode 100644 client/src/data-provider/Messages/private.ts diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 43068a3520b..694c4e8e31e 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -34,6 +34,7 @@ const { getConversationWriteContext, savePrivateTextMessage, stampPrivateTextMessage, + deferPrivateTextStart, requirePrivateTextPersistence, persistedReasoningOverrideFields, } = require('@librechat/api'); @@ -636,17 +637,14 @@ class BaseClient { }); } - // A created frame enables Stop and owner reads on other replicas. Protected - // turns must have their atomic user/conversation write committed first. - await requirePrivateTextPersistence(this.options.req, () => - this.saveMessageToDatabase(userMessage, saveOptions, user), + this.privateTextStart = deferPrivateTextStart( + this.options.req, + opts?.onStart, + userMessage, + responseMessageId, + !requestConvoId && parentMessageId === Constants.NO_PARENT, ); - if (typeof opts?.onStart === 'function') { - const isNewConvo = !requestConvoId && parentMessageId === Constants.NO_PARENT; - opts.onStart(userMessage, responseMessageId, isNewConvo); - } - return { ...opts, user, @@ -1053,9 +1051,6 @@ class BaseClient { await balanceReservations.track(balanceAdmission); } - await requirePrivateTextPersistence(this.options.req, () => - userMessagePersistence != null ? userMessagePersistence.start() : userMessagePromise, - ); completionResult = await this.sendCompletion(payload, opts); } catch (error) { if (userMessagePersistence?.isPending()) { @@ -1069,6 +1064,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) { diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 10dd5de6e73..947dafafd17 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -3,6 +3,10 @@ const BaseClientClass = require('../BaseClient'); const { ContentFilterError, createPrivateTextIngress, + createModelBoundChatModelCallback, + getPrivateTextAdmission, + getPrivateTextInspectionTokens, + assertModelBoundContent, resolveTurnDeliveryRouting, buildSteerMedia, Tokenizer, @@ -1735,65 +1739,131 @@ describe('BaseClient', () => { ); }); - test('protected created events wait for the atomic user write and fail closed on write failure', async () => { + 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 res = { status: jest.fn().mockReturnThis(), json: jest.fn() }; const next = jest.fn(); createPrivateTextIngress({ - getFilters: () => ({ - messages: { - pii: { - action: 'redact', - fields: ['text'], - starterPatterns: [], - customPatterns: [ - { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, - ], - }, - }, - }), - getLegacyPii: () => undefined, + getFilters: () => filters, + getLegacyPii: () => legacyPii, getKey: () => 'ab'.repeat(32), - })(req, res, next); + })(req, { status: jest.fn().mockReturnThis(), json: jest.fn() }, next); expect(next).toHaveBeenCalledTimes(1); - const client = Object.create(BaseClientClass.prototype); - client.options = { req }; - client.sender = 'Agent'; - client.resolveStartUserMessage = jest.fn(() => ({ - messageId: 'private-user', - conversationId: 'private-conversation', - text: req.body.text, - isCreatedByUser: true, - })); - client.setMessageOptions = jest.fn(async () => ({ - user: 'owner', - saveOptions: {}, - conversationId: 'private-conversation', - responseMessageId: 'private-response', - parentMessageId: Constants.NO_PARENT, - })); + 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 = jest.fn((message) => { + client.saveMessageToDatabase.mockImplementationOnce((message) => { written = message; return committed.promise; }); const onStart = jest.fn(); - const started = client.handleStartMethods(req.body.text, { onStart }); - await Promise.resolve(); - await Promise.resolve(); - expect(client.saveMessageToDatabase).toHaveBeenCalledTimes(1); + 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(); - committed.resolve({ message: { ...written } }); - await started; + 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({}); - await expect(client.handleStartMethods(req.body.text, { onStart })).rejects.toThrow(); - expect(onStart).toHaveBeenCalledTimes(1); + const onStart = jest.fn(); + await expect(client.sendMessage(req.body.text, { onStart })).rejects.toThrow(); + expect(onStart).not.toHaveBeenCalled(); + expect(provider).not.toHaveBeenCalled(); + }); + + 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 () => { diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 468d0262a10..f5c39b77177 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -159,6 +159,7 @@ const { filterFilesByEndpointRuntimeConfig, createModelBoundChatModelCallback: createModelBoundContentCallback, getPrivateTextInspectionTokens, + getPrivateTextAdmission, createInitialModelBoundAdmissionCallback, hasModelBoundContentProtection, assertResumeRuntimeContentAllowed, @@ -2039,6 +2040,11 @@ class AgentClient extends BaseClient { }, { onContentRejected: persistence?.cancel, + onContentAllowed: getPrivateTextAdmission( + this.options.req, + persistence?.start, + this.privateTextStart, + ), }, ); } diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index b2b5bad2325..93d3d4e04e5 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -451,27 +451,25 @@ describe('Agent Abort Endpoint', () => { await options.beforePublish(abortResult); return abortResult; }); - mockGetPrivateMessageTexts.mockResolvedValueOnce([ - { - messageId: userMessageId, - text, - privacyRevision: 'recovered-revision', - privateText: 'v1:protected', - }, - ]); + mockSaveMessage.mockImplementationOnce(async (_ctx, message, metadata) => { + expect(metadata.insertOnly).toBe(true); + return { ...message, _id: 'protected-parent', privacyRevision: 'recovered-revision' }; + }); const response = await request(app).post('/api/agents/chat/abort').send({ conversationId }); expect(response.status).toBe(200); - expect(mockGetPrivateMessageTexts).toHaveBeenCalledWith( + expect(mockGetPrivateMessageTexts).not.toHaveBeenCalled(); + expect(mockHasPersistedPrivateText).toHaveBeenCalledWith( expect.objectContaining({ userId: 'test-user-123', conversationId, - messageIds: [userMessageId], + messageId: userMessageId, + privacyRevision: 'recovered-revision', + text, }), ); expect(abortResult.finalEvent.requestMessage.privacyRevision).toBe('recovered-revision'); - expect(mockHasPersistedPrivateText).not.toHaveBeenCalled(); - expect(mockSaveMessage).toHaveBeenCalledTimes(1); + expect(mockSaveMessage).toHaveBeenCalledTimes(2); expect(mockSaveMessage).toHaveBeenCalledWith( expect.anything(), expect.objectContaining({ messageId: 'protected-response' }), diff --git a/client/src/components/Chat/Messages/PrivateText.spec.tsx b/client/src/components/Chat/Messages/PrivateText.spec.tsx index 76754777c99..dab3c2f8f6d 100644 --- a/client/src/components/Chat/Messages/PrivateText.spec.tsx +++ b/client/src/components/Chat/Messages/PrivateText.spec.tsx @@ -1,6 +1,9 @@ import { dataService } from 'librechat-data-provider'; import { render, screen, waitFor, act } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; import type { TMessage } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import { getOwnerQueryClient } from '~/data-provider/Messages/private'; import { OwnerTextProvider, PrivateText } from './PrivateText'; let mockOwnerId = 'owner'; @@ -9,7 +12,10 @@ jest.mock('~/hooks/AuthContext', () => ({ useAuthContext: () => ({ user: { id: mockOwnerId, tenantId: mockTenantId } }), })); jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); -jest.mock('librechat-data-provider', () => ({ dataService: { getOwnerMessageTexts: jest.fn() } })); +jest.mock('librechat-data-provider', () => ({ + ...jest.requireActual('librechat-data-provider'), + dataService: { getOwnerMessageTexts: jest.fn() }, +})); jest.mock('./Content/MessageContent', () => ({ DisplayMessage: ({ text, message }: { text: string; message: TMessage }) => (
@@ -26,6 +32,13 @@ const canonical = Object.freeze({ privacyRevision: 'turn', }) as TMessage; const load = dataService.getOwnerMessageTexts as jest.Mock; +let applicationClient: QueryClient; +function Wrapper({ children }: { children: ReactNode }) { + return {children}; +} +function renderView(ui: ReactNode) { + return render(ui, { wrapper: Wrapper }); +} const original = { canonicalText: canonical.text, messageId: 'message', @@ -58,11 +71,12 @@ beforeEach(() => { mockOwnerId = 'owner'; mockTenantId = 'tenant-a'; load.mockReset(); + applicationClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); }); it('renders originals without mutating canonical model/export input, and reloads from the private API', async () => { load.mockResolvedValue({ messages: [original] }); - const first = render(); + const first = renderView(); expect(await screen.findByText('alice@example.com')).toBeInTheDocument(); expect(screen.getByTestId('canonical')).not.toHaveTextContent('alice@example.com'); expect(screen.getByTestId('standard-user-renderer')).toHaveAttribute( @@ -72,7 +86,10 @@ it('renders originals without mutating canonical model/export input, and reloads expect(canonical.text).toBe('[EMAIL_1_turn]'); expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_hidden'); first.unmount(); - render(); + await waitFor(() => + expect(getOwnerQueryClient(applicationClient).getQueryCache().getAll()).toHaveLength(0), + ); + renderView(); expect(await screen.findByText('alice@example.com')).toBeInTheDocument(); expect(load).toHaveBeenCalledTimes(2); }); @@ -83,7 +100,7 @@ it('loads the acknowledged first-turn owner text before navigating away from /ne conversationId: '11111111-1111-4111-8111-111111111111', }; load.mockResolvedValue({ messages: [{ ...original, canonicalText: firstTurn.text }] }); - const view = render(); + const view = renderView(); expect(await screen.findByText(original.text)).toBeInTheDocument(); expect(load).toHaveBeenCalledWith(firstTurn.conversationId, ['message']); expect(load).not.toHaveBeenCalledWith('new', expect.anything()); @@ -107,7 +124,7 @@ it('never mixes protected rows from stale conversations into a first-turn owner conversationId: '22222222-2222-4222-8222-222222222222', }; load.mockResolvedValue({ messages: [{ ...original, messageId: 'fresh-message' }] }); - render(); + renderView(); expect(await screen.findByText(original.text)).toBeInTheDocument(); expect(load).toHaveBeenCalledTimes(1); expect(load).toHaveBeenCalledWith(fresh.conversationId, ['fresh-message']); @@ -115,7 +132,7 @@ it('never mixes protected rows from stale conversations into a first-turn owner it('does not fetch originals for an ordinary transcript', () => { const plain = { ...canonical, privacyRevision: undefined }; - render( + renderView( , @@ -125,7 +142,7 @@ it('does not fetch originals for an ordinary transcript', () => { }); it('renders only filtered text without an owner provider, as on external viewers', () => { - render(); + renderView(); expect(screen.getByText(canonical.text)).toBeInTheDocument(); expect(load).not.toHaveBeenCalled(); }); @@ -137,19 +154,21 @@ it('shows loading then safe unavailable text when decryption or authorization fa finish = resolve; }), ); - render(); + renderView(); expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_loading'); await act(async () => { finish({ messages: [] }); }); expect(screen.getByText(canonical.text)).toBeInTheDocument(); - expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'); + await waitFor(() => + expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), + ); }); it('retries a provisional empty owner read once on turn completion, without repeatedly polling', async () => { load.mockResolvedValueOnce({ messages: [] }); load.mockResolvedValueOnce({ messages: [original] }); - const view = render(); + const view = renderView(); await waitFor(() => expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), ); @@ -165,7 +184,7 @@ it('retries a provisional empty owner read once on turn completion, without repe it('does not automatically retry an old missing owner row on unrelated submission transitions', async () => { load.mockResolvedValue({ messages: [] }); - const view = render(); + const view = renderView(); await waitFor(() => expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), ); @@ -177,7 +196,7 @@ it('does not automatically retry an old missing owner row on unrelated submissio it('offers a safe retry after a transient owner-text request failure', async () => { load.mockRejectedValueOnce(new Error('temporary outage')); load.mockResolvedValueOnce({ messages: [original] }); - render(); + renderView(); const retry = await screen.findByRole('button', { name: 'com_ui_private_text_retry' }); expect(screen.getByText(canonical.text)).toBeInTheDocument(); @@ -193,7 +212,7 @@ it('offers a safe retry after a transient owner-text request failure', async () it('does not cache a failed decryption as if it contained a usable original', async () => { load.mockResolvedValueOnce({ messages: [{ ...original, text: undefined }] }); load.mockResolvedValueOnce({ messages: [original] }); - render(); + renderView(); const retry = await screen.findByRole('button', { name: 'com_ui_private_text_retry' }); expect(screen.getByText(canonical.text)).toBeInTheDocument(); await act(async () => retry.click()); @@ -203,7 +222,7 @@ it('does not cache a failed decryption as if it contained a usable original', as it('rejects stale revisions instead of restoring a previous original', async () => { load.mockResolvedValue({ messages: [{ ...original, revision: 'old-revision' }] }); - render(); + renderView(); await waitFor(() => expect(screen.getByRole('status')).toHaveTextContent('com_ui_private_text_unavailable'), ); @@ -219,7 +238,7 @@ it('clears the visible original immediately on account switching, ignoring late }), ); load.mockResolvedValue({ messages: [] }); - const view = render(); + const view = renderView(); mockOwnerId = 'another-owner'; view.rerender(); await act(async () => { @@ -236,7 +255,7 @@ it('batches selected private rows and never loads ordinary messages', async () = messageId: `message-${index}`, })); messages.push({ ...canonical, messageId: 'plain', privacyRevision: undefined }); - render(); + renderView(); await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); expect(load.mock.calls.map(([, ids]) => ids.length)).toEqual([50, 1]); expect(load.mock.calls.flatMap(([, ids]) => ids)).not.toContain('plain'); @@ -255,7 +274,7 @@ it('loads batches concurrently, publishes completed batches, and only fetches ne ...canonical, messageId: `message-${index}`, })); - const view = render(); + const view = renderView(); await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); expect(await screen.findByText(original.text)).toBeInTheDocument(); await act(async () => { @@ -273,7 +292,7 @@ it('loads batches concurrently, publishes completed batches, and only fetches ne it('invalidates an already rendered original when the canonical message changes', async () => { load.mockResolvedValue({ messages: [original] }); - const view = render(); + const view = renderView(); expect(await screen.findByText(original.text)).toBeInTheDocument(); view.rerender(); expect(screen.queryByText(original.text)).not.toBeInTheDocument(); @@ -284,7 +303,7 @@ it('invalidates an already rendered original when the canonical message changes' it('clears originals when tenant identity changes even if the user ID is unchanged', async () => { load.mockResolvedValueOnce({ messages: [original] }); - const view = render(); + const view = renderView(); expect(await screen.findByText(original.text)).toBeInTheDocument(); load.mockResolvedValue({ messages: [] }); mockTenantId = 'tenant-b'; @@ -292,3 +311,50 @@ it('clears originals when tenant identity changes even if the user ID is unchang expect(screen.queryByText(original.text)).not.toBeInTheDocument(); await waitFor(() => expect(load).toHaveBeenCalledTimes(2)); }); + +it('deduplicates concurrent panes and keeps originals out of the application query cache', async () => { + let finish!: (value: { messages: (typeof original)[] }) => void; + load.mockReturnValue( + new Promise((resolve) => { + finish = resolve; + }), + ); + const panes = renderView( + <> + + + , + ); + await waitFor(() => expect(load).toHaveBeenCalledTimes(1)); + await act(async () => finish({ messages: [original] })); + expect(await screen.findAllByText(original.text)).toHaveLength(2); + expect( + JSON.stringify( + applicationClient + .getQueryCache() + .getAll() + .map((query) => query.state.data), + ), + ).not.toContain(original.text); + panes.unmount(); + await waitFor(() => + expect(getOwnerQueryClient(applicationClient).getQueryCache().getAll()).toHaveLength(0), + ); +}); + +it('bounds active requests to three batches and cancels queued work on unmount', async () => { + const finishers: Array<(value: { messages: [] }) => void> = []; + load.mockImplementation(() => new Promise((resolve) => finishers.push(resolve))); + const messages = Array.from({ length: 201 }, (_, index) => ({ + ...canonical, + messageId: `bounded-${index}`, + })); + const view = renderView(); + await waitFor(() => expect(load).toHaveBeenCalledTimes(3)); + view.unmount(); + await act(async () => finishers.forEach((finish) => finish({ messages: [] }))); + expect(load).toHaveBeenCalledTimes(3); + await waitFor(() => + expect(getOwnerQueryClient(applicationClient).getQueryCache().getAll()).toHaveLength(0), + ); +}); diff --git a/client/src/components/Chat/Messages/PrivateText.tsx b/client/src/components/Chat/Messages/PrivateText.tsx index 2a9b3f06475..95e3fee300e 100644 --- a/client/src/components/Chat/Messages/PrivateText.tsx +++ b/client/src/components/Chat/Messages/PrivateText.tsx @@ -1,16 +1,12 @@ -import { - createContext, - lazy, - Suspense, - useContext, - useEffect, - useMemo, - useRef, - useState, -} from 'react'; -import { dataService } from 'librechat-data-provider'; +import { createContext, lazy, Suspense, useContext } from 'react'; +import { QueryClientProvider, useQueryClient } from '@tanstack/react-query'; import type { TMessage } from 'librechat-data-provider'; import type { ReactNode } from 'react'; +import { + OwnerQueryContext, + getOwnerQueryClient, + useOwnerMessageTexts, +} from '~/data-provider/Messages/private'; import { useAuthContext } from '~/hooks/AuthContext'; import { useLocalize } from '~/hooks'; const DisplayMessage = lazy(async () => ({ @@ -22,18 +18,12 @@ interface Original { revision: string; text?: string; } -interface MissingOriginal { - canonicalText: string; - revision: string; - retryOnCompletion: boolean; -} interface OwnerTextState { - scope: string; messages: ReadonlyMap; loading: boolean; retry?: () => void; } -const empty: OwnerTextState = { scope: '', messages: new Map(), loading: false }; +const empty: OwnerTextState = { messages: new Map(), loading: false }; const OwnerTextContext = createContext(empty); interface OwnerTextProviderProps { @@ -56,7 +46,19 @@ export function OwnerTextProvider(props: OwnerTextProviderProps) { protectedMessage.conversationId !== 'new' ? protectedMessage.conversationId : props.conversationId; - return ; + return ; +} + +function OwnerQueries(props: OwnerTextProviderProps) { + const applicationClient = useQueryClient(); + return ( + + + + ); } function ActiveOwnerTextProvider({ @@ -66,137 +68,14 @@ function ActiveOwnerTextProvider({ children, }: OwnerTextProviderProps) { const { user } = useAuthContext(); - const selection = useMemo( - () => - JSON.stringify( - (messages ?? []) - .filter( - (message) => - message.isCreatedByUser && - message.privacyRevision && - message.conversationId === conversationId, - ) - .map((message) => [message.messageId, message.privacyRevision, message.text]) - .sort(), - ), - [messages, conversationId], - ); - const scope = JSON.stringify([user?.id, user?.tenantId, conversationId, selection]); - const [state, setState] = useState(empty); - const [retryAttempt, setRetryAttempt] = useState(0); - const cached = useRef<{ - scope: string; - messages: Map; - missing: Map; - retryAttempt: number; - }>({ scope: '', messages: new Map(), missing: new Map(), retryAttempt: 0 }); - useEffect(() => { - let cancelled = false; - const selected = JSON.parse(selection) as Array<[string, string, string]>; - if (!user?.id || !conversationId || selected.length === 0) { - cached.current = { scope: '', messages: new Map(), missing: new Map(), retryAttempt }; - setState(empty); - return; - } - const ownerScope = JSON.stringify([user.id, user.tenantId, conversationId]); - if (cached.current.scope !== ownerScope) { - cached.current = { scope: ownerScope, messages: new Map(), missing: new Map(), retryAttempt }; - } - const forceRetry = cached.current.retryAttempt !== retryAttempt; - cached.current.retryAttempt = retryAttempt; - const originals = new Map(); - const missing = new Map(); - const pending: Array<[string, string, string]> = []; - for (const [id, revision, text] of selected) { - const prior = cached.current.messages.get(id); - if (prior?.revision === revision && prior.canonicalText === text && prior.text != null) { - originals.set(id, prior); - continue; - } - const missed = cached.current.missing.get(id); - if (missed?.revision === revision && missed.canonicalText === text) { - if (!forceRetry && (isSubmitting || !missed.retryOnCompletion)) { - missing.set(id, missed); - continue; - } - } - pending.push([id, revision, text]); - } - // Do not retain originals or failed reads from removed or edited messages. - cached.current.messages = originals; - cached.current.missing = missing; - const retry = () => setRetryAttempt((attempt) => attempt + 1); - setState({ scope, messages: new Map(originals), loading: pending.length > 0, retry }); - if (pending.length === 0) { - return; - } - let next = 0; - const load = async () => { - const workers = Array.from( - { length: Math.min(3, Math.ceil(pending.length / 50)) }, - async () => { - while (next < pending.length) { - const start = next; - next += 50; - const batch = pending.slice(start, start + 50); - const expected = new Map(batch.map(([id, revision, text]) => [id, { revision, text }])); - try { - const result = await dataService.getOwnerMessageTexts( - conversationId, - batch.map(([id]) => id), - ); - if (cancelled) { - return; - } - for (const message of result.messages) { - const match = expected.get(message.messageId); - if ( - match?.revision === message.revision && - match.text === message.canonicalText && - typeof message.text === 'string' - ) { - const original = { - revision: message.revision, - text: message.text, - canonicalText: message.canonicalText, - }; - originals.set(message.messageId, original); - cached.current.messages.set(message.messageId, original); - } - } - for (const [id, revision, text] of batch) { - if (!originals.has(id)) { - missing.set(id, { - revision, - canonicalText: text, - retryOnCompletion: isSubmitting, - }); - } - } - setState({ scope, messages: new Map(originals), loading: true, retry }); - } catch { - if (cancelled) { - return; - } - for (const [id, revision, text] of batch) { - missing.set(id, { revision, canonicalText: text, retryOnCompletion: isSubmitting }); - } - } - } - }, - ); - await Promise.all(workers); - if (!cancelled) { - setState({ scope, messages: new Map(originals), loading: false, retry }); - } - }; - void load(); - return () => { - cancelled = true; - }; - }, [scope, selection, conversationId, user?.id, user?.tenantId, isSubmitting, retryAttempt]); - const visible = state.scope === scope ? state : empty; - return {children}; + const state = useOwnerMessageTexts({ + messages, + conversationId, + isSubmitting, + userId: user?.id, + tenantId: user?.tenantId, + }); + return {children}; } /** No owner-view data is passed to edit, copy/export, retry, or prompt-building callbacks. */ diff --git a/client/src/data-provider/Messages/private.ts b/client/src/data-provider/Messages/private.ts new file mode 100644 index 00000000000..9afe14ca475 --- /dev/null +++ b/client/src/data-provider/Messages/private.ts @@ -0,0 +1,152 @@ +import { createContext, useEffect, useMemo } from 'react'; +import { QueryKeys, dataService } from 'librechat-data-provider'; +import { QueryClient, useQuery, useQueryClient } from '@tanstack/react-query'; +import type { TMessage } from 'librechat-data-provider'; + +export const OwnerQueryContext = createContext(undefined); +const clients = new WeakMap(); + +/** Shared by panes, but excluded from ordinary caches, devtools and persistence. */ +export function getOwnerQueryClient(applicationClient: QueryClient): QueryClient { + let client = clients.get(applicationClient); + if (client == null) { + client = new QueryClient({ + defaultOptions: { queries: { cacheTime: 0, retry: false, networkMode: 'always' } }, + }); + clients.set(applicationClient, client); + } + return client; +} + +interface Original { + messageId: string; + canonicalText: string; + revision: string; + text: string; +} +interface OwnerTexts { + messages: Original[]; + provisional: boolean; +} + +export function useOwnerMessageTexts({ + messages, + userId, + tenantId, + conversationId, + isSubmitting, +}: { + messages: readonly TMessage[] | null; + userId?: string; + tenantId?: string; + conversationId?: string; + isSubmitting: boolean; +}) { + const client = useQueryClient({ context: OwnerQueryContext }); + const selected = useMemo( + () => + (messages ?? []) + .filter( + (message) => + message.isCreatedByUser && + message.privacyRevision && + typeof message.text === 'string' && + message.conversationId === conversationId, + ) + .map((message) => ({ + messageId: message.messageId, + revision: message.privacyRevision!, + canonicalText: message.text, + })) + .sort((a, b) => a.messageId.localeCompare(b.messageId)), + [messages, conversationId], + ); + const scope = [QueryKeys.ownerMessageTexts, userId, tenantId ?? '', conversationId]; + const queryKey = [...scope, selected]; + const query = useQuery({ + context: OwnerQueryContext, + queryKey, + enabled: Boolean(userId && conversationId && selected.length), + cacheTime: 0, + staleTime: Infinity, + retry: false, + refetchOnWindowFocus: false, + refetchOnReconnect: false, + queryFn: async ({ signal }) => { + const expected = new Map(selected.map((message) => [message.messageId, message])); + const originals = new Map(); + // Reuse only validated rows still held by active queries in this exact owner scope. + for (const [, prior] of client.getQueriesData(scope)) { + for (const message of prior?.messages ?? []) { + const match = expected.get(message.messageId); + if ( + match?.revision === message.revision && + match.canonicalText === message.canonicalText + ) { + originals.set(message.messageId, message); + } + } + } + const pending = selected.filter((message) => !originals.has(message.messageId)); + const result = (): OwnerTexts => ({ + messages: [...originals.values()], + provisional: isSubmitting && originals.size < selected.length, + }); + const publish = () => { + if (!signal?.aborted) { + client.setQueryData(queryKey, result()); + } + }; + publish(); + let next = 0; + let failed = false; + // Transport batches are bounded; React Query owns deduplication, cancellation and retry. + await Promise.all( + Array.from({ length: Math.min(3, Math.ceil(pending.length / 50)) }, async () => { + while (next < pending.length && !signal?.aborted) { + const batch = pending.slice(next, (next += 50)); + try { + const response = await dataService.getOwnerMessageTexts( + conversationId!, + batch.map((message) => message.messageId), + ); + if (signal?.aborted) { + return; + } + const batchIds = new Set(batch.map((message) => message.messageId)); + for (const message of response.messages) { + const match = expected.get(message.messageId); + if ( + batchIds.has(message.messageId) && + match?.revision === message.revision && + match.canonicalText === message.canonicalText && + typeof message.text === 'string' + ) { + originals.set(message.messageId, { ...message, text: message.text }); + } + } + publish(); + } catch { + failed = true; + } + } + }), + ); + if (failed) { + // Never retain HTTP payloads or upstream error text in a query error. + throw new Error('owner_text_unavailable'); + } + return result(); + }, + }); + useEffect(() => { + if (!isSubmitting && query.data?.provisional && !query.isFetching) { + void query.refetch(); + } + }, [isSubmitting, query.data?.provisional, query.isFetching, query.refetch]); + const originals = useMemo( + () => new Map((query.data?.messages ?? []).map((message) => [message.messageId, message])), + [query.data], + ); + return { messages: originals, loading: query.isFetching, retry: () => void query.refetch() }; +} diff --git a/packages/api/src/middleware/modelBoundContent.spec.ts b/packages/api/src/middleware/modelBoundContent.spec.ts index e708eedbb08..17388cc3a00 100644 --- a/packages/api/src/middleware/modelBoundContent.spec.ts +++ b/packages/api/src/middleware/modelBoundContent.spec.ts @@ -6197,3 +6197,40 @@ describe('assertModelBoundProviderContent', () => { expect(onAllowed).not.toHaveBeenCalled(); }); }); + +it('admits only after every native batch passes and awaits the protected commit', async () => { + const onContentAllowed = jest.fn(async () => {}); + const callback = createModelBoundChatModelCallback( + { + filters: { + messages: { + pii: { + fields: ['text', 'content_part'], + starterPatterns: [], + customPatterns: [{ id: 'secret', label: 'Secret', regex: 'BLOCK-ME' }], + }, + }, + }, + }, + { onContentAllowed }, + ); + expect(() => + callback.handleChatModelStart(undefined, [ + [{ role: 'user', content: 'safe' }], + [{ role: 'user', content: 'BLOCK-ME' }], + ]), + ).toThrow(); + expect(onContentAllowed).not.toHaveBeenCalled(); + let finish!: () => void; + onContentAllowed.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const pending = callback.handleChatModelStart(undefined, [[{ role: 'user', content: 'safe' }]]); + expect(pending).toBeInstanceOf(Promise); + finish(); + await pending; + expect(onContentAllowed).toHaveBeenCalledTimes(1); +}); diff --git a/packages/api/src/middleware/modelBoundContent.ts b/packages/api/src/middleware/modelBoundContent.ts index 74388e3256f..53f8a9dcc9d 100644 --- a/packages/api/src/middleware/modelBoundContent.ts +++ b/packages/api/src/middleware/modelBoundContent.ts @@ -681,7 +681,7 @@ export interface ModelBoundChatModelCallback { readonly handleChatModelStart: ( llm: object | undefined, messageBatches: readonly (readonly ModelBoundProviderMessage[])[], - ) => void; + ) => void | Promise; } export interface InitialModelBoundAdmissionCallback { @@ -3238,7 +3238,10 @@ function snapshotBoundedSourceFileIds( /** Creates a run-stable callback shared by root, summary, and subagent model clients. */ export function createModelBoundChatModelCallback( input: Omit, - options: { readonly onContentRejected?: (error: unknown) => void } = {}, + options: { + readonly onContentRejected?: (error: unknown) => void; + readonly onContentAllowed?: () => void | Promise; + } = {}, ): ModelBoundChatModelCallback { const storedMessageSnapshot = snapshotBoundedProviderArray(input.storedMessages); const resolvedFileSnapshot = snapshotBoundedProviderArray(input.resolvedFiles); @@ -3323,6 +3326,7 @@ export function createModelBoundChatModelCallback( options.onContentRejected?.(error); throw new FatalModelBoundPolicyError(error); } + return options.onContentAllowed?.(); }, }); return callback; diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 06a940ebc6d..aa00a1e9d66 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -171,7 +171,7 @@ describe('private text submission boundary', () => { ]); const store = { saveMessage, - getPersistedPrivateTextId: jest.fn(async () => 'protected-row-id'), + getPersistedPrivateTextId: jest.fn(async (): Promise => 'protected-row-id'), getPrivateMessageTexts, }; const finalEvent = { requestMessage: { messageId: message.messageId, privacyRevision: '' } }; @@ -186,15 +186,41 @@ describe('private text submission boundary', () => { finalEvent, ), ).toEqual({ _id: undefined }); - expect(saveMessage).not.toHaveBeenCalled(); + expect(saveMessage).toHaveBeenCalledWith({ userId: 'owner' }, older, { insertOnly: true }); expect(finalEvent.requestMessage.privacyRevision).toBe(message.privacyRevision); - getPrivateMessageTexts.mockResolvedValueOnce([]); + store.getPersistedPrivateTextId.mockResolvedValueOnce(null); await expect( saveAbortedUserMessage(store, { userId: 'owner' }, older, undefined, 'tenant-a'), ).rejects.toThrow('private value'); - expect(saveMessage).not.toHaveBeenCalled(); + expect(saveMessage).toHaveBeenCalledTimes(2); }); + it.each([true, false])( + 'stops a revisionless literal placeholder, existing row: %s', + async (existing) => { + const message = { + messageId: 'literal-token', + conversationId: 'conversation', + isCreatedByUser: true, + text: `[EMAIL_1_${'a'.repeat(32)}]`, + }; + const saveMessage = jest.fn(async () => + Object.assign(message as IMessage, { _id: existing ? 'existing' : 'inserted' }), + ); + const getPersistedPrivateTextId = jest.fn(async () => null); + expect( + await saveAbortedUserMessage( + { saveMessage, getPersistedPrivateTextId }, + { userId: 'owner' }, + message, + undefined, + ), + ).toEqual({ _id: existing ? 'existing' : 'inserted' }); + expect(saveMessage).toHaveBeenCalledWith({ userId: 'owner' }, message, { insertOnly: true }); + expect(getPersistedPrivateTextId).not.toHaveBeenCalled(); + }, + ); + it('stamps a protected preliminary job message before the created event', () => { const { req, message } = submit(); const preliminary = stampPreliminaryPrivateTextMessage(req, { diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 0b62163a3a5..1f26ec24bbf 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -25,7 +25,6 @@ interface Capture { } const captures = new WeakMap(); -const PRIVATE_PLACEHOLDER = /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]/; const CONTROL_ROUTES = new Set([ 'abort', @@ -373,6 +372,7 @@ export async function savePrivateTextMessage( export async function requirePrivateTextPersistence( req: object | undefined, start: () => Promise<{ message?: PrivateTextMessage | null } | undefined>, + onPersisted?: () => void, ): Promise { const capture = req == null ? undefined : captures.get(req); if (capture == null) { @@ -385,6 +385,43 @@ export async function requirePrivateTextPersistence( ) { throw unavailable(); } + onPersisted?.(); +} + +/** Ordinary startup stays immediate; protected revisions are announced only after admission. */ +export function deferPrivateTextStart( + req: object | undefined, + onStart: ((message: PrivateTextMessage, responseId: string, isNew: boolean) => void) | undefined, + message: PrivateTextMessage, + responseId: string, + isNew: boolean, +): (() => void) | undefined { + if (typeof onStart !== 'function') { + return; + } + if (req == null || !captures.has(req)) { + onStart(message, responseId, isNew); + return; + } + let announced = false; + return () => { + if (!announced) { + announced = true; + onStart(message, responseId, isNew); + } + }; +} + +/** Runs only after the exact native payload passes its policy callback. */ +export function getPrivateTextAdmission( + req: object | undefined, + start: (() => Promise<{ message?: PrivateTextMessage | null } | undefined>) | undefined, + onPersisted?: () => void, +): (() => Promise) | undefined { + if (req == null || !captures.has(req) || start == null) { + return; + } + return () => requirePrivateTextPersistence(req, start, onPersisted); } /** @@ -393,10 +430,7 @@ export async function requirePrivateTextPersistence( * omit the revision, so their prerequisite is an insert-only write. */ export async function saveAbortedUserMessage( - store: Pick< - MessageMethods, - 'saveMessage' | 'getPersistedPrivateTextId' | 'getPrivateMessageTexts' - >, + store: Pick, ctx: Parameters[0], message: Parameters[1], metadata: Parameters[2], @@ -408,57 +442,35 @@ export async function saveAbortedUserMessage( if (!message.messageId || !message.conversationId) { throw unavailable(); } - if (typeof message.text !== 'string' || !PRIVATE_PLACEHOLDER.test(message.text)) { - const saved = await store.saveMessage(ctx, message, { ...metadata, insertOnly: true }); - if ( - saved == null || - saved.messageId !== message.messageId || - saved.conversationId !== message.conversationId || - saved.text !== message.text - ) { + const saved = await store.saveMessage(ctx, message, { ...metadata, insertOnly: true }); + if ( + saved == null || + saved.messageId !== message.messageId || + saved.conversationId !== message.conversationId || + saved.text !== message.text + ) { + throw unavailable(); + } + if (typeof saved.privacyRevision === 'string' && saved.privacyRevision.length > 0) { + if (typeof saved.text !== 'string') { throw unavailable(); } - if (typeof saved.privacyRevision === 'string' && saved.privacyRevision.length > 0) { - if (typeof saved.text !== 'string') { - throw unavailable(); - } - const exists = await store.getPersistedPrivateTextId({ - userId: ctx.userId, - tenantId, - conversationId: message.conversationId, - messageId: message.messageId, - text: saved.text, - privacyRevision: saved.privacyRevision, - }); - if (!exists) { - throw unavailable(); - } - if (finalEvent?.requestMessage != null) { - finalEvent.requestMessage.privacyRevision = saved.privacyRevision; - } + const exists = await store.getPersistedPrivateTextId({ + userId: ctx.userId, + tenantId, + conversationId: message.conversationId, + messageId: message.messageId, + text: saved.text, + privacyRevision: saved.privacyRevision, + }); + if (!exists) { + throw unavailable(); + } + if (finalEvent?.requestMessage != null) { + finalEvent.requestMessage.privacyRevision = saved.privacyRevision; } - return { _id: saved._id }; - } - const rows = await store.getPrivateMessageTexts({ - userId: ctx.userId, - tenantId, - conversationId: message.conversationId, - messageIds: [message.messageId], - }); - const row = rows.find( - (candidate) => - candidate.messageId === message.messageId && - candidate.text === message.text && - candidate.privacyRevision && - candidate.privateText, - ); - if (row == null) { - throw unavailable(); - } - if (finalEvent?.requestMessage != null) { - finalEvent.requestMessage.privacyRevision = row.privacyRevision; } - return { _id: row._id }; + return { _id: saved._id }; } if (!message.messageId || !message.conversationId || typeof message.text !== 'string') { throw unavailable(); diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index f81b97e5e7d..ede3cfc0770 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -1,5 +1,6 @@ export enum QueryKeys { messages = 'messages', + ownerMessageTexts = 'ownerMessageTexts', sharedMessages = 'sharedMessages', sharedStartupConfig = 'sharedStartupConfig', sharedLinks = 'sharedLinks', From 4d2804ec0a8f54c6026f928dee0db87d761a3e48 Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 02:55:43 +0000 Subject: [PATCH 13/18] fix: Track Private Query Effect Dependencies --- client/src/data-provider/Messages/private.ts | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/client/src/data-provider/Messages/private.ts b/client/src/data-provider/Messages/private.ts index 9afe14ca475..3f4f05e4d11 100644 --- a/client/src/data-provider/Messages/private.ts +++ b/client/src/data-provider/Messages/private.ts @@ -139,14 +139,15 @@ export function useOwnerMessageTexts({ return result(); }, }); + const { data, isFetching, refetch } = query; useEffect(() => { - if (!isSubmitting && query.data?.provisional && !query.isFetching) { - void query.refetch(); + if (!isSubmitting && data?.provisional && !isFetching) { + void refetch(); } - }, [isSubmitting, query.data?.provisional, query.isFetching, query.refetch]); + }, [isSubmitting, data?.provisional, isFetching, refetch]); const originals = useMemo( - () => new Map((query.data?.messages ?? []).map((message) => [message.messageId, message])), - [query.data], + () => new Map((data?.messages ?? []).map((message) => [message.messageId, message])), + [data], ); - return { messages: originals, loading: query.isFetching, retry: () => void query.refetch() }; + return { messages: originals, loading: isFetching, retry: () => void refetch() }; } From e8621b7caf30d1b4d780e8fbfff640b854310ce3 Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 12:04:42 +0000 Subject: [PATCH 14/18] fix: Fence Protected Admission and Screenshot Targets --- api/app/clients/specs/BaseClient.test.js | 37 ++- .../components/Chat/Messages/MessagesView.tsx | 7 +- .../Messages/__tests__/MessagesView.spec.tsx | 12 + .../__tests__/screenshot.spec.tsx | 238 ++++++++++++++++++ .../__tests__/useExportConversation.spec.tsx | 9 +- .../Conversations/useExportConversation.ts | 38 ++- client/src/hooks/ScreenshotContext.tsx | 37 ++- e2e/specs/mock/private-text.spec.ts | 39 +++ .../src/protection/private/submission.spec.ts | 10 + .../api/src/protection/private/submission.ts | 5 +- 10 files changed, 412 insertions(+), 20 deletions(-) create mode 100644 client/src/hooks/Conversations/__tests__/screenshot.spec.tsx diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 947dafafd17..e9937e7b497 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -1784,10 +1784,10 @@ describe('BaseClient', () => { privateTextTokens: getPrivateTextInspectionTokens(client.modelBoundStoredMessages), }, { - onContentRejected: client.modelBoundUserMessagePersistence.cancel, + onContentRejected: client.modelBoundUserMessagePersistence?.cancel, onContentAllowed: getPrivateTextAdmission( req, - client.modelBoundUserMessagePersistence.start, + client.modelBoundUserMessagePersistence?.start, client.privateTextStart, ), }, @@ -1834,6 +1834,39 @@ describe('BaseClient', () => { 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, { diff --git a/client/src/components/Chat/Messages/MessagesView.tsx b/client/src/components/Chat/Messages/MessagesView.tsx index 8e0e80c9616..97b4dd9cc8f 100644 --- a/client/src/components/Chat/Messages/MessagesView.tsx +++ b/client/src/components/Chat/Messages/MessagesView.tsx @@ -161,7 +161,12 @@ function MessagesViewContent({
) : ( <> -
+
{FLAT_THREAD && threadRows ? ( { expect(screen.getByTestId('pending-steers')).toHaveAttribute('data-conversation-id', 'convo-2'); }); + it('pins the screenshot identity to the rendered transcript and replaces it on navigation', () => { + const view = render(); + const initialTarget = screen.getByTestId('screenshot-target'); + expect(initialTarget).toHaveAttribute('data-conversation-id', 'convo-1'); + const destinationTree = [{ ...messageTree[0], conversationId: 'convo-2' }]; + view.rerender(); + const destinationTarget = screen.getByTestId('screenshot-target'); + expect(destinationTarget).toHaveAttribute('data-conversation-id', 'convo-2'); + expect(destinationTarget).not.toBe(initialTarget); + expect(initialTarget.isConnected).toBe(false); + }); + it('keeps recovery visible while the message tree is temporarily empty', () => { render(); diff --git a/client/src/hooks/Conversations/__tests__/screenshot.spec.tsx b/client/src/hooks/Conversations/__tests__/screenshot.spec.tsx new file mode 100644 index 00000000000..3a287d9244c --- /dev/null +++ b/client/src/hooks/Conversations/__tests__/screenshot.spec.tsx @@ -0,0 +1,238 @@ +import { Provider } from 'jotai'; +import download from 'downloadjs'; +import { toCanvas } from 'html-to-image'; +import { QueryKeys } from 'librechat-data-provider'; +import { act, renderHook } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { MemoryRouter, Route, Routes, useNavigate, useParams } from 'react-router-dom'; +import type { TMessage, TConversation } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import { completeProgressiveRowMounts } from '~/hooks/Messages/useProgressiveRowMount'; +import { ScreenshotProvider, useScreenshot } from '~/hooks/ScreenshotContext'; +import useExportConversation from '../useExportConversation'; + +const mockShowToast = jest.fn(); +jest.mock('@librechat/client', () => ({ + ...jest.requireActual('@librechat/client'), + useToastContext: () => ({ showToast: mockShowToast }), +})); +jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); +jest.mock('downloadjs', () => jest.fn()); +jest.mock('html-to-image', () => ({ toCanvas: jest.fn() })); +jest.mock('~/hooks/Messages/useProgressiveRowMount', () => ({ + completeProgressiveRowMounts: jest.fn(), +})); + +const sourceId = '11111111-1111-4111-8111-111111111111'; +const destinationId = '22222222-2222-4222-8222-222222222222'; +const clean = { + conversationId: sourceId, + messageId: 'safe', + text: 'Safe text', + isCreatedByUser: true, + createdAt: '2026-10-02T11:00:00', +} as TMessage; + +function Target() { + const { conversationId } = useParams(); + const { screenshotTargetRef } = useScreenshot(); + return ( +
+ ); +} +function setup() { + const client = new QueryClient(); + client.setQueryData([QueryKeys.messages, sourceId], [clean]); + client.setQueryData( + [QueryKeys.messages, destinationId], + [{ ...clean, conversationId: destinationId, privacyRevision: 'protected' }], + ); + const hook = renderHook( + () => { + const { conversationId } = useParams(); + const navigate = useNavigate(); + const { exportConversation } = useExportConversation({ + conversation: { conversationId } as TConversation, + filename: 'screenshot', + type: 'screenshot', + includeOptions: false, + exportBranches: false, + recursive: false, + }); + return { navigate, exportConversation }; + }, + { + wrapper: ({ children }: { children: ReactNode }) => ( + + + + + + + + {children} + + } + /> + + + + + + ), + }, + ); + const target = document.querySelector('[data-conversation-id]')!; + Object.defineProperties(target, { + scrollWidth: { value: 100 }, + scrollHeight: { value: 100 }, + }); + return { ...hook, client, target }; +} +function canvas() { + const element = document.createElement('canvas'); + jest.spyOn(element, 'toBlob').mockImplementation((callback) => { + callback(new Blob(['image'], { type: 'image/png' })); + }); + return element; +} +beforeEach(() => { + jest.mocked(completeProgressiveRowMounts).mockResolvedValue(undefined); + jest.mocked(toCanvas).mockResolvedValue(canvas()); +}); + +it('downloads an unchanged unprotected capture', async () => { + const { result, target } = setup(); + await act(async () => { + await result.current.exportConversation(); + }); + expect(toCanvas).toHaveBeenCalledWith(target, expect.anything()); + expect(download).toHaveBeenCalledTimes(1); +}); + +it('rejects navigation to a protected transcript during progressive mounting', async () => { + let finish!: () => void; + jest.mocked(completeProgressiveRowMounts).mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result, target } = setup(); + const exporting = result.current.exportConversation(); + act(() => result.current.navigate(`/c/${destinationId}`)); + const destination = document.querySelector('[data-conversation-id]')!; + Object.defineProperties(destination, { + scrollWidth: { value: 100 }, + scrollHeight: { value: 100 }, + }); + expect(target.isConnected).toBe(false); + await act(async () => { + finish(); + await exporting; + }); + expect(toCanvas).not.toHaveBeenCalled(); + expect(download).not.toHaveBeenCalled(); + expect(mockShowToast).toHaveBeenCalled(); +}); + +it('discards a capture when the target changes during image cloning', async () => { + let finish!: (result: HTMLCanvasElement) => void; + jest.mocked(toCanvas).mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result } = setup(); + const exporting = result.current.exportConversation(); + await act(async () => { + await Promise.resolve(); + }); + expect(toCanvas).toHaveBeenCalledTimes(1); + act(() => result.current.navigate(`/c/${destinationId}`)); + await act(async () => { + finish(canvas()); + await exporting; + }); + expect(download).not.toHaveBeenCalled(); +}); + +it('checks newly protected text after progressive mounting and before cloning', async () => { + let finish!: () => void; + jest.mocked(completeProgressiveRowMounts).mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result, client } = setup(); + const exporting = result.current.exportConversation(); + act(() => + client.setQueryData( + [QueryKeys.messages, sourceId], + [{ ...clean, privacyRevision: 'protected' }], + ), + ); + await act(async () => { + finish(); + await exporting; + }); + expect(toCanvas).not.toHaveBeenCalled(); + expect(download).not.toHaveBeenCalled(); + expect(mockShowToast).toHaveBeenCalledWith( + expect.objectContaining({ + message: 'com_nav_export_screenshot_private_text', + }), + ); +}); + +it('refuses a target whose rendered conversation differs from the export scope', async () => { + const { result, target } = setup(); + target.dataset.conversationId = destinationId; + await act(async () => { + await result.current.exportConversation(); + }); + expect(toCanvas).not.toHaveBeenCalled(); + expect(download).not.toHaveBeenCalled(); +}); + +it('refuses a reused DOM node when its transcript identity changes during mounting', async () => { + let finish!: () => void; + jest.mocked(completeProgressiveRowMounts).mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result, target } = setup(); + const exporting = result.current.exportConversation(); + target.dataset.conversationId = destinationId; + await act(async () => { + finish(); + await exporting; + }); + expect(toCanvas).not.toHaveBeenCalled(); + expect(download).not.toHaveBeenCalled(); +}); + +it('does not revive a cancelled capture after navigating away and back', async () => { + let finish!: (result: HTMLCanvasElement) => void; + jest.mocked(toCanvas).mockReturnValueOnce( + new Promise((resolve) => { + finish = resolve; + }), + ); + const { result, target } = setup(); + const exporting = result.current.exportConversation(); + await act(async () => { + await Promise.resolve(); + }); + act(() => result.current.navigate(`/c/${destinationId}`)); + act(() => result.current.navigate(`/c/${sourceId}`)); + expect(document.querySelector('[data-conversation-id]')).not.toBe(target); + await act(async () => { + finish(canvas()); + await exporting; + }); + expect(download).not.toHaveBeenCalled(); +}); diff --git a/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx b/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx index af3dfca614e..715739472cc 100644 --- a/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx +++ b/client/src/hooks/Conversations/__tests__/useExportConversation.spec.tsx @@ -12,6 +12,7 @@ import useExportConversation from '../useExportConversation'; const mockGetMessages = jest.fn(); const mockShowToast = jest.fn(); const mockCaptureScreenshot = jest.fn(); +const mockScreenshotRef = { current: document.createElement('div') }; jest.mock('librechat-data-provider', () => { const actual = jest.requireActual('librechat-data-provider'); @@ -29,8 +30,12 @@ jest.mock('@librechat/client', () => ({ })); jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key })); jest.mock('~/hooks/ScreenshotContext', () => ({ - useScreenshot: () => ({ captureScreenshot: mockCaptureScreenshot }), + useScreenshot: () => ({ + captureScreenshot: mockCaptureScreenshot, + screenshotTargetRef: mockScreenshotRef, + }), ScreenshotLimitError: class ScreenshotLimitError extends Error {}, + ScreenshotTargetError: class ScreenshotTargetError extends Error {}, })); jest.mock('downloadjs', () => jest.fn()); jest.mock('export-from-json', () => @@ -96,6 +101,8 @@ async function readBlob(blob: Blob): Promise { beforeEach(() => { jest.clearAllMocks(); + mockScreenshotRef.current.dataset.conversationId = conversationId; + document.body.append(mockScreenshotRef.current); }); it('exports acknowledged canonical text rather than a pending private cache value', async () => { diff --git a/client/src/hooks/Conversations/useExportConversation.ts b/client/src/hooks/Conversations/useExportConversation.ts index 80900ebda53..5fa8986ee39 100644 --- a/client/src/hooks/Conversations/useExportConversation.ts +++ b/client/src/hooks/Conversations/useExportConversation.ts @@ -6,7 +6,11 @@ import { useToastContext } from '@librechat/client'; import { useQueryClient } from '@tanstack/react-query'; import { buildTree, Constants, dataService, QueryKeys } from 'librechat-data-provider'; import type { TConversation, TMessage, TPreset } from 'librechat-data-provider'; -import { ScreenshotLimitError, useScreenshot } from '~/hooks/ScreenshotContext'; +import { + ScreenshotLimitError, + ScreenshotTargetError, + useScreenshot, +} from '~/hooks/ScreenshotContext'; import useBuildMessageTree from '~/hooks/Messages/useBuildMessageTree'; import { isUnacknowledgedUserMessage } from '~/utils/messages'; import { NotificationSeverity } from '~/common'; @@ -37,17 +41,19 @@ export default function useExportConversation({ }) { const queryClient = useQueryClient(); const { showToast } = useToastContext(); - const { captureScreenshot } = useScreenshot(); + const { captureScreenshot, screenshotTargetRef } = useScreenshot(); const buildMessageTree = useBuildMessageTree(); const localize = useLocalize(); const { conversationId: paramId } = useParams(); + const screenshotConversationId = + paramId === 'new' ? paramId : (conversation?.conversationId ?? paramId ?? ''); const getCachedMessages = useCallback(() => { - const queryParam = - paramId === 'new' ? paramId : (conversation?.conversationId ?? paramId ?? ''); - return queryClient.getQueryData([QueryKeys.messages, queryParam]) ?? []; - }, [paramId, conversation?.conversationId, queryClient]); + return ( + queryClient.getQueryData([QueryKeys.messages, screenshotConversationId]) ?? [] + ); + }, [screenshotConversationId, queryClient]); const getMessageTree = useCallback(async () => { const conversationId = conversation?.conversationId; @@ -85,11 +91,27 @@ export default function useExportConversation({ if (refuseUnsafeScreenshot()) { return; } + const target = screenshotTargetRef?.current; let data: Blob; try { - data = await captureScreenshot(); + data = await captureScreenshot( + (node) => + node === target && + node.dataset.conversationId === screenshotConversationId && + !screenshotWouldExposePrivateText(), + ); + if ( + !target?.isConnected || + screenshotTargetRef?.current !== target || + target.dataset.conversationId !== screenshotConversationId + ) { + throw new ScreenshotTargetError(); + } } catch (err) { - console.error('Failed to capture screenshot', err); + if (refuseUnsafeScreenshot()) { + return; + } + console.error('Failed to capture screenshot'); showToast({ message: localize( err instanceof ScreenshotLimitError diff --git a/client/src/hooks/ScreenshotContext.tsx b/client/src/hooks/ScreenshotContext.tsx index e99c177a5c2..6a0ddc43d6f 100644 --- a/client/src/hooks/ScreenshotContext.tsx +++ b/client/src/hooks/ScreenshotContext.tsx @@ -23,6 +23,13 @@ export class ScreenshotLimitError extends Error { } } +export class ScreenshotTargetError extends Error { + constructor() { + super('Screenshot target changed or is unavailable.'); + this.name = 'ScreenshotTargetError'; + } +} + const ScreenshotContext = createContext({}); export const useScreenshot = () => { @@ -78,17 +85,33 @@ export const useScreenshot = () => { return blob; }; - const captureScreenshot = async (): Promise => { + const captureScreenshot = async (canCapture?: (node: HTMLElement) => boolean): Promise => { if (ref instanceof Function) { throw new Error('Ref callback is not supported.'); } - /** A capture taken while a long thread is still progressively mounting - * would clone a truncated DOM; force the remaining rows in first. */ - await completeProgressiveRowMounts(); - if (ref?.current) { - return takeScreenShot(ref.current); + const node = ref?.current; + if (!node) { + throw new ScreenshotTargetError(); } - throw new Error('Ref is not attached to any element.'); + const conversationId = node.dataset.conversationId; + const assertTarget = () => { + if ( + !conversationId || + !node.isConnected || + ref?.current !== node || + node.dataset.conversationId !== conversationId || + canCapture?.(node) === false + ) { + throw new ScreenshotTargetError(); + } + }; + assertTarget(); + /** Pin the transcript before mounting or cloning can yield to navigation. */ + await completeProgressiveRowMounts(); + assertTarget(); + const image = await takeScreenShot(node); + assertTarget(); + return image; }; return { screenshotTargetRef: ref, captureScreenshot }; diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index 4f8189b2a90..8883c6dc8cb 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -267,3 +267,42 @@ test('the first protected owner view loads from its server ID while generation i } } }); + +test('an unchanged unprotected transcript still downloads a PNG screenshot', async ({ + page, + request, +}) => { + const token = await loginAdmin(request); + let conversationId: string | undefined; + try { + await page.goto('/c/new'); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const response = await sendMessageAndWaitForCompletion(page, 'Ordinary screenshot capture'); + conversationId = (await response.json()).conversationId as string; + await expect(page.getByTestId('screenshot-target')).toHaveAttribute( + 'data-conversation-id', + conversationId!, + ); + await page.getByRole('button', { name: 'Export/Share' }).click(); + await page.getByRole('menuitem', { name: 'Export' }).click(); + const dialog = page.getByRole('dialog', { name: 'Export conversation' }); + await dialog.getByTestId('dropdown-menu').click(); + await page.getByRole('option', { name: 'screenshot (.png)' }).click(); + const [image] = await Promise.all([ + page.waitForEvent('download'), + dialog.getByRole('button', { name: 'Export', exact: true }).click(), + ]); + expect(image.suggestedFilename()).toMatch(/\.png$/); + const bytes = await readFile(await image.path()); + expect([...bytes.subarray(0, 8)]).toEqual([137, 80, 78, 71, 13, 10, 26, 10]); + } finally { + if (conversationId) { + await requestResult(request, { + path: '/api/convos', + token, + method: 'DELETE', + data: { arg: { conversationId } }, + }); + } + } +}); diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index aa00a1e9d66..68474604485 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -11,6 +11,7 @@ import { isPreDenialTextSubmission, isPrivateTextChatSubmission, getPreinspectedPrivateText, + getPrivateTextAdmission, getPrivateTextInspectionTokens, privateTextBinding, } from './submission'; @@ -442,6 +443,15 @@ describe('private text submission boundary', () => { ).rejects.toThrow(); }); + it('requires a persistence gate for a captured turn before allowing any model callback', () => { + const { req } = submit(); + expect(() => getPrivateTextAdmission(req, undefined)).toThrow( + expect.objectContaining({ code: 'content_filter_block' }), + ); + expect(getPrivateTextAdmission(undefined, undefined)).toBeUndefined(); + expect(getPrivateTextAdmission({}, undefined)).toBeUndefined(); + }); + it('rejects stale or swallowed persistence results, including a duplicate ID with different text', async () => { const { req, message } = submit(); const save: MessageMethods['saveMessage'] = jest.fn(async () => undefined); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 1f26ec24bbf..5a978170257 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -418,9 +418,12 @@ export function getPrivateTextAdmission( start: (() => Promise<{ message?: PrivateTextMessage | null } | undefined>) | undefined, onPersisted?: () => void, ): (() => Promise) | undefined { - if (req == null || !captures.has(req) || start == null) { + if (req == null || !captures.has(req)) { return; } + if (start == null) { + throw unavailable(); + } return () => requirePrivateTextPersistence(req, start, onPersisted); } From 6a30e2ae0d5ff901826a2109b98f35df43070d5f Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 13:16:05 +0000 Subject: [PATCH 15/18] fix: Share Protected Admission and Preserve Native Token Trust --- api/app/clients/BaseClient.js | 5 +- api/app/clients/specs/BaseClient.test.js | 2 +- api/server/controllers/agents/client.js | 11 +- api/server/controllers/agents/client.test.js | 55 +++++++++ api/server/utils/import/fork.js | 90 ++++++++------ api/server/utils/import/fork.spec.js | 111 ++++++++++++++++- api/server/utils/import/importBatchBuilder.js | 5 +- api/server/utils/import/importers.js | 1 + e2e/specs/mock/private-text.spec.ts | 112 ++++++++++++++++++ packages/api/src/agents/memory.spec.ts | 86 ++++++++++++++ packages/api/src/agents/memory.ts | 14 +++ packages/api/src/protection/index.ts | 1 + packages/api/src/protection/private/copy.ts | 99 ++++++++++++++++ .../src/protection/private/submission.spec.ts | 39 ++++++ .../api/src/protection/private/submission.ts | 79 ++++++++++-- packages/api/src/shared-links/protection.ts | 1 + .../src/methods/message.privateText.spec.ts | 75 ++++++++++++ packages/data-schemas/src/methods/message.ts | 52 ++++++-- .../data-schemas/src/methods/share.test.ts | 1 + packages/data-schemas/src/methods/share.ts | 9 +- packages/data-schemas/src/schema/message.ts | 1 + packages/data-schemas/src/types/message.ts | 1 + 22 files changed, 787 insertions(+), 63 deletions(-) create mode 100644 packages/api/src/protection/private/copy.ts diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 694c4e8e31e..bb5b5841177 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -36,6 +36,7 @@ const { stampPrivateTextMessage, deferPrivateTextStart, requirePrivateTextPersistence, + rejectPrivateTextAdmission, persistedReasoningOverrideFields, } = require('@librechat/api'); const { @@ -1053,6 +1054,7 @@ class BaseClient { completionResult = await this.sendCompletion(payload, opts); } catch (error) { + rejectPrivateTextAdmission(this.options.req); if (userMessagePersistence?.isPending()) { if (isContentFilterError(error)) { userMessagePersistence.cancel(); @@ -1287,7 +1289,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); diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index e9937e7b497..4b9a618ef71 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -1602,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'); }); diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index f5c39b77177..1996670c139 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -160,6 +160,8 @@ const { createModelBoundChatModelCallback: createModelBoundContentCallback, getPrivateTextInspectionTokens, getPrivateTextAdmission, + requirePrivateTextAdmission, + rejectPrivateTextAdmission, createInitialModelBoundAdmissionCallback, hasModelBoundContentProtection, assertResumeRuntimeContentAllowed, @@ -3505,6 +3507,7 @@ class AgentClient extends BaseClient { */ async runMemory(messages) { try { + await requirePrivateTextAdmission(this.options.req, this.abortController?.signal); if (this.processMemory == null) { return; } @@ -3576,7 +3579,12 @@ class AgentClient extends BaseClient { }); } const bufferMessage = new HumanMessage(limitedMemoryInput); - return await this.processMemory([bufferMessage], filteredMessages); + return await this.processMemory( + [bufferMessage], + filteredMessages, + getPrivateTextInspectionTokens(this.modelBoundStoredMessages ?? []), + this.abortController?.signal, + ); } catch (error) { logger.error('Memory Agent failed to process memory', getSafeErrorMetadata(error)); } @@ -5210,6 +5218,7 @@ class AgentClient extends BaseClient { }); } } finally { + rejectPrivateTextAdmission(this.options.req); /** An aborted/erroring run can still have completed compaction before * the failure; retain that model-visible state for actor reconciliation. */ await this.options.runFiles?.close(); diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index c03c096cd69..c0009e4f236 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -7509,6 +7509,61 @@ describe('AgentClient - titleConvo', () => { client.responseMessageId = 'response-123'; }); + it.each(['failure', 'reject', 'abort', 'success'])( + 'gates automatic extraction on protected admission: %s', + async (outcome) => { + const { HumanMessage } = require('@librechat/agents/langchain/messages'); + const api = require('@librechat/api'); + const filters = { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }; + mockReq.body = { text: 'Remember alice@example.com', clientRequestId: 'memory-private' }; + mockReq.path = '/'; + api.createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(mockReq, { status: jest.fn().mockReturnThis(), json: jest.fn() }, jest.fn()); + const message = api.stampPrivateTextMessage(mockReq, { + text: mockReq.body.text, + isCreatedByUser: true, + conversationId: 'conversation', + messageId: 'user-message', + }); + client.setModelBoundStoredMessages([message]); + client.abortController = new AbortController(); + const extraction = client.runMemory([new HumanMessage(message.text)]); + await Promise.resolve(); + expect(mockProcessMemory).not.toHaveBeenCalled(); + if (outcome === 'abort') { + client.abortController.abort(); + } else if (outcome === 'reject') { + api.rejectPrivateTextAdmission(mockReq); + } else if (outcome === 'failure') { + await expect(api.getPrivateTextAdmission(mockReq, async () => ({}))()).rejects.toThrow(); + } else { + await api.getPrivateTextAdmission(mockReq, async () => ({ message }))(); + } + await extraction; + expect(mockProcessMemory).toHaveBeenCalledTimes(outcome === 'success' ? 1 : 0); + if (outcome === 'success') { + expect(mockProcessMemory.mock.calls[0][2]).toEqual( + api.getPrivateTextInspectionTokens([message]), + ); + expect(mockProcessMemory.mock.calls[0][3]).toBe(client.abortController.signal); + } + }, + ); + it('should filter out image URLs from message content', async () => { const { HumanMessage, AIMessage } = require('@librechat/agents/langchain/messages'); const messages = [ diff --git a/api/server/utils/import/fork.js b/api/server/utils/import/fork.js index 381217b21d6..dc141a63199 100644 --- a/api/server/utils/import/fork.js +++ b/api/server/utils/import/fork.js @@ -1,9 +1,12 @@ const { v4: uuidv4 } = require('uuid'); const { cloneLineage, + saveCopiedMessage, + createNativeCopyPreflight, withoutTraceRefs, isTemporaryRecord, getAllMessagesUpToParent, + transferNativeCopyProvenance, } = require('@librechat/api'); const { logger, tenantStorage } = require('@librechat/data-schemas'); const { EModelEndpoint, Constants, ForkOptions } = require('librechat-data-provider'); @@ -24,7 +27,7 @@ const BaseClient = require('~/app/clients/BaseClient'); function cloneMessagesWithTimestamps( messagesToClone, importBatchBuilder, - { detachSubagentRuntime = false } = {}, + { detachSubagentRuntime = false, nativeCopy = false } = {}, ) { const { entries, idMapping } = cloneLineage(messagesToClone, uuidv4); for (const { source, messageId, parentMessageId, createdAt } of entries) { @@ -36,12 +39,13 @@ function cloneMessagesWithTimestamps( }; delete clonedMessage.privateText; delete clonedMessage.privacyRevision; + delete clonedMessage.privateTextTokens; if (detachSubagentRuntime) { delete clonedMessage.subagentTask; delete clonedMessage.subagentTranscript; } - importBatchBuilder.saveMessage(clonedMessage); + saveCopiedMessage(importBatchBuilder, source, clonedMessage, nativeCopy); } return idMapping; @@ -80,10 +84,13 @@ async function forkConversation({ }) { try { const originalConvo = await getConvo(requestUserId, originalConvoId); - let originalMessages = await getMessages({ - user: requestUserId, - conversationId: originalConvoId, - }); + let originalMessages = await getMessages( + { + user: requestUserId, + conversationId: originalConvoId, + }, + '+privateTextTokens', + ); let targetMessageId = targetId; if (splitAtTarget && !latestMessageId) { @@ -121,6 +128,7 @@ async function forkConversation({ * durable child executor. Preserve visible history while dropping the * task protocol and private serialized model transcript. */ detachSubagentRuntime: originalConvo.subagentThread != null, + nativeCopy: true, }); const result = importBatchBuilder.finishConversation( @@ -140,10 +148,13 @@ async function forkConversation({ } const conversation = await getConvo(requestUserId, result.conversation.conversationId); - const messages = await getMessages({ - user: requestUserId, - conversationId: conversation.conversationId, - }); + const messages = await getMessages( + { + user: requestUserId, + conversationId: conversation.conversationId, + }, + '-privateTextTokens', + ); return { conversation, @@ -373,7 +384,7 @@ async function forkSharedConversation({ // or share file URLs into the new conversation while file serving is off. const share = await getSharedMessages(shareId, shareResourceId, { snapshotFiles, - preflight: sharedContentPreflight, + preflight: createNativeCopyPreflight(sharedContentPreflight), }); if (!share?.messages?.length) { return null; @@ -419,14 +430,18 @@ async function forkSharedConversation({ } const messageIds = new Set(sourceMessages.map((message) => message.messageId)); - const messagesToClone = sourceMessages.map(({ model: _model, ...message }) => - stripSharedFileIds({ - ...message, - parentMessageId: - message.parentMessageId != null && messageIds.has(message.parentMessageId) - ? message.parentMessageId - : Constants.NO_PARENT, - }), + const messagesToClone = sourceMessages.map((source) => + transferNativeCopyProvenance( + source, + stripSharedFileIds({ + ...source, + model: undefined, + parentMessageId: + source.parentMessageId != null && messageIds.has(source.parentMessageId) + ? source.parentMessageId + : Constants.NO_PARENT, + }), + ), ); /** @@ -460,7 +475,7 @@ async function forkSharedConversation({ ); importBatchBuilder.startConversation(endpoint); - cloneMessagesWithTimestamps(messagesToClone, importBatchBuilder); + cloneMessagesWithTimestamps(messagesToClone, importBatchBuilder, { nativeCopy: true }); const result = importBatchBuilder.finishConversation(share.title, new Date(), {}, model); await importBatchBuilder.saveBatch(); @@ -471,10 +486,13 @@ async function forkSharedConversation({ }); const conversation = await getConvo(requestUserId, result.conversation.conversationId); - const messages = await getMessages({ - user: requestUserId, - conversationId: conversation.conversationId, - }); + const messages = await getMessages( + { + user: requestUserId, + conversationId: conversation.conversationId, + }, + '-privateTextTokens', + ); return { conversation, @@ -509,10 +527,13 @@ async function duplicateConversation({ throw new Error('Conversation not found'); } - const originalMessages = await getMessages({ - user: userId, - conversationId, - }); + const originalMessages = await getMessages( + { + user: userId, + conversationId, + }, + '+privateTextTokens', + ); const messagesToClone = getMessagesUpToTargetLevel( originalMessages, @@ -526,7 +547,7 @@ async function duplicateConversation({ importBatchBuilder.sourceIsTemporary = isTemporaryRecord(originalConvo); importBatchBuilder.startConversation(originalConvo.endpoint ?? EModelEndpoint.openAI); - cloneMessagesWithTimestamps(messagesToClone, importBatchBuilder); + cloneMessagesWithTimestamps(messagesToClone, importBatchBuilder, { nativeCopy: true }); const duplicateTitle = title || originalConvo.title; const result = importBatchBuilder.finishConversation(duplicateTitle, new Date(), originalConvo); @@ -539,10 +560,13 @@ async function duplicateConversation({ }); const conversation = await getConvo(userId, result.conversation.conversationId); - const messages = await getMessages({ - user: userId, - conversationId: conversation.conversationId, - }); + const messages = await getMessages( + { + user: userId, + conversationId: conversation.conversationId, + }, + '-privateTextTokens', + ); return { conversation, diff --git a/api/server/utils/import/fork.spec.js b/api/server/utils/import/fork.spec.js index c71602a71bf..8bf578fcfef 100644 --- a/api/server/utils/import/fork.spec.js +++ b/api/server/utils/import/fork.spec.js @@ -755,7 +755,7 @@ describe('forkSharedConversation', () => { expect(getSharedMessages).toHaveBeenCalledWith('share123', 'resource123', { snapshotFiles: undefined, - preflight: undefined, + preflight: expect.any(Function), }); const savedMessages = bulkSaveMessages.mock.calls[0][0]; @@ -891,7 +891,7 @@ describe('forkSharedConversation', () => { expect(getSharedMessages).toHaveBeenCalledWith('share123', 'resource123', { snapshotFiles: false, - preflight: undefined, + preflight: expect.any(Function), }); }); @@ -905,9 +905,12 @@ describe('forkSharedConversation', () => { sharedContentPreflight, }); + const options = getSharedMessages.mock.calls[0][2]; + await options.preflight(mockShare, { canonicalMessages: [] }); + expect(sharedContentPreflight).toHaveBeenCalledWith(mockShare, { canonicalMessages: [] }); expect(getSharedMessages).toHaveBeenCalledWith('share123', 'resource123', { snapshotFiles: undefined, - preflight: sharedContentPreflight, + preflight: expect.any(Function), }); }); @@ -1911,3 +1914,105 @@ describe('cloneMessagesWithTimestamps', () => { expect(clonedMessages.length).toBe(complexMessages.length); }); }); + +describe('native copy protected token admission', () => { + beforeEach(() => jest.clearAllMocks()); + const revision = 'a'.repeat(32); + const token = `[EMAIL_1_${revision}]`; + const filters = { + messages: { + pii: { + fields: ['text'], + starterPatterns: [], + customPatterns: [{ id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}' }], + }, + }, + }; + const source = { + messageId: 'protected-source', + conversationId: 'source', + parentMessageId: Constants.NO_PARENT, + text: token, + isCreatedByUser: true, + privacyRevision: revision, + }; + it.each(['fork', 'duplicate'])( + 'admits %s with exact trusted tokens and persists only canonical provenance', + async (kind) => { + getConvo.mockResolvedValue({ + conversationId: 'source', + endpoint: 'agents', + title: 'Safe copy', + }); + getMessages.mockResolvedValue([source]); + const result = + kind === 'fork' + ? await forkConversation({ + originalConvoId: 'source', + targetMessageId: source.messageId, + requestUserId: 'owner', + filters, + }) + : await duplicateConversation({ conversationId: 'source', userId: 'owner', filters }); + expect(result).toBeDefined(); + const call = bulkSaveMessages.mock.calls.at(-1); + expect(call[0][0]).not.toHaveProperty('privacyRevision'); + expect(call[0][0]).not.toHaveProperty('privateTextTokens'); + expect(call[2].privateTextTokens.get(call[0][0].messageId)).toEqual([token]); + }, + ); + it("retains a copied row's server-only token provenance across a second native copy", async () => { + getConvo.mockResolvedValue({ + conversationId: 'copy', + endpoint: 'agents', + title: 'Second copy', + }); + getMessages.mockResolvedValue([ + { ...source, privacyRevision: undefined, privateTextTokens: [token] }, + ]); + await duplicateConversation({ conversationId: 'copy', userId: 'owner', filters }); + const [rows, , metadata] = bulkSaveMessages.mock.calls.at(-1); + expect(metadata.privateTextTokens.get(rows[0].messageId)).toEqual([token]); + expect(rows[0]).not.toHaveProperty('privateTextTokens'); + }); + + it('uses share-authorized canonical provenance without putting it in the shared DTO', async () => { + const share = { + title: 'Shared', + messages: [{ ...source, privacyRevision: undefined }], + conversationId: 'anonymous', + shareId: 'shared', + }; + getSharedMessages.mockImplementationOnce(async (_id, _resource, options) => { + await options.preflight(share, { canonicalMessages: [source] }); + return share; + }); + getConvo.mockResolvedValue({ conversationId: 'copy', endpoint: 'agents' }); + getMessages.mockResolvedValue([]); + const loadAppConfig = jest.fn(async () => ({ filters })); + await forkSharedConversation({ shareId: 'shared', requestUserId: 'viewer', loadAppConfig }); + const [rows, , metadata] = bulkSaveMessages.mock.calls.at(-1); + expect(metadata.privateTextTokens.get(rows[0].messageId)).toEqual([token]); + expect(JSON.stringify(share)).not.toContain('privateTextTokens'); + expect(rows[0]).not.toHaveProperty('privacyRevision'); + }); + + it('does not grant token trust to untrusted use of the generic clone helper', async () => { + const builder = createImportBatchBuilder('owner', undefined, filters); + builder.startConversation(); + cloneMessagesWithTimestamps([source], builder); + builder.finishConversation('Safe title'); + await expect(builder.saveBatch()).rejects.toMatchObject({ code: 'content_filter_block' }); + }); + it('rejects raw credentials adjacent to native-copy placeholders', async () => { + getConvo.mockResolvedValue({ + conversationId: 'source', + endpoint: 'agents', + title: 'Safe copy', + }); + getMessages.mockResolvedValue([{ ...source, text: `${token} ${'b'.repeat(32)}` }]); + await expect( + duplicateConversation({ conversationId: 'source', userId: 'owner', filters }), + ).rejects.toMatchObject({ code: 'content_filter_block' }); + }); +}); diff --git a/api/server/utils/import/importBatchBuilder.js b/api/server/utils/import/importBatchBuilder.js index f7421fce57c..3ab878f2986 100644 --- a/api/server/utils/import/importBatchBuilder.js +++ b/api/server/utils/import/importBatchBuilder.js @@ -7,6 +7,8 @@ const { executeConversationImportWrites, resolveImportRetentionFields, resolveImportTagCounts, + getNativeCopyInspectionTokens, + saveNativeCopyMessages, } = require('@librechat/api'); const { getTenantId, @@ -202,6 +204,7 @@ class ImportBatchBuilder { { user: { id: this.requestUserId }, getFiles, + privateTextTokens: getNativeCopyInspectionTokens(this.messages), ...(this.legacyPii == null ? {} : { legacyPii: this.legacyPii }), }, ); @@ -220,7 +223,7 @@ class ImportBatchBuilder { try { await executeConversationImportWrites({ saveConversations: () => bulkSaveConvos(this.conversations), - saveMessages: () => bulkSaveMessages(this.messages, true), + saveMessages: () => saveNativeCopyMessages(bulkSaveMessages, this.messages), updateTagCounts: () => bulkIncrementTagCounts(this.requestUserId, tags), deleteMessages: () => deleteImportedMessages(cleanupScope), deleteConversations: () => deleteImportedConversations(cleanupScope), diff --git a/api/server/utils/import/importers.js b/api/server/utils/import/importers.js index ee211c67888..9a995257c82 100644 --- a/api/server/utils/import/importers.js +++ b/api/server/utils/import/importers.js @@ -70,6 +70,7 @@ function sanitizeImportedMessage(message) { delete importable.contextMeta; delete importable.privateText; delete importable.privacyRevision; + delete importable.privateTextTokens; return { ...importable, isUserSubmitted: true, diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index 8883c6dc8cb..66cc63e7bf9 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -306,3 +306,115 @@ test('an unchanged unprotected transcript still downloads a PNG screenshot', asy } } }); + +test('native copies retain canonical token trust without copying owner originals', async ({ + page, + request, +}) => { + test.setTimeout(120000); + const token = await loginAdmin(request); + const copies: string[] = []; + await setRuntimeFilters(request, token, { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}', category: 'credential' }, + ], + }, + }, + }); + try { + await page.goto('/c/new'); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const response = await sendMessageAndWaitForCompletion(page, original); + const source = (await response.json()).conversationId as string; + copies.push(source); + const sourceMessages = await fetchJson(page, `/api/messages/${source}`, token); + const sourceUser = sourceMessages.find((message) => message.isCreatedByUser)!; + const duplicate = await requestResult(request, { + path: '/api/convos/duplicate', + token, + method: 'POST', + data: { conversationId: source }, + }); + expect(duplicate.ok).toBe(true); + const copied = duplicate.body as { + conversation: { conversationId: string }; + messages: TMessage[]; + }; + const duplicateId = copied.conversation.conversationId; + copies.push(duplicateId); + const fork = await requestResult(request, { + path: '/api/convos/fork', + token, + method: 'POST', + data: { conversationId: source, messageId: sourceMessages.at(-1)!.messageId }, + }); + expect(fork.ok).toBe(true); + copies.push((fork.body as typeof copied).conversation.conversationId); + expect(JSON.stringify([duplicate.body, fork.body])).not.toContain('alice@example.com'); + expect(JSON.stringify([duplicate.body, fork.body])).not.toContain('privateText'); + expect(JSON.stringify([duplicate.body, fork.body])).not.toContain('privacyRevision'); + await withMongo(async (db) => { + const row = await db + .collection('messages') + .findOne({ conversationId: duplicateId, isCreatedByUser: true }); + expect(row?.text).toBe(sourceUser.text); + expect(row?.privateTextTokens).toEqual([ + sourceUser.text.match(/\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]/)![0], + ]); + expect(row?.privacyRevision).toBeUndefined(); + expect(row?.privateText).toBeUndefined(); + }); + await page.goto(`/c/${duplicateId}`); + await sendMessageAndWaitForCompletion(page, 'Safe copied-history follow-up'); + await expect( + messagesView(page).getByText('E2E private model input verified', { exact: true }).last(), + ).toBeVisible(); + const canonical = await fetchJson(page, `/api/messages/${duplicateId}`, token); + expect(JSON.stringify(canonical)).not.toContain('privateText'); + expect(JSON.stringify(canonical)).not.toContain('alice@example.com'); + const share = await requestResult(request, { + path: `/api/share/${duplicateId}`, + token, + method: 'POST', + data: {}, + }); + expect(share.ok).toBe(true); + const shareId = (share.body as { shareId: string }).shareId; + const published = await requestResult(request, { path: `/api/share/${shareId}`, token }); + expect(published.ok).toBe(true); + expect(published.text).not.toContain('privateText'); + const sharedFork = await requestResult(request, { + path: `/api/share/${shareId}/fork`, + token, + method: 'POST', + data: {}, + }); + expect(sharedFork.ok).toBe(true); + copies.push((sharedFork.body as typeof copied).conversation.conversationId); + expect(sharedFork.text).not.toContain('privateText'); + const repeated = await requestResult(request, { + path: '/api/convos/duplicate', + token, + method: 'POST', + data: { conversationId: duplicateId }, + }); + expect(repeated.ok).toBe(true); + copies.push((repeated.body as typeof copied).conversation.conversationId); + } finally { + await restoreRuntimeFilters(request, token); + for (const conversationId of copies) { + await requestResult(request, { + path: '/api/convos', + token, + method: 'DELETE', + data: { arg: { conversationId } }, + }); + } + } +}); diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 3f1285207af..70cfe0ba70f 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -1292,3 +1292,89 @@ describe('memory model-bound content preflight', () => { expect(Run.create).not.toHaveBeenCalled(); }); }); + +describe('automatic memory canonical token trust', () => { + const token = `[EMAIL_1_${'a'.repeat(32)}]`; + const filters: FiltersConfig = { + messages: { + pii: { + fields: ['text', 'content_part'], + starterPatterns: [], + customPatterns: [{ id: 'hex', label: 'Credential', regex: '[a-f0-9]{32}' }], + }, + }, + }; + it.each([true, false])('admits only verified generated tokens, trusted: %s', async (trusted) => { + await processMemory({ + res: {} as Response, + userId: 'owner', + messageId: 'response', + conversationId: 'conversation', + messages: [new HumanMessage(token)], + inspectionMessages: [new HumanMessage(token)], + privateTextTokens: trusted ? new Set([token]) : undefined, + memory: '', + instructions: 'Extract explicitly requested memories', + filters, + setMemory: jest.fn(), + deleteMemory: jest.fn(), + }); + expect(Run.create).toHaveBeenCalledTimes(trusted ? 1 : 0); + }); + it('still rejects raw credentials surrounding a trusted token', async () => { + await processMemory({ + res: {} as Response, + userId: 'owner', + messageId: 'response', + conversationId: 'conversation', + messages: [new HumanMessage(`${token} ${'b'.repeat(32)}`)], + privateTextTokens: new Set([token]), + memory: '', + instructions: 'Extract', + filters, + setMemory: jest.fn(), + deleteMemory: jest.fn(), + }); + expect(Run.create).not.toHaveBeenCalled(); + }); +}); + +it('does not invoke automatic extraction when its generation has already aborted', async () => { + const controller = new AbortController(); + controller.abort(); + await processMemory({ + res: {} as Response, + userId: 'owner', + messageId: 'response', + conversationId: 'conversation', + messages: [new HumanMessage('Safe')], + memory: '', + instructions: 'Extract', + signal: controller.signal, + setMemory: jest.fn(), + deleteMemory: jest.fn(), + }); + expect(Run.create).not.toHaveBeenCalled(); +}); + +it('forwards generation cancellation into the automatic memory SDK run', async () => { + const controller = new AbortController(); + const processStream = jest.fn(async () => 'success'); + (Run.create as jest.Mock).mockReturnValueOnce({ processStream }); + await processMemory({ + res: {} as Response, + userId: 'owner', + messageId: 'response', + conversationId: 'conversation', + messages: [new HumanMessage('Safe')], + memory: '', + instructions: 'Extract', + signal: controller.signal, + setMemory: jest.fn(), + deleteMemory: jest.fn(), + }); + expect(processStream).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ signal: controller.signal }), + ); +}); diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 97f62332a79..1e4db222636 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -761,6 +761,8 @@ export async function processMemory({ deleteMemory, messages, inspectionMessages, + privateTextTokens, + signal, memory, memoryEntries, messageId, @@ -788,6 +790,8 @@ export async function processMemory({ conversationId: string; messages: BaseMessage[]; inspectionMessages?: BaseMessage[]; + privateTextTokens?: ReadonlySet; + signal?: AbortSignal; validKeys?: string[]; instructions: string; /** Canonical rows preserve key/value granularity for field-scoped policy. */ @@ -807,6 +811,7 @@ export async function processMemory({ tenantId?: string; }): Promise<(TAttachment | null)[] | undefined> { try { + signal?.throwIfAborted(); const submittedMessages = (inspectionMessages ?? messages).filter( (message) => message._getType() !== 'ai', ); @@ -821,6 +826,7 @@ export async function processMemory({ } assertModelBoundContent({ filters, + privateTextTokens, submittedMessages, agents: [{ instructions, model_parameters: llmConfig }], memories, @@ -994,7 +1000,9 @@ ${memory ?? 'No existing memories'}`; returnContent: true, }); + signal?.throwIfAborted(); const config = { + ...(signal != null && { signal }), runName: 'MemoryRun', configurable: { user_id: userId, @@ -1066,6 +1074,8 @@ export async function createMemoryProcessor({ ( messages: BaseMessage[], inspectionMessages?: BaseMessage[], + privateTextTokens?: ReadonlySet, + signal?: AbortSignal, ) => Promise<(TAttachment | null)[] | undefined>, ] > { @@ -1101,6 +1111,8 @@ export async function createMemoryProcessor({ async function ( messages: BaseMessage[], inspectionMessages?: BaseMessage[], + privateTextTokens?: ReadonlySet, + signal?: AbortSignal, ): Promise<(TAttachment | null)[] | undefined> { try { return await processMemory({ @@ -1109,6 +1121,8 @@ export async function createMemoryProcessor({ agentId, messages, inspectionMessages, + privateTextTokens, + signal, validKeys, llmConfig, messageId, diff --git a/packages/api/src/protection/index.ts b/packages/api/src/protection/index.ts index bfc6d02532b..c962ed4f146 100644 --- a/packages/api/src/protection/index.ts +++ b/packages/api/src/protection/index.ts @@ -12,3 +12,4 @@ export * from './adapters/messages'; export * from './adapters/submissions'; export * from './private/submission'; export * from './private/view'; +export * from './private/copy'; diff --git a/packages/api/src/protection/private/copy.ts b/packages/api/src/protection/private/copy.ts new file mode 100644 index 00000000000..4ec88e8ea2a --- /dev/null +++ b/packages/api/src/protection/private/copy.ts @@ -0,0 +1,99 @@ +import type { MessageMethods } from '@librechat/data-schemas'; +import type { ShareContentPreflight } from '../../shared-links/protection'; +import { getPrivateTextInspectionTokens } from './submission'; + +type CopyMessage = { + readonly isCreatedByUser?: boolean; + readonly messageId?: string; + readonly text?: string; + readonly privacyRevision?: string; + readonly privateTextTokens?: readonly string[]; +}; +const copiedTokens = new WeakMap(); + +/** Only native callers supply server-loaded source rows. External imports use the default. */ +export function saveCopiedMessage( + builder: { saveMessage: (message: T) => CopyMessage }, + source: CopyMessage, + clone: T, + nativeCopy = false, +): void { + const message = builder.saveMessage(clone); + if (!nativeCopy) { + return; + } + const trusted = copiedTokens.get(source) ?? [...getPrivateTextInspectionTokens([source])]; + const tokens = trusted.filter((token) => message.text?.includes(token)); + if (tokens.length > 0) { + copiedTokens.set(message, tokens); + } +} + +export function getNativeCopyProvenance(messages: readonly CopyMessage[]): { + privateTextTokens: ReadonlyMap; +} { + const tokens = new Map(); + for (const message of messages) { + const trusted = copiedTokens.get(message); + if (message.messageId && trusted != null) { + tokens.set(message.messageId, trusted); + } + } + return { privateTextTokens: tokens }; +} + +export function getNativeCopyInspectionTokens( + messages: readonly CopyMessage[], +): ReadonlySet { + const tokens = new Set(); + for (const message of messages) { + for (const token of copiedTokens.get(message) ?? []) { + tokens.add(token); + } + } + return tokens; +} + +/** No extra metadata or persistence work for ordinary imports. */ +export function saveNativeCopyMessages( + save: MessageMethods['bulkSaveMessages'], + messages: Parameters[0], +): Promise { + const provenance = getNativeCopyProvenance(messages); + return provenance.privateTextTokens.size > 0 + ? save(messages, true, provenance) + : save(messages, true); +} + +export function transferNativeCopyProvenance(source: object, target: T): T { + const tokens = copiedTokens.get(source); + if (tokens != null) { + copiedTokens.set(target, tokens); + } + return target; +} + +/** Public snapshots gain only internal provenance from their authorized canonical read. */ +export function createNativeCopyPreflight( + preflight?: ShareContentPreflight, +): ShareContentPreflight { + return async (snapshot, context) => { + await preflight?.(snapshot, context); + const trusted = getPrivateTextInspectionTokens(context?.canonicalMessages ?? []); + for (const message of snapshot.messages) { + const tokens: string[] = []; + if (message.isCreatedByUser === true && typeof message.text === 'string') { + for (const match of message.text.matchAll( + /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]/g, + )) { + if (trusted.has(match[0])) { + tokens.push(match[0]); + } + } + } + if (tokens.length > 0) { + copiedTokens.set(message, tokens); + } + } + }; +} diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index 68474604485..f0b33b6ff7e 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -12,6 +12,8 @@ import { isPrivateTextChatSubmission, getPreinspectedPrivateText, getPrivateTextAdmission, + requirePrivateTextAdmission, + rejectPrivateTextAdmission, getPrivateTextInspectionTokens, privateTextBinding, } from './submission'; @@ -452,6 +454,43 @@ describe('private text submission boundary', () => { expect(getPrivateTextAdmission({}, undefined)).toBeUndefined(); }); + it.each(['failure', 'rejection', 'abort'])( + 'suppresses side effects on protected %s', + async (outcome) => { + const { req, message } = submit(); + const abortController = new AbortController(); + const sideEffect = jest.fn(); + const pending = requirePrivateTextAdmission(req, abortController.signal).then(sideEffect); + await Promise.resolve(); + expect(sideEffect).not.toHaveBeenCalled(); + if (outcome === 'abort') { + abortController.abort(); + } else { + rejectPrivateTextAdmission(req); + } + await expect(pending).rejects.toThrow(); + expect(sideEffect).not.toHaveBeenCalled(); + expect(message.text).not.toContain(original); + }, + ); + + it('holds optional side effects until the admitted atomic write finishes', async () => { + const { req, message } = submit(); + const sideEffect = jest.fn(); + const waiting = requirePrivateTextAdmission(req).then(sideEffect); + let finish!: (result: { message: typeof message }) => void; + const write = new Promise<{ message: typeof message }>((resolve) => { + finish = resolve; + }); + const admitted = getPrivateTextAdmission(req, () => write)!(); + await Promise.resolve(); + expect(sideEffect).not.toHaveBeenCalled(); + finish({ message }); + await admitted; + await waiting; + expect(sideEffect).toHaveBeenCalledTimes(1); + }); + it('rejects stale or swallowed persistence results, including a duplicate ID with different text', async () => { const { req, message } = submit(); const save: MessageMethods['saveMessage'] = jest.fn(async () => undefined); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 5a978170257..b221d65695e 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -13,6 +13,7 @@ interface PrivateTextMessage { isCreatedByUser?: boolean; text?: string; privacyRevision?: string; + privateTextTokens?: readonly string[]; } interface Capture { @@ -22,6 +23,8 @@ interface Capture { readonly text: string; readonly envelope: string; readonly cipher: PrivateTextCipher; + readonly admission: Promise; + readonly admit: (allowed: boolean) => void; } const captures = new WeakMap(); @@ -232,7 +235,11 @@ export function createPrivateTextIngress(options: { body.text.includes(marker) ? marker : `[${category}_${index}_${revision}]`, ); const envelope = cipher.seal(body.text, [userId, tenantId, revision]); - captures.set(req, { userId, tenantId, revision, text, envelope, cipher }); + let admit!: (allowed: boolean) => void; + const admission = new Promise((resolve) => { + admit = resolve; + }); + captures.set(req, { userId, tenantId, revision, text, envelope, cipher, admission, admit }); body.text = text; next(); } catch { @@ -278,16 +285,26 @@ export function getPrivateTextInspectionTokens( return tokens; } for (const message of messages) { - if ( - message?.isCreatedByUser !== true || - typeof message.text !== 'string' || - !/^[a-f0-9]{32}$/.test(message.privacyRevision ?? '') - ) { + if (message?.isCreatedByUser !== true || typeof message.text !== 'string') { continue; } if (message.text.length > 524288) { return new Set(); } + if ((message.privateTextTokens?.length ?? 0) > 4096) { + return new Set(); + } + for (const token of message.privateTextTokens ?? []) { + if ( + /^\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]$/.test(token) && + message.text.includes(token) + ) { + tokens.add(token); + if (tokens.size > 4096) { + return new Set(); + } + } + } for (const match of message.text.matchAll( /\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_([a-f0-9]{32})\]/g, )) { @@ -378,14 +395,52 @@ export async function requirePrivateTextPersistence( if (capture == null) { return; } - const result = await start(); - if ( - result?.message?.privacyRevision !== capture.revision || - result.message.text !== capture.text - ) { + try { + const result = await start(); + if ( + result?.message?.privacyRevision !== capture.revision || + result.message.text !== capture.text + ) { + throw unavailable(); + } + capture.admit(true); + onPersisted?.(); + } catch (error) { + capture.admit(false); + throw error; + } +} + +/** Optional model side effects share the turn's admission decision, without starting its write. */ +export async function requirePrivateTextAdmission( + req: object | undefined, + signal?: AbortSignal, +): Promise { + const capture = req == null ? undefined : captures.get(req); + if (capture == null) { + return; + } + if (signal?.aborted) { throw unavailable(); } - onPersisted?.(); + let abort!: () => void; + const cancelled = new Promise((resolve) => { + abort = () => resolve(false); + signal?.addEventListener('abort', abort, { once: true }); + }); + try { + if (!(await Promise.race([capture.admission, cancelled])) || signal?.aborted) { + throw unavailable(); + } + } finally { + signal?.removeEventListener('abort', abort); + } +} + +export function rejectPrivateTextAdmission(req: object | undefined): void { + if (req != null) { + captures.get(req)?.admit(false); + } } /** Ordinary startup stays immediate; protected revisions are announced only after admission. */ diff --git a/packages/api/src/shared-links/protection.ts b/packages/api/src/shared-links/protection.ts index a9c2334f44d..dafe96032ff 100644 --- a/packages/api/src/shared-links/protection.ts +++ b/packages/api/src/shared-links/protection.ts @@ -120,6 +120,7 @@ export type ShareContentPreflight = ( readonly text?: string; readonly isCreatedByUser?: boolean; readonly privacyRevision?: string; + readonly privateTextTokens?: readonly string[]; }[]; }, ) => Promise; diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index cfc0f2018ad..587eef47919 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -355,3 +355,78 @@ it.each([ ).toBe(false); }); }); + +it('persists native-copy token provenance only from metadata and clears it on canonical edits', async () => { + const token = `[EMAIL_1_${'a'.repeat(32)}]`; + const conversationId = uuid(); + const messageId = uuid(); + const message = { + messageId, + conversationId, + user: 'owner', + text: token, + isCreatedByUser: true, + privateTextTokens: [token], + privacyRevision: 'forged', + privateText: 'forged', + }; + await tenant('tenant-a', async () => { + await methods.bulkSaveMessages([message]); + expect( + (await methods.getMessages({ messageId, user: 'owner' }, '+privateTextTokens'))[0], + ).not.toHaveProperty('privateTextTokens'); + await methods.bulkSaveMessages([message], true, { + privateTextTokens: new Map([[messageId, [token]]]), + }); + const internal = ( + await methods.getMessages({ messageId, user: 'owner' }, '+privateTextTokens') + )[0]; + expect(internal.privateTextTokens).toEqual([token]); + expect(internal).not.toHaveProperty('privacyRevision'); + expect(internal).not.toHaveProperty('privateText'); + for (const select of [undefined, CLIENT_MESSAGE_SELECT]) { + expect( + (await methods.getMessages({ messageId, user: 'owner' }, select))[0], + ).not.toHaveProperty('privateTextTokens'); + } + await methods.updateMessageText('owner', { messageId, text: 'Clean edit' }); + expect( + (await methods.getMessages({ messageId, user: 'owner' }, '+privateTextTokens'))[0], + ).not.toHaveProperty('privateTextTokens'); + }); +}); + +it.each(['saveMessage', 'recordMessage', 'updateMessage', 'bulkSaveMessages'] as const)( + 'clears native-copy provenance when %s overwrites canonical text', + async (writer) => { + const token = `[EMAIL_1_${'a'.repeat(32)}]`; + const messageId = uuid(); + const conversationId = uuid(); + await tenant('tenant-a', async () => { + await methods.bulkSaveMessages( + [{ user: 'owner', conversationId, messageId, text: token, isCreatedByUser: true }], + true, + { privateTextTokens: new Map([[messageId, [token]]]) }, + ); + const edited = { + conversationId, + messageId, + text: 'Edited text', + isCreatedByUser: true, + privateTextTokens: [token], + }; + if (writer === 'saveMessage') { + await methods.saveMessage({ userId: 'owner' }, edited); + } else if (writer === 'recordMessage') { + await methods.recordMessage({ user: 'owner', ...edited }); + } else if (writer === 'updateMessage') { + await methods.updateMessage('owner', edited); + } else { + await methods.bulkSaveMessages([{ user: 'owner', ...edited }]); + } + expect( + (await methods.getMessages({ user: 'owner', messageId }, '+privateTextTokens'))[0], + ).not.toHaveProperty('privateTextTokens'); + }); + }, +); diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index eaad2eaacd1..5d0ec86234b 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -310,7 +310,11 @@ function buildMessageSaveUpdate( ...((options.unsetContextMeta || options.unsetPrivateText) && { $unset: { ...(options.unsetContextMeta && { contextMeta: 1 }), - ...(options.unsetPrivateText && { privateText: 1, privacyRevision: 1 }), + ...(options.unsetPrivateText && { + privateText: 1, + privacyRevision: 1, + ...(update.privateTextTokens === undefined && { privateTextTokens: 1 }), + }), }, }), }; @@ -403,7 +407,11 @@ async function findOneAndMergeMessageProvenance( ...((options.unsetContextMeta || options.unsetPrivateText) && { $unset: { ...(options.unsetContextMeta && { contextMeta: 1 }), - ...(options.unsetPrivateText && { privateText: 1, privacyRevision: 1 }), + ...(options.unsetPrivateText && { + privateText: 1, + privacyRevision: 1, + ...(update.privateTextTokens === undefined && { privateTextTokens: 1 }), + }), }, }), }, @@ -513,6 +521,7 @@ export const CLIENT_MESSAGE_SELECT: string = [ '-summary', '-summaryTokenCount', '-privateText', + '-privateTextTokens', '-contextMeta', '-langfuseSampled', '-langfuseDestinationIds', @@ -803,6 +812,7 @@ export interface MessageMethods { bulkSaveMessages( messages: Array>, overrideTimestamp?: boolean, + provenance?: { privateTextTokens: ReadonlyMap }, ): Promise; recordMessage(params: { user: string; @@ -1140,6 +1150,7 @@ export function createMessageMethods( }; delete update.privateText; delete update.privacyRevision; + delete update.privateTextTokens; if (metadata?.privateText != null) { if (params.isCreatedByUser !== true || typeof params.text !== 'string') { throw new Error('Private text requires a user message.'); @@ -1351,6 +1362,7 @@ export function createMessageMethods( async function bulkSaveMessages( messages: Array>, overrideTimestamp = false, + provenance?: { privateTextTokens: ReadonlyMap }, ) { try { const Message = mongoose.models.Message as Model; @@ -1358,22 +1370,37 @@ export function createMessageMethods( const normalizedMessage = sanitizeMessageUpdate(message); delete normalizedMessage.privateText; delete normalizedMessage.privacyRevision; - const provenance = capNormalizedProvenance( + delete normalizedMessage.privateTextTokens; + const tokens = provenance?.privateTextTokens.get(String(message.messageId)); + const text = message.text; + if ( + tokens?.length && + tokens.length <= 4096 && + message.isCreatedByUser === true && + typeof text === 'string' + ) { + normalizedMessage.privateTextTokens = tokens.filter( + (token) => + /^\[(?:EMAIL|PHONE|NAME|CREDENTIAL|CUSTOM)_\d+_[a-f0-9]{32}\]$/.test(token) && + text.includes(token), + ); + } + const submittedProvenance = capNormalizedProvenance( normalizeUserSubmittedPaths(message.userSubmittedPaths), normalizeUserSubmittedMessageFieldPaths(message.userSubmittedMessageFieldPaths), ); - if (provenance.userSubmittedPaths.length > 0) { - normalizedMessage.userSubmittedPaths = provenance.userSubmittedPaths; + if (submittedProvenance.userSubmittedPaths.length > 0) { + normalizedMessage.userSubmittedPaths = submittedProvenance.userSubmittedPaths; } else { delete normalizedMessage.userSubmittedPaths; } - if (provenance.userSubmittedMessageFieldPaths.length > 0) { + if (submittedProvenance.userSubmittedMessageFieldPaths.length > 0) { normalizedMessage.userSubmittedMessageFieldPaths = - provenance.userSubmittedMessageFieldPaths; + submittedProvenance.userSubmittedMessageFieldPaths; } else { delete normalizedMessage.userSubmittedMessageFieldPaths; } - if (provenance.promoteWholeMessage) { + if (submittedProvenance.promoteWholeMessage) { normalizedMessage.isUserSubmitted = true; } return { @@ -1385,7 +1412,11 @@ export function createMessageMethods( update: { $set: normalizedMessage, $inc: { __v: 1 }, - $unset: { privateText: 1, privacyRevision: 1 }, + $unset: { + privateText: 1, + privacyRevision: 1, + ...(normalizedMessage.privateTextTokens == null && { privateTextTokens: 1 }), + }, }, timestamps: !overrideTimestamp, upsert: true, @@ -1446,6 +1477,7 @@ export function createMessageMethods( for (const op of guarded) { await writeMessage(op.updateOne.filter, op.updateOne.update.$set, { upsert: true, + unsetPrivateText: true, timestamps: !overrideTimestamp, onWrite: (inserted, id) => { if (inserted) { @@ -1508,6 +1540,7 @@ export function createMessageMethods( } = rest; delete safeRest.privateText; delete safeRest.privacyRevision; + delete safeRest.privateTextTokens; const message = { user, endpoint, @@ -2461,6 +2494,7 @@ export function createMessageMethods( const { messageId, ...update } = message; delete update.privateText; delete update.privacyRevision; + delete update.privateTextTokens; const updatedMessage = await writeMessage({ messageId, user: userId }, update, { upsert: false, unsetPrivateText: Object.prototype.hasOwnProperty.call(update, 'text'), diff --git a/packages/data-schemas/src/methods/share.test.ts b/packages/data-schemas/src/methods/share.test.ts index 868b276f0a2..a004e4695b3 100644 --- a/packages/data-schemas/src/methods/share.test.ts +++ b/packages/data-schemas/src/methods/share.test.ts @@ -77,6 +77,7 @@ describe('Share Methods', () => { conversationId: { type: String, required: true }, user: { type: String, required: true }, text: String, + privateTextTokens: { type: [String], select: false, default: undefined }, isCreatedByUser: Boolean, isUserSubmitted: Boolean, userSubmittedPaths: { diff --git a/packages/data-schemas/src/methods/share.ts b/packages/data-schemas/src/methods/share.ts index ff7a27aee34..bd8f7e0ddc6 100644 --- a/packages/data-schemas/src/methods/share.ts +++ b/packages/data-schemas/src/methods/share.ts @@ -82,6 +82,7 @@ export type SharedMessagesPreflight = ( readonly text?: string; readonly isCreatedByUser?: boolean; readonly privacyRevision?: string; + readonly privateTextTokens?: readonly string[]; }[]; }, ) => void | Promise; @@ -965,7 +966,7 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { const share = (await query .populate({ path: 'messages', - select: CLIENT_MESSAGE_SELECT, + select: `${CLIENT_MESSAGE_SELECT.replace(' -privateTextTokens', '')} +privateTextTokens`, }) .select('-__v') .lean()) as (t.ISharedLink & { messages: t.IMessage[] }) | null; @@ -1282,7 +1283,10 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { }) .select('-_id -__v -user') .lean() as Promise, - Message.find({ conversationId, user }).sort({ createdAt: 1 }).lean(), + Message.find({ conversationId, user }) + .select('+privateTextTokens') + .sort({ createdAt: 1 }) + .lean(), ]); if (existingShare) { @@ -1457,6 +1461,7 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { } const updatedMessages = await Message.find({ conversationId: share.conversationId, user }) + .select('+privateTextTokens') .sort({ createdAt: 1 }) .lean(); diff --git a/packages/data-schemas/src/schema/message.ts b/packages/data-schemas/src/schema/message.ts index 450e5d6961e..6f0552fb5a5 100644 --- a/packages/data-schemas/src/schema/message.ts +++ b/packages/data-schemas/src/schema/message.ts @@ -61,6 +61,7 @@ const messageSchema: Schema = new Schema( }, privateText: { type: String, select: false }, privacyRevision: { type: String }, + privateTextTokens: { type: [String], select: false, default: undefined }, isCreatedByUser: { type: Boolean, required: true, diff --git a/packages/data-schemas/src/types/message.ts b/packages/data-schemas/src/types/message.ts index 517defdcf00..f6fcd3f8d39 100644 --- a/packages/data-schemas/src/types/message.ts +++ b/packages/data-schemas/src/types/message.ts @@ -63,6 +63,7 @@ export interface IMessage extends Document { /** Authenticated ciphertext, available only through the owner-view read. */ privateText?: string; privacyRevision?: string; + privateTextTokens?: string[]; isCreatedByUser: boolean; /** True when the complete stored row came from outside the model. */ isUserSubmitted?: boolean; From c7dee4f016c3b0eaf84c0ae753ff20e25af88260 Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 13:23:17 +0000 Subject: [PATCH 16/18] test: Wait for Private Message Concurrency Indexes --- packages/data-schemas/src/methods/message.privateText.spec.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/data-schemas/src/methods/message.privateText.spec.ts b/packages/data-schemas/src/methods/message.privateText.spec.ts index 587eef47919..816e1b8198f 100644 --- a/packages/data-schemas/src/methods/message.privateText.spec.ts +++ b/packages/data-schemas/src/methods/message.privateText.spec.ts @@ -20,6 +20,7 @@ beforeAll(async () => { server = await MongoMemoryServer.create(); Object.assign(mongoose.models, createModels(mongoose)); await mongoose.connect(server.getUri()); + await mongoose.models.Message.createIndexes(); }); afterAll(async () => { await mongoose.disconnect(); @@ -162,6 +163,7 @@ it('inserts a missing Stop prerequisite once without overwriting a stored or con }, ), ]); + expect(await mongoose.models.Message.countDocuments({ messageId: concurrentId })).toBe(1); const concurrent = await mongoose.models.Message.findOne({ messageId: concurrentId }) .select('+privateText') .lean(); From 02eee18971d47f98f8cc5fd750d28ab56e190fbb Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 13:32:57 +0000 Subject: [PATCH 17/18] fix: Preserve Protected Admission Across Recovery and Titles --- .../request.partialDisconnect.spec.js | 2 + .../__tests__/request.resumeMetadata.spec.js | 67 ++++++++++++++++++ api/server/controllers/agents/client.js | 1 + api/server/controllers/agents/client.test.js | 59 ++++++++++++++++ api/server/controllers/agents/request.js | 41 ++++++----- e2e/specs/mock/private-text.spec.ts | 70 +++++++++++++++++++ .../src/protection/private/submission.spec.ts | 64 +++++++++++++++++ .../api/src/protection/private/submission.ts | 50 +++++++++++-- 8 files changed, 331 insertions(+), 23 deletions(-) diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index aa21c26b499..18a88284973 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -48,6 +48,8 @@ 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), diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index bc9acac96a2..c8dcd02f819 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -257,6 +257,8 @@ 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, getSteerRecoveryFailure: jest.requireActual( '../../../../../packages/api/src/stream/SteerRecovery', @@ -3687,6 +3689,71 @@ describe('ResumableAgentController resume metadata', () => { } } + it.each(['history', 'exact-model'])( + 'does not persist a protected turn rejected by %s policy during terminal recovery', + async (boundary) => { + const api = jest.requireActual('@librechat/api'); + const req = createFailedRequest({ + text: 'alice@example.com', + clientRequestId: `private-error-${boundary}`, + }); + req.path = '/'; + const filters = { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }; + const next = jest.fn(); + api.createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, { status: jest.fn().mockReturnThis(), json: jest.fn() }, next); + expect(next).toHaveBeenCalledTimes(1); + const rejection = new api.ContentFilterError({ + detectorId: 'pii-pattern', + ruleId: 'secret', + label: 'secret', + source: 'message', + field: 'text', + provenance: 'user', + fragmentId: boundary, + fragmentPath: '/text', + }); + const client = { + options: {}, + sendMessage: jest.fn(async () => { + throw rejection; + }), + }; + await AgentController( + req, + createResumableResponse(), + jest.fn(), + jest.fn().mockResolvedValue({ client }), + null, + ); + await flushBackgroundGeneration(); + expect(client.sendMessage).toHaveBeenCalledTimes(1); + expect(mockGenerationJobManager.completeJob).toHaveBeenCalledWith( + conversationId, + rejection.message, + 1000, + expect.objectContaining({ beforeErrorPublication: expect.any(Function) }), + ); + expect(mockSaveMessage).not.toHaveBeenCalled(); + expect(mockSaveConvo).not.toHaveBeenCalled(); + expect(mockStampConvoLastResponse).not.toHaveBeenCalled(); + }, + ); + it('persists an initialization failure before terminal error publication', async () => { const events = []; mockSaveConvo.mockImplementation(async () => { diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 1996670c139..5b40833e20b 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -6142,6 +6142,7 @@ class AgentClient extends BaseClient { }); try { + await requirePrivateTextAdmission(req, abortController.signal); const titleResult = await this.run.generateTitle({ provider, clientOptions, diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index c0009e4f236..7a62a1c31bc 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -3956,6 +3956,65 @@ describe('AgentClient - titleConvo', () => { expect(mockRun.generateTitle).toHaveBeenCalled(); }); + it.each(['success', 'write-failure', 'policy-rejection', 'abort'])( + 'waits for protected persistence before an immediate title model call: %s', + async (outcome) => { + const api = jest.requireActual('@librechat/api'); + const filters = { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + ], + }, + }, + }; + mockReq.path = '/'; + mockReq.body.text = 'alice@example.com'; + mockReq.body.clientRequestId = 'protected-title'; + const next = jest.fn(); + api.createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(mockReq, { status: jest.fn().mockReturnThis(), json: jest.fn() }, next); + expect(next).toHaveBeenCalledTimes(1); + const message = api.stampPrivateTextMessage(mockReq, { + messageId: 'title-user', + conversationId: 'title-conversation', + text: mockReq.body.text, + isCreatedByUser: true, + }); + const abortController = new AbortController(); + const title = client.titleConvo({ text: message.text, abortController, immediate: true }); + await new Promise((resolve) => setImmediate(resolve)); + expect(mockRun.generateTitle).not.toHaveBeenCalled(); + expect(client.recordCollectedUsage).not.toHaveBeenCalled(); + if (outcome === 'abort') { + abortController.abort(); + } else if (outcome === 'policy-rejection') { + api.rejectPrivateTextAdmission(mockReq); + } else if (outcome === 'write-failure') { + await expect( + api.requirePrivateTextPersistence(mockReq, async () => ({})), + ).rejects.toThrow(); + } else { + const write = deferred(); + const admission = api.requirePrivateTextPersistence(mockReq, () => write.promise); + await Promise.resolve(); + expect(mockRun.generateTitle).not.toHaveBeenCalled(); + write.resolve({ message }); + await admission; + } + await title; + expect(mockRun.generateTitle).toHaveBeenCalledTimes(outcome === 'success' ? 1 : 0); + expect(client.recordCollectedUsage).toHaveBeenCalledTimes(outcome === 'success' ? 1 : 0); + }, + ); + it('passes empty contentParts in immediate mode (title from the user input only)', async () => { client.contentParts = [{ type: 'text', text: 'Streaming response so far' }]; const abortController = new AbortController(); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index f876e56470c..f378f1ad83d 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -59,6 +59,7 @@ const { getFailedTurnTraceFields, resolveFailedTurnContent, savePrivateTextMessage, + savePrivateTextErrorTurn, stampPreliminaryPrivateTextMessage, announceReply, announceErrorTurn, @@ -3369,17 +3370,19 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit ownsScheduledFailure = (await GenerationJobManager.completeJob(streamId, generationError, jobCreatedAt, { beforeErrorPublication: () => - saveErrorTurn(req, { - conversationId, - endpointOption, - isNewConvo, - errorText: generationError, - liveUserMessage: userMessage, - liveResponseMessageId, - runCreated: client?.run != null, - sender: client?.sender, - initialAgentId: verifiedInitialAgentId, - }), + savePrivateTextErrorTurn(req, error, () => + saveErrorTurn(req, { + conversationId, + endpointOption, + isNewConvo, + errorText: generationError, + liveUserMessage: userMessage, + liveResponseMessageId, + runCreated: client?.run != null, + sender: client?.sender, + initialAgentId: verifiedInitialAgentId, + }), + ), })) === true; /** A true completion means this owner won the terminal CAS and * the beforeErrorPublication barrier above finished. Only that @@ -3581,13 +3584,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit const completionPromise = persistInitializationError ? GenerationJobManager.completeJob(streamId, initializationError, jobCreatedAt, { beforeErrorPublication: () => - saveErrorTurn(req, { - conversationId, - endpointOption, - isNewConvo, - errorText: initializationError, - initialAgentId: verifiedInitialAgentId, - }), + savePrivateTextErrorTurn(req, error, () => + saveErrorTurn(req, { + conversationId, + endpointOption, + isNewConvo, + errorText: initializationError, + initialAgentId: verifiedInitialAgentId, + }), + ), }) : GenerationJobManager.completeJob(streamId, initializationError, jobCreatedAt); initializationFinalized = diff --git a/e2e/specs/mock/private-text.spec.ts b/e2e/specs/mock/private-text.spec.ts index 66cc63e7bf9..93ebe493821 100644 --- a/e2e/specs/mock/private-text.spec.ts +++ b/e2e/specs/mock/private-text.spec.ts @@ -418,3 +418,73 @@ test('native copies retain canonical token trust without copying owner originals } } }); + +test('a later history-policy rejection never persists the transformed protected turn', async ({ + page, + request, +}) => { + test.setTimeout(90000); + const token = await loginAdmin(request); + let conversationId: string | undefined; + try { + await page.goto('/c/new'); + await selectMockEndpoint(page, MOCK_ENDPOINTS[0]); + const seeded = await sendMessageAndWaitForCompletion(page, 'Safe admission-history seed'); + conversationId = (await seeded.json()).conversationId as string; + const prior = await fetchJson(page, `/api/messages/${conversationId}`, token); + const user = prior.find((message) => message.isCreatedByUser)!; + await withMongo(async (db) => { + await db.collection('messages').updateOne( + { conversationId, messageId: user.messageId }, + { + $set: { text: 'FORBIDDEN_HISTORICAL_MARKER' }, + }, + ); + }); + await page.reload(); + await setRuntimeFilters(request, token, { + messages: { + pii: { + action: 'redact', + fields: ['text'], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' }, + { id: 'history', label: 'Restricted history', regex: 'FORBIDDEN_HISTORICAL_MARKER' }, + ], + }, + }, + }); + const completion = page.waitForResponse(async (response) => { + if ( + response.request().method() !== 'GET' || + !new URL(response.url()).pathname.includes('/api/agents/chat/stream/') + ) { + return false; + } + return true; + }); + const admitted = await sendMessage(page, 'E2E_REJECTED_PRIVATE_TEXT: alice@example.com'); + expect(admitted.ok()).toBe(true); + const stream = await completion; + const eventText = await stream.text(); + expect(eventText).toContain('error'); + await withMongo(async (db) => { + const rows = await db.collection('messages').find({ conversationId }).toArray(); + expect(rows).toHaveLength(prior.length); + expect(JSON.stringify(rows)).not.toContain('E2E_REJECTED_PRIVATE_TEXT'); + expect(JSON.stringify(rows)).not.toContain('alice@example.com'); + expect(rows.every((row) => row.privateText == null)).toBe(true); + }); + } finally { + await restoreRuntimeFilters(request, token); + if (conversationId) { + await requestResult(request, { + path: '/api/convos', + token, + method: 'DELETE', + data: { arg: { conversationId } }, + }); + } + } +}); diff --git a/packages/api/src/protection/private/submission.spec.ts b/packages/api/src/protection/private/submission.spec.ts index f0b33b6ff7e..808c6340c95 100644 --- a/packages/api/src/protection/private/submission.spec.ts +++ b/packages/api/src/protection/private/submission.spec.ts @@ -14,6 +14,7 @@ import { getPrivateTextAdmission, requirePrivateTextAdmission, rejectPrivateTextAdmission, + savePrivateTextErrorTurn, getPrivateTextInspectionTokens, privateTextBinding, } from './submission'; @@ -22,6 +23,7 @@ import { assertModelBoundContent, } from '../../middleware/modelBoundContent'; import { createMessageFilterPii } from '../../middleware/messageFilterPii'; +import { ContentFilterError } from '../../middleware/contentFilter'; import { createPrivateTextCipher } from './crypto'; import { createPrivateTextView } from './view'; @@ -491,6 +493,68 @@ describe('private text submission boundary', () => { expect(sideEffect).toHaveBeenCalledTimes(1); }); + function policyError() { + return new ContentFilterError({ + detectorId: 'pii-pattern', + ruleId: 'secret', + label: 'secret', + source: 'message', + field: 'text', + provenance: 'user', + fragmentId: 'history', + fragmentPath: '/text', + }); + } + + it('rejects terminal recovery before any write when protected content admission failed', async () => { + const { req, message } = submit(); + const saveErrorTurn = jest.fn(async () => {}); + await savePrivateTextErrorTurn(req, policyError(), saveErrorTurn); + expect(saveErrorTurn).not.toHaveBeenCalled(); + const save = jest.fn(async () => message as IMessage); + await expect(savePrivateTextMessage(save, req, { userId: 'owner' }, message)).rejects.toThrow(); + expect(save).not.toHaveBeenCalled(); + const start = jest.fn(async () => ({ message })); + await expect(requirePrivateTextPersistence(req, start)).rejects.toThrow(); + expect(start).not.toHaveBeenCalled(); + }); + + it('cannot revive failed persistence through terminal recovery or a later admission', async () => { + const { req, message } = submit(); + await expect(requirePrivateTextPersistence(req, async () => ({}))).rejects.toThrow(); + const recovery = jest.fn(async () => {}); + await savePrivateTextErrorTurn(req, new Error('persistence failure'), recovery); + expect(recovery).not.toHaveBeenCalled(); + await expect(requirePrivateTextPersistence(req, async () => ({ message }))).rejects.toThrow(); + }); + + it('does not announce a write that resolves after an admission rejection', async () => { + const { req, message } = submit(); + let finish!: (result: { message: typeof message }) => void; + const write = new Promise<{ message: typeof message }>((resolve) => { + finish = resolve; + }); + const announce = jest.fn(); + const admission = requirePrivateTextPersistence(req, () => write, announce); + rejectPrivateTextAdmission(req); + finish({ message }); + await expect(admission).rejects.toThrow(); + expect(announce).not.toHaveBeenCalled(); + }); + + it('retains ordinary error recovery and already-admitted protected turns', async () => { + const ordinary = jest.fn(async () => {}); + await savePrivateTextErrorTurn({}, policyError(), ordinary); + expect(ordinary).toHaveBeenCalledTimes(1); + const { req, message } = submit(); + await requirePrivateTextPersistence(req, async () => ({ message })); + rejectPrivateTextAdmission(req); // Terminal cleanup cannot undo a committed admission. + const recovery = jest.fn(async () => {}); + await savePrivateTextErrorTurn(req, policyError(), recovery); + expect(recovery).toHaveBeenCalledTimes(1); + await expect(requirePrivateTextAdmission(req)).resolves.toBeUndefined(); + }); + it('rejects stale or swallowed persistence results, including a duplicate ID with different text', async () => { const { req, message } = submit(); const save: MessageMethods['saveMessage'] = jest.fn(async () => undefined); diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index b221d65695e..1d66ddbbe2d 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -2,7 +2,7 @@ import type { FiltersConfig, MessageFilterPiiConfig } from 'librechat-data-provi import type { RequestHandler, Request, Response } from 'express'; import type { MessageMethods } from '@librechat/data-schemas'; import type { PrivateTextCipher } from './crypto'; -import { ContentFilterError } from '../../middleware/contentFilter'; +import { ContentFilterError, isContentFilterError } from '../../middleware/contentFilter'; import { createPiiTextTransformer } from '../transform'; import { createPrivateTextCipher } from './crypto'; import { inspectContent } from '../runtime'; @@ -25,6 +25,7 @@ interface Capture { readonly cipher: PrivateTextCipher; readonly admission: Promise; readonly admit: (allowed: boolean) => void; + readonly admissionState: () => boolean | undefined; } const captures = new WeakMap(); @@ -235,11 +236,28 @@ export function createPrivateTextIngress(options: { body.text.includes(marker) ? marker : `[${category}_${index}_${revision}]`, ); const envelope = cipher.seal(body.text, [userId, tenantId, revision]); - let admit!: (allowed: boolean) => void; + let resolveAdmission!: (allowed: boolean) => void; + let admissionResult: boolean | undefined; const admission = new Promise((resolve) => { - admit = resolve; + resolveAdmission = resolve; + }); + const admit = (allowed: boolean) => { + if (admissionResult === undefined) { + admissionResult = allowed; + resolveAdmission(allowed); + } + }; + captures.set(req, { + userId, + tenantId, + revision, + text, + envelope, + cipher, + admission, + admit, + admissionState: () => admissionResult, }); - captures.set(req, { userId, tenantId, revision, text, envelope, cipher, admission, admit }); body.text = text; next(); } catch { @@ -345,7 +363,7 @@ export async function savePrivateTextMessage( if (capture == null || message.isCreatedByUser !== true) { return save(...args); } - if (message.text !== capture.text) { + if (capture.admissionState() === false || message.text !== capture.text) { throw unavailable(); } if ( @@ -396,8 +414,12 @@ export async function requirePrivateTextPersistence( return; } try { + if (capture.admissionState() === false) { + throw unavailable(); + } const result = await start(); if ( + capture.admissionState() === false || result?.message?.privacyRevision !== capture.revision || result.message.text !== capture.text ) { @@ -443,6 +465,24 @@ export function rejectPrivateTextAdmission(req: object | undefined): void { } } +/** Recovery may retain ordinary failures, but cannot revive an unadmitted protected turn. */ +export async function savePrivateTextErrorTurn( + req: object | undefined, + error: unknown, + save: () => Promise, +): Promise { + const capture = req == null ? undefined : captures.get(req); + if (capture != null) { + if (capture.admissionState() !== true && isContentFilterError(error)) { + capture.admit(false); + } + if (capture.admissionState() === false) { + return; + } + } + await save(); +} + /** Ordinary startup stays immediate; protected revisions are announced only after admission. */ export function deferPrivateTextStart( req: object | undefined, From 7b9013e45a949e3aa1d44485c79d024d0539baed Mon Sep 17 00:00:00 2001 From: Lia Date: Fri, 2 Oct 2026 13:54:47 +0000 Subject: [PATCH 18/18] fix: Coordinate Protected Root Admission and Early Stop --- api/app/clients/BaseClient.js | 18 +- api/app/clients/specs/BaseClient.test.js | 29 ++ api/server/controllers/agents/client.js | 41 +- .../api/src/middleware/modelBoundContent.ts | 18 +- packages/api/src/protection/index.ts | 1 + .../src/protection/private/admission.spec.ts | 385 ++++++++++++++++++ .../api/src/protection/private/admission.ts | 188 +++++++++ .../api/src/protection/private/submission.ts | 24 ++ 8 files changed, 676 insertions(+), 28 deletions(-) create mode 100644 packages/api/src/protection/private/admission.spec.ts create mode 100644 packages/api/src/protection/private/admission.ts diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index bb5b5841177..f02af652d4b 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -37,6 +37,7 @@ const { deferPrivateTextStart, requirePrivateTextPersistence, rejectPrivateTextAdmission, + bindPrivateTextPersistenceAbort, persistedReasoningOverrideFields, } = require('@librechat/api'); const { @@ -996,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 { diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index 4b9a618ef71..e6578ed3af7 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -1834,6 +1834,35 @@ describe('BaseClient', () => { 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) => { diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 5b40833e20b..6b8818d5932 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -159,7 +159,9 @@ const { filterFilesByEndpointRuntimeConfig, createModelBoundChatModelCallback: createModelBoundContentCallback, getPrivateTextInspectionTokens, - getPrivateTextAdmission, + getPrivateTextModelHooks, + createPrivateTextInitialAdmissionCallback, + withPrivateTextAdmissionConfig, requirePrivateTextAdmission, rejectPrivateTextAdmission, createInitialModelBoundAdmissionCallback, @@ -2026,7 +2028,7 @@ class AgentClient extends BaseClient { * at every chat-model call instead. */ assertBuiltModelBoundContent() {} - createModelBoundChatModelCallback() { + createModelBoundChatModelCallback(initialAdmission) { const fileProjection = BaseClient.prototype.getModelBoundFileProjection.call(this); const persistence = BaseClient.prototype.getModelBoundUserMessagePersistence.call(this); return createModelBoundContentCallback( @@ -2040,19 +2042,28 @@ class AgentClient extends BaseClient { resolvedFiles: fileProjection.resolvedFiles, sourceFileProjectionOverflowed: fileProjection.overflowed, }, - { - onContentRejected: persistence?.cancel, - onContentAllowed: getPrivateTextAdmission( - this.options.req, - persistence?.start, - this.privateTextStart, - ), - }, + getPrivateTextModelHooks( + this.options.req, + initialAdmission, + persistence?.start, + persistence?.cancel, + this.privateTextStart, + ), ); } createInitialModelBoundAdmissionCallback(startingAgentIds) { const persistence = BaseClient.prototype.getModelBoundUserMessagePersistence.call(this); + const protectedAdmission = createPrivateTextInitialAdmissionCallback(this.options.req, { + agentIds: startingAgentIds, + start: persistence?.start, + cancel: persistence?.cancel, + onPersisted: this.privateTextStart, + signal: this.abortController?.signal, + }); + if (protectedAdmission != null) { + return protectedAdmission; + } if (persistence == null || !persistence.isPending() || startingAgentIds.length === 0) { return undefined; } @@ -4826,16 +4837,16 @@ class AgentClient extends BaseClient { if (this.agentConfigs && this.agentConfigs.size > 0) { agents.push(...this.agentConfigs.values()); } - const modelBoundCallback = - AgentClient.prototype.createModelBoundChatModelCallback.call(this); const initialModelBoundAdmission = AgentClient.prototype.createInitialModelBoundAdmissionCallback.call( this, AgentClient.getStartingAgentIds(agents), ); - if (initialModelBoundAdmission != null) { - config.callbacks = [initialModelBoundAdmission]; - } + const modelBoundCallback = AgentClient.prototype.createModelBoundChatModelCallback.call( + this, + initialModelBoundAdmission, + ); + config = withPrivateTextAdmissionConfig(config, initialModelBoundAdmission); // TODO: needs to be added as part of AgentContext initialization // const noSystemModelRegex = [/\b(o1-preview|o1-mini|amazon\.titan-text)\b/gi]; diff --git a/packages/api/src/middleware/modelBoundContent.ts b/packages/api/src/middleware/modelBoundContent.ts index 53f8a9dcc9d..b89d5101661 100644 --- a/packages/api/src/middleware/modelBoundContent.ts +++ b/packages/api/src/middleware/modelBoundContent.ts @@ -681,6 +681,11 @@ export interface ModelBoundChatModelCallback { readonly handleChatModelStart: ( llm: object | undefined, messageBatches: readonly (readonly ModelBoundProviderMessage[])[], + runId?: string, + parentRunId?: string, + extraParams?: Record, + tags?: string[], + metadata?: Record, ) => void | Promise; } @@ -3240,7 +3245,11 @@ export function createModelBoundChatModelCallback( input: Omit, options: { readonly onContentRejected?: (error: unknown) => void; - readonly onContentAllowed?: () => void | Promise; + readonly onContentAllowed?: (context: { + runId?: string; + parentRunId?: string; + metadata?: Record; + }) => void | Promise; } = {}, ): ModelBoundChatModelCallback { const storedMessageSnapshot = snapshotBoundedProviderArray(input.storedMessages); @@ -3269,6 +3278,11 @@ export function createModelBoundChatModelCallback( handleChatModelStart: ( _llm: object | undefined, messageBatches: readonly (readonly ModelBoundProviderMessage[])[], + runId?: string, + parentRunId?: string, + _extraParams?: Record, + _tags?: string[], + metadata?: Record, ) => { let messageBatchCount = 0; let messageBatchesOverflowed = false; @@ -3326,7 +3340,7 @@ export function createModelBoundChatModelCallback( options.onContentRejected?.(error); throw new FatalModelBoundPolicyError(error); } - return options.onContentAllowed?.(); + return options.onContentAllowed?.({ runId, parentRunId, metadata }); }, }); return callback; diff --git a/packages/api/src/protection/index.ts b/packages/api/src/protection/index.ts index c962ed4f146..18efce9a68d 100644 --- a/packages/api/src/protection/index.ts +++ b/packages/api/src/protection/index.ts @@ -13,3 +13,4 @@ export * from './adapters/submissions'; export * from './private/submission'; export * from './private/view'; export * from './private/copy'; +export * from './private/admission'; diff --git a/packages/api/src/protection/private/admission.spec.ts b/packages/api/src/protection/private/admission.spec.ts new file mode 100644 index 00000000000..c90fbb72527 --- /dev/null +++ b/packages/api/src/protection/private/admission.spec.ts @@ -0,0 +1,385 @@ +import { Run, Providers, FakeChatModel } from '@librechat/agents'; +import { HumanMessage } from '@librechat/agents/langchain/messages'; +import { CallbackManager } from '@langchain/core/callbacks/manager'; +import { BaseCallbackHandler } from '@langchain/core/callbacks/base'; +import type { Request } from 'express'; +import { + createPrivateTextIngress, + stampPrivateTextMessage, + requirePrivateTextAdmission, + savePrivateTextErrorTurn, + bindPrivateTextPersistenceAbort, +} from './submission'; +import { + createPrivateTextInitialAdmissionCallback, + getPrivateTextModelHooks, + withPrivateTextAdmissionConfig, +} from './admission'; +import { createModelBoundChatModelCallback } from '../../middleware/modelBoundContent'; + +function turn() { + const filters = { + messages: { + pii: { + action: 'redact' as const, + fields: ['text' as const, 'content_part' as const], + starterPatterns: [], + customPatterns: [ + { id: 'email', label: 'Email', regex: 'alice@example\\.com', category: 'email' as const }, + { id: 'history', label: 'History', regex: 'RESTRICTED-HISTORY' }, + ], + }, + }, + }; + const req = { + path: '/', + user: { id: 'owner' }, + body: { + text: 'alice@example.com', + clientRequestId: 'parallel-private', + }, + } as unknown as Request; + createPrivateTextIngress({ + getFilters: () => filters, + getLegacyPii: () => undefined, + getKey: () => 'ab'.repeat(32), + })(req, {} as never, jest.fn()); + const message = stampPrivateTextMessage(req, { + messageId: 'user', + conversationId: 'conversation', + text: req.body.text, + isCreatedByUser: true, + }); + return { filters, req, message }; +} +function setup(roots = ['a', 'b']) { + const { req, message, filters } = turn(); + const controller = new AbortController(); + const start = jest.fn(async () => ({ message })); + const cancel = jest.fn(); + const created = jest.fn(); + const initial = createPrivateTextInitialAdmissionCallback(req, { + agentIds: roots, + start, + cancel, + onPersisted: created, + signal: controller.signal, + })!; + const content = createModelBoundChatModelCallback( + { filters }, + getPrivateTextModelHooks(req, initial, start, cancel, created), + ); + const manager = new CallbackManager(); + manager.addHandler( + Object.assign( + BaseCallbackHandler.fromMethods({ + handleChainStart: initial.handleChainStart, + handleChainEnd: initial.handleChainEnd, + handleChainError: initial.handleChainError, + handleLLMError: initial.handleLLMError, + }), + { + raiseError: true, + awaitHandlers: true, + }, + ), + true, + ); + const graph = manager.handleChainStart( + { lc: 1, type: 'not_implemented', id: ['graph'] }, + {}, + 'graph', + ); + const model = async (agent: string, text = 'Safe native input') => { + const graphRun = await graph; + const nodeManager = graphRun.getChild(); + const metadata = { agentId: agent, langgraph_node: `agent=${agent}` }; + nodeManager.addMetadata(metadata, true); + const node = await nodeManager.handleChainStart( + { lc: 1, type: 'not_implemented', id: ['agent'] }, + {}, + `node-${agent}`, + undefined, + undefined, + metadata, + `agent=${agent}`, + ); + const modelManager = node.getChild(); + modelManager.addHandler( + Object.assign( + BaseCallbackHandler.fromMethods({ + handleChatModelStart: (llm, batches, runId, parentRunId, extra, tags, metadata) => + content.handleChatModelStart( + llm, + batches.map((batch) => + batch.map((message) => ({ + role: message.getType(), + content: message.content, + })), + ), + runId, + parentRunId, + extra, + tags, + metadata, + ), + }), + { + raiseError: true, + awaitHandlers: true, + }, + ), + true, + ); + // Actual LangChain dispatch supplies ancestry and metadata to intrinsic callbacks. + return modelManager.handleChatModelStart( + { lc: 1, type: 'not_implemented', id: ['model'] }, + [[new HumanMessage(text)]], + `model-${agent}`, + undefined, + undefined, + undefined, + metadata, + ); + }; + return { req, message, start, cancel, created, initial, content, controller, model }; +} + +it('holds a safe root until a sibling input passes, then commits once before either model', async () => { + const state = setup(); + const invoke = jest.fn(); + const first = state.model('a').then(invoke); + await new Promise((resolve) => setImmediate(resolve)); + expect(state.start).not.toHaveBeenCalled(); + expect(invoke).not.toHaveBeenCalled(); + await state.model('b').then(invoke); + await first; + expect(state.start).toHaveBeenCalledTimes(1); + expect(state.created).toHaveBeenCalledTimes(1); + expect(invoke).toHaveBeenCalledTimes(2); + expect(state.created.mock.invocationCallOrder[0]).toBeLessThan( + invoke.mock.invocationCallOrder[0], + ); +}); + +it('a sibling rejection releases every waiting root without saving, invoking, or recovery', async () => { + const state = setup(); + const invoke = jest.fn(); + const first = state.model('a').then(invoke); + const observed = first.catch((error) => error); + await new Promise((resolve) => setImmediate(resolve)); + await expect(state.model('b', 'RESTRICTED-HISTORY')).rejects.toThrow(); + expect(await observed).toBeInstanceOf(Error); + expect(state.start).not.toHaveBeenCalled(); + expect(state.created).not.toHaveBeenCalled(); + expect(invoke).not.toHaveBeenCalled(); + const recover = jest.fn(async () => {}); + await savePrivateTextErrorTurn(state.req, new Error('rejected'), recover); + expect(recover).not.toHaveBeenCalled(); + await expect(requirePrivateTextAdmission(state.req)).rejects.toThrow(); +}); + +it('cancels a waiting parallel root on Stop, without beginning protected persistence', async () => { + const state = setup(); + const invoke = jest.fn(); + const first = state.model('a').then(invoke); + const observed = first.catch((error) => error); + await new Promise((resolve) => setImmediate(resolve)); + state.controller.abort(); + expect(await observed).toBeInstanceOf(Error); + expect(state.start).not.toHaveBeenCalled(); + expect(invoke).not.toHaveBeenCalled(); +}); + +it('does not accept a summary or unidentified model as initial-root admission', async () => { + const state = setup(); + await expect( + state.content.handleChatModelStart( + undefined, + [[{ role: 'user', content: 'Safe' }]], + 'summary', + 'graph', + undefined, + undefined, + { summarization: true }, + ), + ).rejects.toThrow(); + expect(state.start).not.toHaveBeenCalled(); +}); + +it('releases waiting roots when a peer needs a later summarization superstep', async () => { + const state = setup(); + const waiting = state.model('a').catch((error) => error); + await new Promise((resolve) => setImmediate(resolve)); + state.initial.handleChainStart( + undefined, + {}, + 'b-detour', + 'graph', + undefined, + { langgraph_node: 'agent=b' }, + undefined, + 'agent=b', + ); + state.initial.handleChainEnd({ summarizationRequest: {} }, 'b-detour'); + expect(await waiting).toBeInstanceOf(Error); + expect(state.start).not.toHaveBeenCalled(); +}); + +it('raises concurrency to the root-wave size without changing ordinary run configuration', () => { + const state = setup(['a', 'b', 'c']); + expect(withPrivateTextAdmissionConfig({ maxConcurrency: 1 }, state.initial).maxConcurrency).toBe( + 3, + ); + expect(withPrivateTextAdmissionConfig({}, state.initial).maxConcurrency).toBeUndefined(); + expect(withPrivateTextAdmissionConfig({ maxConcurrency: 8 }, state.initial).maxConcurrency).toBe( + 8, + ); + const ordinary = { maxConcurrency: 1 }; + expect(withPrivateTextAdmissionConfig(ordinary)).toBe(ordinary); +}); + +it.each([false, true])( + 'cancels an unadmitted protected write on Stop, pre-aborted: %s', + (aborted) => { + const { req } = turn(); + const controller = new AbortController(); + const start = jest.fn(); + const cancel = jest.fn(); + if (aborted) { + controller.abort(); + } + bindPrivateTextPersistenceAbort(req, controller.signal, start, cancel); + if (!aborted) { + controller.abort(); + } + expect(start).not.toHaveBeenCalled(); + expect(cancel).toHaveBeenCalledTimes(1); + }, +); + +it('retains ordinary Stop durability', () => { + const controller = new AbortController(); + const start = jest.fn(); + const cancel = jest.fn(); + bindPrivateTextPersistenceAbort({}, controller.signal, start, cancel); + controller.abort(); + expect(start).toHaveBeenCalledTimes(1); + expect(cancel).not.toHaveBeenCalled(); +}); + +it.each([false, true])( + 'coordinates the actual SDK parallel graph, rejected sibling: %s', + async (rejectSibling) => { + const { req, message, filters } = turn(); + const start = jest.fn(async () => ({ message })); + const created = jest.fn(); + const cancel = jest.fn(); + const initial = createPrivateTextInitialAdmissionCallback(req, { + agentIds: ['a', 'b'], + start, + cancel, + onPersisted: created, + })!; + const inspect = createModelBoundChatModelCallback( + { + filters: { + ...filters, + agentInstructions: { + pii: { + starterPatterns: [], + customPatterns: [{ id: 'history', label: 'Restricted', regex: 'RESTRICTED-HISTORY' }], + }, + }, + }, + }, + getPrivateTextModelHooks(req, initial, start, cancel, created), + ); + const invoke = jest.fn(); + const model = new FakeChatModel({ responses: ['Safe reply'], sleep: 0 }); + const stream = model._streamResponseChunks.bind(model); + jest.spyOn(model, '_streamResponseChunks').mockImplementation(async function* (...args) { + expect(created).toHaveBeenCalledTimes(1); + invoke(); + yield* stream(...args); + }); + // Same registration adapter used by the production model-client boundary. + model.callbacks = [ + Object.assign( + BaseCallbackHandler.fromMethods({ + handleChatModelStart: (llm, batches, runId, parentRunId, extra, tags, metadata) => + inspect.handleChatModelStart( + llm, + batches.map((batch) => + batch.map((entry) => ({ + role: entry.getType(), + content: + rejectSibling && metadata?.agentId === 'b' + ? 'RESTRICTED-HISTORY' + : entry.content, + })), + ), + runId, + parentRunId, + extra, + tags, + metadata, + ), + }), + { raiseError: true, awaitHandlers: true }, + ), + ]; + const run = await Run.create({ + runId: `private-parallel-${rejectSibling}`, + skipCleanup: true, + tokenCounter: () => 1, + graphConfig: { + type: 'multi-agent', + edges: [], + agents: ['a', 'b'].map((agentId) => ({ + agentId, + provider: Providers.OPENAI, + tools: [], + instructions: + rejectSibling && agentId === 'b' ? 'RESTRICTED-HISTORY' : 'Safe instructions', + maxContextTokens: 100_000, + clientOptions: { apiKey: 'fixture-key', streaming: false }, + })), + }, + }); + if (run.Graph == null) { + throw new Error('SDK graph was not initialized'); + } + run.Graph.overrideModel = model; + const config = withPrivateTextAdmissionConfig( + { configurable: { thread_id: 'private-parallel' }, version: 'v2' as const }, + initial, + ); + const execution = run.processStream( + { messages: [new HumanMessage('Safe turn')] }, + { + ...config, + callbacks: [ + Object.assign( + BaseCallbackHandler.fromMethods({ + handleChainStart: initial.handleChainStart, + handleChainEnd: initial.handleChainEnd, + handleChainError: initial.handleChainError, + handleLLMError: initial.handleLLMError, + }), + { raiseError: true, awaitHandlers: true }, + ), + ], + }, + ); + if (rejectSibling) { + await execution.catch(() => undefined); + expect(start).not.toHaveBeenCalled(); + expect(invoke).not.toHaveBeenCalled(); + } else { + await execution; + expect(start).toHaveBeenCalledTimes(1); + expect(invoke).toHaveBeenCalledTimes(2); + } + }, +); diff --git a/packages/api/src/protection/private/admission.ts b/packages/api/src/protection/private/admission.ts new file mode 100644 index 00000000000..246a00b83ae --- /dev/null +++ b/packages/api/src/protection/private/admission.ts @@ -0,0 +1,188 @@ +import type { InitialModelBoundAdmissionCallback } from '../../middleware/modelBoundContent'; +import type { getPrivateTextAdmission } from './submission'; +import { getPrivateTextAdmission as getAdmission, rejectPrivateTextAdmission } from './submission'; +import { ContentFilterError } from '../../middleware/contentFilter'; + +type ModelContext = { + runId?: string; + parentRunId?: string; + metadata?: Record; +}; +export interface ProtectedInitialAdmission extends InitialModelBoundAdmissionCallback { + readonly requiredConcurrency: number; + readonly admitModel: (context: ModelContext) => Promise | void; + readonly reject: (error?: unknown) => void; +} +function blocked(): ContentFilterError { + return new ContentFilterError({ + detectorId: 'pii-pattern', + ruleId: 'private-text', + label: 'private value that could not be protected', + source: 'message', + field: 'text', + provenance: 'user', + fragmentId: 'chat.text', + fragmentPath: '/text', + }); +} + +/** Every starting root must pass its exact native input before any root can invoke its model. */ +export function createPrivateTextInitialAdmissionCallback( + req: object | undefined, + options: { + agentIds: readonly string[]; + start: Parameters[1]; + cancel?: () => unknown; + onPersisted?: () => void; + signal?: AbortSignal; + }, +): ProtectedInitialAdmission | undefined { + const commit = getAdmission(req, options.start, options.onPersisted); + if (commit == null) { + return; + } + const pending = new Set(options.agentIds); + const parents = new Map(); + const rootNodes = new Map(); + let settle!: () => void; + let fail!: (error: unknown) => void; + const ready = new Promise((resolve, reject) => { + settle = resolve; + fail = reject; + }); + // A chain can fail before any model handler waits on the barrier. + void ready.catch(() => {}); + let state: 'pending' | 'committing' | 'admitted' | 'rejected' = 'pending'; + const detach = () => { + options.signal?.removeEventListener('abort', abort); + parents.clear(); + rootNodes.clear(); + }; + const reject = (error: unknown = blocked()) => { + if (state === 'admitted' || state === 'rejected') { + return; + } + state = 'rejected'; + rejectPrivateTextAdmission(req); + options.cancel?.(); + detach(); + fail(error); + }; + const abort = () => reject(); + options.signal?.addEventListener('abort', abort, { once: true }); + if (options.signal?.aborted || pending.size === 0) { + reject(); + } + const findRoot = (parentRunId: string | undefined): string | undefined => { + const seen = new Set(); + let id = parentRunId; + while (id != null && !seen.has(id)) { + seen.add(id); + const agent = rootNodes.get(id); + if (agent != null) { + return agent; + } + id = parents.get(id); + } + return; + }; + const admitModel = ({ parentRunId, metadata }: ModelContext) => { + if (state === 'admitted') { + return; + } + // Summary detours need a later graph superstep. Waiting would deadlock roots + // already at their native boundary; invoking would bypass turn admission. + if (metadata?.summarization === true) { + reject(); + return ready; + } + const agentId = findRoot(parentRunId); + if ( + !agentId || + metadata?.agentId !== agentId || + metadata.langgraph_node !== `agent=${agentId}` + ) { + reject(); + return ready; + } + pending.delete(agentId); + if (pending.size === 0 && state === 'pending') { + state = 'committing'; + void commit().then(() => { + if (state !== 'rejected') { + state = 'admitted'; + detach(); + settle(); + } + }, reject); + } + return ready; + }; + return Object.freeze({ + name: 'librechat-initial-model-bound-admission', + raiseError: true, + awaitHandlers: true, + requiredConcurrency: pending.size, + admitModel, + reject, + handleChainStart: (_chain, _inputs, runId, parentRunId, _tags, metadata, _runType, runName) => { + if (state !== 'pending') { + return; + } + if (parentRunId != null) { + parents.set(runId, parentRunId); + } + for (const agentId of pending) { + if (runName === `agent=${agentId}` && metadata?.langgraph_node === runName) { + rootNodes.set(runId, agentId); + break; + } + } + }, + handleChatModelStart: () => {}, + handleLLMEnd: () => {}, + handleLLMError: (error) => reject(error), + handleChainEnd: (_output, runId) => { + const agentId = rootNodes.get(runId); + if (agentId != null && pending.has(agentId)) { + reject(); + } + rootNodes.delete(runId); + parents.delete(runId); + }, + handleChainError: (error) => reject(error), + } satisfies ProtectedInitialAdmission); +} + +/** Ordinary resumed runs retain their legacy hook; captured runs use the coordinated barrier. */ +export function getPrivateTextModelHooks( + req: object | undefined, + initial: InitialModelBoundAdmissionCallback | ProtectedInitialAdmission | undefined, + start: Parameters[1], + cancel?: () => unknown, + onPersisted?: () => void, +): { + onContentAllowed?: (context: ModelContext) => void | Promise; + onContentRejected?: (error: unknown) => void; +} { + if (initial != null && 'admitModel' in initial) { + return { onContentAllowed: initial.admitModel, onContentRejected: initial.reject }; + } + return { onContentAllowed: getAdmission(req, start, onPersisted), onContentRejected: cancel }; +} + +export function withPrivateTextAdmissionConfig( + config: T & { maxConcurrency?: number }, + initial?: InitialModelBoundAdmissionCallback | ProtectedInitialAdmission, +): T & { callbacks?: InitialModelBoundAdmissionCallback[]; maxConcurrency?: number } { + return initial == null + ? config + : { + ...config, + callbacks: [initial], + ...('requiredConcurrency' in initial && + config.maxConcurrency != null && { + maxConcurrency: Math.max(config.maxConcurrency, initial.requiredConcurrency), + }), + }; +} diff --git a/packages/api/src/protection/private/submission.ts b/packages/api/src/protection/private/submission.ts index 1d66ddbbe2d..5921f901a33 100644 --- a/packages/api/src/protection/private/submission.ts +++ b/packages/api/src/protection/private/submission.ts @@ -465,6 +465,30 @@ export function rejectPrivateTextAdmission(req: object | undefined): void { } } +/** Stop before protected admission owns no persisted turn. Ordinary Stop retains its writer. */ +export function bindPrivateTextPersistenceAbort( + req: object | undefined, + signal: AbortSignal | undefined, + start: () => unknown, + cancel: () => unknown, +): () => void { + const abort = () => { + const capture = req == null ? undefined : captures.get(req); + if (capture != null && capture.admissionState() !== true) { + capture.admit(false); + cancel(); + return; + } + start(); + }; + if (signal?.aborted) { + abort(); + return () => {}; + } + signal?.addEventListener('abort', abort, { once: true }); + return () => signal?.removeEventListener('abort', abort); +} + /** Recovery may retain ordinary failures, but cannot revive an unadmitted protected turn. */ export async function savePrivateTextErrorTurn( req: object | undefined,