From b399ad83707ad38cf151e34e2db96eb9d229cd18 Mon Sep 17 00:00:00 2001 From: Oliver Faust <90081405+lidonius1122@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:24:43 +0100 Subject: [PATCH 1/8] =?UTF-8?q?=F0=9F=93=84=20fix:=20Serve=20Stored=20Text?= =?UTF-8?q?=20for=20"Upload=20as=20Text"=20File=20Downloads=20(#14723)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 📄 fix: Serve Stored Text for Text-Source File Downloads "Upload as Text" attachments store extracted content in the DB with source 'text'; OCR uploads persist the OCR strategy name (e.g. 'mistral_ocr') as a filepath placeholder since no backing file exists. The download route resolved these records to the local strategy and passed the placeholder to fs.createReadStream, which failed with ENOENT — and the response was never ended after the stream error, so the request hung until the client timed out. Serve the stored text directly as a .txt download for text-source files (re-fetched by _id, as getFiles excludes 'text' by default), and end the response on stream errors: 500 without the download headers before headers are sent, otherwise abort the truncated response so clients detect the failure. * fix: Preserve text-source preview semantics * fix: Complete text-source download coverage * fix: Tie text downloads to blob lifecycle * fix: Isolate preview downloads and share text snapshots * fix: Keep shared previews in share scope * style: Sort text download imports --------- Co-authored-by: Danny Avila --- api/server/routes/__tests__/share.spec.js | 22 +++ api/server/routes/files/files.js | 30 ++++ api/server/routes/files/files.test.js | 154 ++++++++++++++++++ api/server/routes/share.js | 17 +- .../components/Chat/Input/InFlightSteers.tsx | 1 + .../Messages/Content/FilePreviewDialog.tsx | 114 +++++-------- .../Chat/Messages/Content/Files.tsx | 1 + .../Chat/Messages/Content/Parts/LogLink.tsx | 10 +- .../Chat/Messages/Content/Parts/SteerPart.tsx | 1 + .../Content/Parts/__tests__/LogLink.test.tsx | 59 +++++++ .../Chat/Messages/Content/RetrievalCall.tsx | 10 ++ .../__tests__/FilePreviewDialog.test.ts | 25 +++ .../Content/__tests__/RetrievalCall.test.tsx | 16 +- .../Chat/Messages/Content/preview.ts | 79 +++++++++ client/src/components/Web/Citation.tsx | 3 + .../Web/__tests__/Citation.test.tsx | 16 +- client/src/data-provider/Files/queries.ts | 24 ++- .../src/utils/__tests__/downloadFile.test.ts | 54 +++++- client/src/utils/downloadFile.ts | 33 +++- .../data-schemas/src/methods/share.test.ts | 23 ++- packages/data-schemas/src/methods/share.ts | 28 +++- 21 files changed, 622 insertions(+), 98 deletions(-) create mode 100644 client/src/components/Chat/Messages/Content/__tests__/FilePreviewDialog.test.ts create mode 100644 client/src/components/Chat/Messages/Content/preview.ts diff --git a/api/server/routes/__tests__/share.spec.js b/api/server/routes/__tests__/share.spec.js index bd82f8c5027..ceb987c9b4a 100644 --- a/api/server/routes/__tests__/share.spec.js +++ b/api/server/routes/__tests__/share.spec.js @@ -945,6 +945,28 @@ describe('share-scoped file routes', () => { expect(response.headers['content-disposition']).toContain('attachment'); }); + it('downloads stored text for a snapshotted text-source file', async () => { + getFiles.mockResolvedValue([{ status: 'ready', text: 'Shared extracted text' }]); + getSharedLinkFile.mockResolvedValue({ + file: { + file_id: 'file-1', + source: 'text', + filepath: 'mistral_ocr', + type: 'application/pdf', + filename: 'report.pdf', + }, + hasSnapshots: true, + }); + + const response = await request(buildApp()).get('/api/share/share-123/files/file-1/download'); + + expect(response.status).toBe(200); + expect(response.headers['content-type']).toContain('text/plain'); + expect(response.headers['content-disposition']).toContain('attachment; report.pdf.txt'); + expect(response.text).toBe('Shared extracted text'); + expect(mockGetStrategyFunctions).not.toHaveBeenCalled(); + }); + it('returns 500 when the backing stream fails before sending bytes', async () => { const failingStream = new Readable({ read() { diff --git a/api/server/routes/files/files.js b/api/server/routes/files/files.js index 4f0caee52a4..612c970576e 100644 --- a/api/server/routes/files/files.js +++ b/api/server/routes/files/files.js @@ -570,6 +570,26 @@ router.get('/download/:userId/:file_id', fileAccess, async (req, res) => { // Access already validated by fileAccess middleware const file = req.fileAccess.file; + // Text-source files store extracted content in the DB; there is no backing file to stream + if (file.source === FileSources.text) { + /** `getFiles` excludes `text` by default, so the authorized record is re-fetched by `_id` */ + const [textFile] = (await db.getFiles({ _id: file._id }, null, { text: 1 })) ?? []; + if (textFile?.text == null) { + logger.warn(`File download requested by user ${userId} has no stored text: ${file_id}`); + return res.status(404).send('No file content found'); + } + const textFilename = file.filename?.toLowerCase().endsWith('.txt') + ? file.filename + : `${file.filename || file_id}.txt`; + res.setHeader('Content-Disposition', getContentDisposition(textFilename)); + res.setHeader('Content-Type', 'text/plain; charset=utf-8'); + res.setHeader( + 'X-File-Metadata', + encodeURIComponent(JSON.stringify(getDownloadFileMetadata(file))), + ); + return res.send(textFile.text); + } + if (checkOpenAIStorage(file.source) && !file.model) { logger.warn(`File download requested by user ${userId} has no associated model: ${file_id}`); return res.status(400).send('The model used when creating this file is not available'); @@ -642,6 +662,16 @@ router.get('/download/:userId/:file_id', fileAccess, async (req, res) => { fileStream.on('error', (streamError) => { logger.error('[DOWNLOAD ROUTE] Stream error:', streamError); + if (res.headersSent) { + if (!res.writableEnded) { + res.destroy(); + } + return; + } + res.removeHeader('Content-Disposition'); + res.removeHeader('Content-Type'); + res.removeHeader('X-File-Metadata'); + res.status(500).send('Error downloading file'); }); setHeaders(); diff --git a/api/server/routes/files/files.test.js b/api/server/routes/files/files.test.js index 4fcdd3a62a1..34894e2a566 100644 --- a/api/server/routes/files/files.test.js +++ b/api/server/routes/files/files.test.js @@ -940,6 +940,160 @@ describe('File Routes - Delete with Agent Access', () => { }), ); }); + + it('serves stored text for text-source files instead of streaming', async () => { + const userFileId = uuidv4(); + const getDownloadStream = jest.fn(); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'screenshot.png', + filepath: FileSources.mistral_ocr, + bytes: 70, + type: 'text/plain', + source: FileSources.text, + text: 'Extracted OCR text', + }); + + const response = await request(app).get(`/files/download/${otherUserId}/${userFileId}`); + + expect(response.status).toBe(200); + expect(response.headers['content-type']).toContain('text/plain'); + expect(response.headers['content-disposition']).toContain('screenshot.png.txt'); + expect(response.text).toBe('Extracted OCR text'); + const metadata = JSON.parse(decodeURIComponent(response.headers['x-file-metadata'])); + expect(metadata).toMatchObject({ file_id: userFileId, source: FileSources.text }); + expect(metadata).not.toHaveProperty('text'); + expect(getDownloadStream).not.toHaveBeenCalled(); + }); + + it('does not append .txt when the text-source filename already ends in .txt', async () => { + const userFileId = uuidv4(); + getStrategyFunctions.mockReturnValue({}); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'NOTES.TXT', + filepath: FileSources.mistral_ocr, + bytes: 20, + type: 'text/plain', + source: FileSources.text, + text: 'plain text notes', + }); + + const response = await request(app).get(`/files/download/${otherUserId}/${userFileId}`); + + expect(response.status).toBe(200); + expect(response.headers['content-disposition']).toContain('filename="NOTES.TXT"'); + expect(response.headers['content-disposition']).not.toContain('NOTES.TXT.txt'); + expect(response.text).toBe('plain text notes'); + }); + + it('returns 404 for text-source files without stored text', async () => { + const userFileId = uuidv4(); + const getDownloadStream = jest.fn(); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'empty.png', + filepath: FileSources.mistral_ocr, + bytes: 0, + type: 'text/plain', + source: FileSources.text, + }); + + const response = await request(app).get(`/files/download/${otherUserId}/${userFileId}`); + + expect(response.status).toBe(404); + expect(response.text).toBe('No file content found'); + expect(getDownloadStream).not.toHaveBeenCalled(); + }); + + it('serves a valid empty stored-text result', async () => { + const userFileId = uuidv4(); + const getDownloadStream = jest.fn(); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'empty.txt', + filepath: '/uploads/empty.txt', + bytes: 0, + type: 'text/plain', + source: FileSources.text, + text: '', + }); + + const response = await request(app).get(`/files/download/${otherUserId}/${userFileId}`); + + expect(response.status).toBe(200); + expect(response.headers['content-type']).toContain('text/plain'); + expect(response.text).toBe(''); + expect(getDownloadStream).not.toHaveBeenCalled(); + }); + + it('responds with 500 when the download stream errors before data is sent', async () => { + const userFileId = uuidv4(); + const erroringStream = new Readable({ + read() { + this.destroy(new Error('ENOENT: no such file or directory')); + }, + }); + const getDownloadStream = jest.fn().mockResolvedValue(erroringStream); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'gone.bin', + filepath: '/uploads/user/gone.bin', + bytes: 5, + type: 'application/octet-stream', + source: FileSources.local, + }); + + const response = await request(app).get(`/files/download/${otherUserId}/${userFileId}`); + + expect(response.status).toBe(500); + expect(response.text).toBe('Error downloading file'); + }); + + it('aborts the response when the download stream errors mid-transfer', async () => { + const userFileId = uuidv4(); + let pushed = false; + const erroringStream = new Readable({ + read() { + if (!pushed) { + pushed = true; + this.push('partial content'); + return; + } + this.destroy(new Error('read failed mid-stream')); + }, + }); + const getDownloadStream = jest.fn().mockResolvedValue(erroringStream); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + + await createFile({ + user: otherUserId, + file_id: userFileId, + filename: 'truncated.bin', + filepath: '/uploads/user/truncated.bin', + bytes: 100, + type: 'application/octet-stream', + source: FileSources.local, + }); + + await expect( + request(app).get(`/files/download/${otherUserId}/${userFileId}`), + ).rejects.toThrow(/aborted|socket hang up|ECONNRESET/i); + }); }); describe('POST /files/usage', () => { diff --git a/api/server/routes/share.js b/api/server/routes/share.js index c44ee8b7e64..c8068e11a17 100644 --- a/api/server/routes/share.js +++ b/api/server/routes/share.js @@ -197,7 +197,6 @@ const resolveShareFile = async (req, res, next) => { /** Stream (or redirect to) a snapshotted file from its original stored object. */ const streamSharedFile = async (req, res, file, requestedDisposition) => { const source = file.source || FileSources.local; - const { getDownloadStream, getDownloadURL } = getStrategyFunctions(source); // An update keeps the shareId, so these URLs are stable across re-publishes. Without // revalidation a viewer's cached copy would outlive a revoked "share files" choice or a @@ -209,6 +208,22 @@ const streamSharedFile = async (req, res, file, requestedDisposition) => { return res.status(304).end(); } + if (source === FileSources.text) { + if (req.liveFile?.text == null) { + return res.status(404).send('No file content found'); + } + const textFilename = file.filename?.toLowerCase().endsWith('.txt') + ? file.filename + : `${file.filename || file.file_id}.txt`; + const disposition = requestedDisposition === 'inline' ? 'inline' : 'attachment'; + res.setHeader('X-Content-Type-Options', 'nosniff'); + res.setHeader('Content-Disposition', getContentDisposition(textFilename, disposition)); + res.setHeader('Content-Type', 'text/plain; charset=utf-8'); + return res.send(req.liveFile.text); + } + + const { getDownloadStream, getDownloadURL } = getStrategyFunctions(source); + // Inline only safe preview types; anything else is forced to attachment. const disposition = requestedDisposition === 'inline' && SAFE_INLINE_TYPES.has(file.type) ? 'inline' : 'attachment'; diff --git a/client/src/components/Chat/Input/InFlightSteers.tsx b/client/src/components/Chat/Input/InFlightSteers.tsx index 31502e40174..633e6250a9a 100644 --- a/client/src/components/Chat/Input/InFlightSteers.tsx +++ b/client/src/components/Chat/Input/InFlightSteers.tsx @@ -552,6 +552,7 @@ const InFlightSteer = memo(function InFlightSteer({ fileId={selectedFile?.file_id} filePath={selectedFile?.filepath} fileType={selectedFile?.type ?? undefined} + fileSource={selectedFile?.source} fileSize={(selectedFile as TFile | null)?.bytes} /> )} diff --git a/client/src/components/Chat/Messages/Content/FilePreviewDialog.tsx b/client/src/components/Chat/Messages/Content/FilePreviewDialog.tsx index 79886c5d73c..958be9f4ab8 100644 --- a/client/src/components/Chat/Messages/Content/FilePreviewDialog.tsx +++ b/client/src/components/Chat/Messages/Content/FilePreviewDialog.tsx @@ -1,10 +1,11 @@ import { useState, useEffect, useCallback, useMemo, useRef } from 'react'; import copy from 'copy-to-clipboard'; -import { useRecoilValue } from 'recoil'; import { Download } from 'lucide-react'; +import { useRecoilValue } from 'recoil'; import { OGDialog, OGDialogContent, OGDialogTitle, OGDialogDescription } from '@librechat/client'; -import { useFileDownload, useSharedFileDownload } from '~/data-provider'; -import { logger, sortPagesByRelevance, triggerDownload } from '~/utils'; +import { getDownloadFilename, logger, sortPagesByRelevance, triggerDownload } from '~/utils'; +import { revokeDownloadURL, useFileDownload, useSharedFileDownload } from '~/data-provider'; +import { getFileExtension, getPreviewKind, shouldUseSharedFileDownload } from './preview'; import CopyButton from '~/components/Messages/Content/CopyButton'; import { useShareContext } from '~/Providers'; import { useLocalize } from '~/hooks'; @@ -20,69 +21,10 @@ interface FilePreviewDialogProps { pages?: number[]; pageRelevance?: Record; fileType?: string; + fileSource?: string; fileSize?: number; } -function getFileExtension(filename: string): string { - const dot = filename.lastIndexOf('.'); - return dot > 0 ? filename.slice(dot + 1).toLowerCase() : ''; -} - -function canPreviewByMime(mime?: string): 'pdf' | 'text' | false { - if (!mime) { - return false; - } - if (mime.includes('pdf')) { - return 'pdf'; - } - if ( - mime.startsWith('text/') || - mime.includes('json') || - mime.includes('xml') || - mime.includes('javascript') || - mime.includes('typescript') || - mime.includes('yaml') || - mime.includes('csv') - ) { - return 'text'; - } - return false; -} - -function canPreviewByExt(filename: string): 'pdf' | 'text' | false { - const ext = getFileExtension(filename); - if (ext === 'pdf') { - return 'pdf'; - } - const textExts = new Set([ - 'txt', - 'md', - 'csv', - 'json', - 'xml', - 'yaml', - 'yml', - 'html', - 'css', - 'js', - 'ts', - 'jsx', - 'tsx', - 'py', - 'rb', - 'java', - 'c', - 'cpp', - 'h', - 'go', - 'rs', - 'sh', - 'sql', - 'log', - ]); - return textExts.has(ext) ? 'text' : false; -} - /** Formats bytes with unit suffix (differs from ~/utils/formatBytes which returns a raw number). */ function formatBytes(bytes: number): string { if (bytes >= 1048576) { @@ -130,22 +72,30 @@ export default function FilePreviewDialog({ onOpenChange, fileName, fileId, - filePath, relevance, pages, pageRelevance, fileType, + fileSource, fileSize, }: FilePreviewDialogProps) { const localize = useLocalize(); const user = useRecoilValue(store.user); const { shareId } = useShareContext(); + // Preview reads revoke their blob after consumption, so they need a separate + // query identity from user-triggered downloads that may be in flight concurrently. const { refetch: downloadOwned } = useFileDownload(user?.id ?? '', fileId, { direct: false }); const { refetch: downloadShared } = useSharedFileDownload(shareId, fileId); - // Use the share route only for snapshotted files (filepath rewritten to the - // share path); otherwise fall back to the owner route. - const useShared = !!shareId && (filePath?.startsWith('/api/share/') ?? false); + const { refetch: previewOwned } = useFileDownload(user?.id ?? '', fileId, { + direct: false, + purpose: 'preview', + }); + const { refetch: previewShared } = useSharedFileDownload(shareId, fileId, 'preview'); + // A shared viewer must stay inside the share-scoped authorization boundary; + // citation and retrieval previews do not carry a rewritten filepath signal. + const useShared = shouldUseSharedFileDownload(shareId, fileId); const downloadFile = useShared ? downloadShared : downloadOwned; + const previewFile = useShared ? previewShared : previewOwned; const [fileContent, setFileContent] = useState(null); const [fileBlobUrl, setFileBlobUrl] = useState(null); @@ -154,7 +104,8 @@ export default function FilePreviewDialog({ const [isCopied, setIsCopied] = useState(false); const loadingRef = useRef(false); - const previewKind = canPreviewByMime(fileType) || canPreviewByExt(fileName); + const previewKind = getPreviewKind(fileName, fileType, fileSource); + const downloadFilename = getDownloadFilename(fileName, fileId, fileSource); const cancelledRef = useRef(false); @@ -168,16 +119,25 @@ export default function FilePreviewDialog({ setPreviewError(false); try { - const result = await downloadFile(); - if (cancelledRef.current || !result.data) { + const result = await previewFile(); + if (!result.data) { if (!cancelledRef.current) { setPreviewError(true); } return; } + if (cancelledRef.current) { + revokeDownloadURL(result.data); + return; + } - const resp = await fetch(result.data); - const blob = await resp.blob(); + let blob: Blob; + try { + const resp = await fetch(result.data); + blob = await resp.blob(); + } finally { + revokeDownloadURL(result.data); + } if (cancelledRef.current) { return; @@ -199,7 +159,7 @@ export default function FilePreviewDialog({ setLoading(false); } } - }, [fileId, previewKind, downloadFile]); + }, [fileId, previewKind, previewFile]); const handleDownload = useCallback(async () => { if (!fileId) { @@ -210,14 +170,14 @@ export default function FilePreviewDialog({ if (!result.data) { return; } - triggerDownload(result.data, fileName); + triggerDownload(result.data, downloadFilename); } catch (err) { logger.error('[FilePreviewDialog] Download failed:', err); } - }, [downloadFile, fileId, fileName]); + }, [downloadFile, downloadFilename, fileId]); useEffect(() => { - if (open && previewKind && !fileContent && !fileBlobUrl) { + if (open && previewKind && fileContent === null && !fileBlobUrl) { loadPreview(); } }, [open, previewKind, fileContent, fileBlobUrl, loadPreview]); @@ -315,7 +275,7 @@ export default function FilePreviewDialog({ className="h-[70vh] w-full rounded-lg border border-border-light" /> )} - {fileContent && ( + {fileContent !== null && ( <>
{ fileId={selectedFile?.file_id} filePath={selectedFile?.filepath} fileType={selectedFile?.type ?? undefined} + fileSource={selectedFile?.source} fileSize={(selectedFile as TFile)?.bytes} /> diff --git a/client/src/components/Chat/Messages/Content/Parts/LogLink.tsx b/client/src/components/Chat/Messages/Content/Parts/LogLink.tsx index 147e11cd905..d6760358df6 100644 --- a/client/src/components/Chat/Messages/Content/Parts/LogLink.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/LogLink.tsx @@ -1,8 +1,8 @@ import React from 'react'; import { useToastContext } from '@librechat/client'; import { FileSources, sharedFileDownload } from 'librechat-data-provider'; +import { getDownloadFilename, isHttpDownloadTarget, triggerDownload } from '~/utils'; import { useCodeOutputDownload, useFileDownload } from '~/data-provider'; -import { isHttpDownloadTarget, triggerDownload } from '~/utils'; import { useShareContext } from '~/Providers'; interface LogLinkProps { @@ -37,6 +37,7 @@ export const isLocallyStoredSource = (source?: string): boolean => { FileSources.s3, FileSources.cloudfront, FileSources.azure_blob, + FileSources.text, ].includes(source as FileSources); }; @@ -53,6 +54,7 @@ export const useAttachmentLink = ({ const useLocalDownload = isLocallyStoredSource(source) && !!file_id && !!user; const { refetch: downloadFromApi } = useFileDownload(user, file_id, { source }); const { refetch: downloadFromUrl } = useCodeOutputDownload(href); + const downloadFilename = getDownloadFilename(filename, file_id, source); /** * Triggers the download and reports whether a file was actually @@ -71,12 +73,12 @@ export const useAttachmentLink = ({ // permission, not owner ACL). Non-snapshotted files fall through so the // original href / code-output path still works when snapshots are disabled. if (shareId && file_id && href.startsWith('/api/share/')) { - triggerDownload(sharedFileDownload(shareId, file_id), filename); + triggerDownload(sharedFileDownload(shareId, file_id), downloadFilename); return true; } if (!useLocalDownload && isHttpDownloadTarget(href)) { - triggerDownload(href, filename); + triggerDownload(href, downloadFilename); return true; } @@ -89,7 +91,7 @@ export const useAttachmentLink = ({ }); return false; } - triggerDownload(stream.data, filename); + triggerDownload(stream.data, downloadFilename); return true; } catch (error) { console.error('Error downloading file:', error); diff --git a/client/src/components/Chat/Messages/Content/Parts/SteerPart.tsx b/client/src/components/Chat/Messages/Content/Parts/SteerPart.tsx index 237235d149d..dc1ab443e27 100644 --- a/client/src/components/Chat/Messages/Content/Parts/SteerPart.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/SteerPart.tsx @@ -133,6 +133,7 @@ const SteerPart = memo(function SteerPart({ fileId={selectedFile?.file_id} filePath={selectedFile?.filepath} fileType={selectedFile?.type ?? undefined} + fileSource={selectedFile?.source} fileSize={(selectedFile as TFile | null)?.bytes} /> )} diff --git a/client/src/components/Chat/Messages/Content/Parts/__tests__/LogLink.test.tsx b/client/src/components/Chat/Messages/Content/Parts/__tests__/LogLink.test.tsx index f4e1257a700..8152714bc2f 100644 --- a/client/src/components/Chat/Messages/Content/Parts/__tests__/LogLink.test.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/__tests__/LogLink.test.tsx @@ -23,6 +23,12 @@ jest.mock('~/Providers', () => ({ })); jest.mock('~/utils', () => ({ + getDownloadFilename: (filename: string, fileId?: string, source?: string) => { + const resolvedFilename = filename || fileId || 'download'; + return source === 'text' && !resolvedFilename.toLowerCase().endsWith('.txt') + ? `${resolvedFilename}.txt` + : resolvedFilename; + }, isHttpDownloadTarget: (target?: string | null) => /^https?:\/\//i.test(target ?? ''), triggerDownload: (...args: Parameters) => mockTriggerDownload(...args), @@ -129,4 +135,57 @@ describe('LogLink download routing', () => { expect(mockDownloadFromUrl).toHaveBeenCalledTimes(1); expect(mockDownloadFromApi).not.toHaveBeenCalled(); }); + + it('uses a text filename for shared text-source downloads', async () => { + mockShareContext = { shareId: 'share-9' }; + const filename = 'report.pdf'; + + render( + + {filename} + , + ); + + fireEvent.click(screen.getByRole('link', { name: filename })); + + await waitFor(() => { + expect(mockTriggerDownload).toHaveBeenCalledWith( + '/api/share/share-9/files/file-1/download', + 'report.pdf.txt', + ); + }); + }); + + it('uses the authorized file route and a text filename for owned text-source files', async () => { + const filename = 'report.pdf'; + mockDownloadFromApi.mockResolvedValue({ data: 'blob:https://app.example.com/text-file' }); + + render( + + {filename} + , + ); + + fireEvent.click(screen.getByRole('link', { name: filename })); + + await waitFor(() => { + expect(mockTriggerDownload).toHaveBeenCalledWith( + 'blob:https://app.example.com/text-file', + 'report.pdf.txt', + ); + }); + expect(mockDownloadFromApi).toHaveBeenCalledTimes(1); + expect(mockDownloadFromUrl).not.toHaveBeenCalled(); + }); }); diff --git a/client/src/components/Chat/Messages/Content/RetrievalCall.tsx b/client/src/components/Chat/Messages/Content/RetrievalCall.tsx index 73a1d67d20e..ea596b97fe4 100644 --- a/client/src/components/Chat/Messages/Content/RetrievalCall.tsx +++ b/client/src/components/Chat/Messages/Content/RetrievalCall.tsx @@ -23,6 +23,7 @@ interface FileSource { pageRelevance: Record; fileType?: string; fileBytes?: number; + fileSource?: string; metadata?: Record; } @@ -68,6 +69,7 @@ function extractFileSources(attachments?: TAttachment[]): FileSource[] { pageRelevance: source.pageRelevance || {}, fileType: (meta?.fileType as string) || undefined, fileBytes: (meta?.fileBytes as number) || undefined, + fileSource: (meta?.storageType as string) || undefined, metadata: meta, }); } @@ -92,6 +94,7 @@ interface DisplayResult { pageRelevance?: Record; fileType?: string; fileBytes?: number; + fileSource?: string; } interface FileMatch { @@ -99,6 +102,7 @@ interface FileMatch { fileName: string; fileType?: string; fileBytes?: number; + fileSource?: string; } function normalizeFilename(filename: string): string { @@ -142,6 +146,7 @@ function buildFileLookup( fileName: source.fileName, fileType: source.fileType, fileBytes: source.fileBytes, + fileSource: source.fileSource, }); } @@ -160,6 +165,7 @@ function buildFileLookup( fileName: file.filename, fileType: file.type ?? undefined, fileBytes: file.bytes, + fileSource: file.source ?? undefined, }); } @@ -181,6 +187,7 @@ function mergeRetrievalResults( pageRelevance: source.pageRelevance, fileType: source.fileType, fileBytes: source.fileBytes, + fileSource: source.fileSource, })); } @@ -197,6 +204,7 @@ function mergeRetrievalResults( content: result.content, fileType: match?.fileType, fileBytes: match?.fileBytes, + fileSource: match?.fileSource, }; }); } @@ -424,6 +432,7 @@ export default function RetrievalCall({ pages: result.pages, pageRelevance: result.pageRelevance, fileType: result.fileType, + fileSource: result.fileSource, }; }, [displayResults, previewIndex]); @@ -522,6 +531,7 @@ export default function RetrievalCall({ pages={previewData?.pages} pageRelevance={previewData?.pageRelevance} fileType={previewData?.fileType} + fileSource={previewData?.fileSource} />
); diff --git a/client/src/components/Chat/Messages/Content/__tests__/FilePreviewDialog.test.ts b/client/src/components/Chat/Messages/Content/__tests__/FilePreviewDialog.test.ts new file mode 100644 index 00000000000..6e8438dbb1f --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/FilePreviewDialog.test.ts @@ -0,0 +1,25 @@ +import { FileSources } from 'librechat-data-provider'; +import { getPreviewKind, shouldUseSharedFileDownload } from '../preview'; +import { getDownloadFilename } from '~/utils/downloadFile'; + +describe('FilePreviewDialog text-source behavior', () => { + it('previews extracted PDF content as text', () => { + expect(getPreviewKind('report.pdf', 'application/pdf', FileSources.text)).toBe('text'); + }); + + it('downloads extracted content with a text extension', () => { + expect(getDownloadFilename('report.pdf', 'file-1', FileSources.text)).toBe('report.pdf.txt'); + expect(getDownloadFilename('notes.txt', 'file-2', FileSources.text)).toBe('notes.txt'); + }); + + it('preserves the original behavior for stored files', () => { + expect(getPreviewKind('report.pdf', 'application/pdf', FileSources.local)).toBe('pdf'); + expect(getDownloadFilename('report.pdf', 'file-3', FileSources.local)).toBe('report.pdf'); + }); + + it('routes any identified file through the share boundary in a shared view', () => { + expect(shouldUseSharedFileDownload('share-1', 'file-1')).toBe(true); + expect(shouldUseSharedFileDownload('share-1', undefined)).toBe(false); + expect(shouldUseSharedFileDownload(undefined, 'file-1')).toBe(false); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/__tests__/RetrievalCall.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/RetrievalCall.test.tsx index 6df45a32105..19b8c3296d4 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/RetrievalCall.test.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/RetrievalCall.test.tsx @@ -82,9 +82,19 @@ jest.mock('~/data-provider', () => ({ jest.mock('../FilePreviewDialog', () => ({ __esModule: true, - default: ({ open, fileId, fileName }: { open: boolean; fileId?: string; fileName: string }) => + default: ({ + open, + fileId, + fileName, + fileSource, + }: { + open: boolean; + fileId?: string; + fileName: string; + fileSource?: string; + }) => open ? ( -
+
{fileName}
) : null, @@ -220,6 +230,7 @@ describe('RetrievalCall - file preview resolution', () => { filename: 'Tutorial Imazing.pdf', bytes: 2048, type: 'application/pdf', + source: 'text', }, ], }); @@ -235,6 +246,7 @@ describe('RetrievalCall - file preview resolution', () => { fireEvent.click(screen.getByRole('button', { name: 'Preview: Tutorial Imazing.pdf' })); expect(screen.getByTestId('file-preview-dialog')).toHaveAttribute('data-file-id', 'file-123'); + expect(screen.getByTestId('file-preview-dialog')).toHaveAttribute('data-file-source', 'text'); }); it('keeps multiple parsed results clickable when only one attachment source is available', () => { diff --git a/client/src/components/Chat/Messages/Content/preview.ts b/client/src/components/Chat/Messages/Content/preview.ts new file mode 100644 index 00000000000..d369cbf597e --- /dev/null +++ b/client/src/components/Chat/Messages/Content/preview.ts @@ -0,0 +1,79 @@ +import { FileSources } from 'librechat-data-provider'; + +type PreviewKind = 'pdf' | 'text' | false; + +const TEXT_EXTENSIONS = new Set([ + 'txt', + 'md', + 'csv', + 'json', + 'xml', + 'yaml', + 'yml', + 'html', + 'css', + 'js', + 'ts', + 'jsx', + 'tsx', + 'py', + 'rb', + 'java', + 'c', + 'cpp', + 'h', + 'go', + 'rs', + 'sh', + 'sql', + 'log', +]); + +export function getFileExtension(filename: string): string { + const dot = filename.lastIndexOf('.'); + return dot > 0 ? filename.slice(dot + 1).toLowerCase() : ''; +} + +export function shouldUseSharedFileDownload(shareId?: string, fileId?: string): boolean { + return !!shareId && !!fileId; +} + +function getPreviewKindByMime(mime?: string): PreviewKind { + if (!mime) { + return false; + } + if (mime.includes('pdf')) { + return 'pdf'; + } + if ( + mime.startsWith('text/') || + mime.includes('json') || + mime.includes('xml') || + mime.includes('javascript') || + mime.includes('typescript') || + mime.includes('yaml') || + mime.includes('csv') + ) { + return 'text'; + } + return false; +} + +function getPreviewKindByExtension(filename: string): PreviewKind { + const extension = getFileExtension(filename); + if (extension === 'pdf') { + return 'pdf'; + } + return TEXT_EXTENSIONS.has(extension) ? 'text' : false; +} + +export function getPreviewKind( + fileName: string, + fileType?: string, + fileSource?: string, +): PreviewKind { + if (fileSource === FileSources.text) { + return 'text'; + } + return getPreviewKindByMime(fileType) || getPreviewKindByExtension(fileName); +} diff --git a/client/src/components/Web/Citation.tsx b/client/src/components/Web/Citation.tsx index d6661e62161..67f3ffafba5 100644 --- a/client/src/components/Web/Citation.tsx +++ b/client/src/components/Web/Citation.tsx @@ -11,6 +11,7 @@ import { useLocalize } from '~/hooks'; interface FileCitationMetadata { fileBytes?: number; fileType?: string; + storageType?: string; } interface FileCitationSource { @@ -282,6 +283,7 @@ export function CompositeCitation(props: CompositeCitationProps) { pages={filePages} pageRelevance={filePageRelevance} fileType={fileMeta?.fileType} + fileSource={fileMeta?.storageType} fileSize={fileMeta?.fileBytes} /> )} @@ -358,6 +360,7 @@ export function Citation(props: CitationComponentProps) { pages={filePages} pageRelevance={filePageRelevance} fileType={fileMeta?.fileType} + fileSource={fileMeta?.storageType} fileSize={fileMeta?.fileBytes} /> )} diff --git a/client/src/components/Web/__tests__/Citation.test.tsx b/client/src/components/Web/__tests__/Citation.test.tsx index 6ec16d4c04d..d25cb2d7611 100644 --- a/client/src/components/Web/__tests__/Citation.test.tsx +++ b/client/src/components/Web/__tests__/Citation.test.tsx @@ -29,9 +29,19 @@ jest.mock('~/hooks', () => ({ jest.mock('~/components/Chat/Messages/Content/FilePreviewDialog', () => ({ __esModule: true, - default: ({ open, fileId, fileName }: { open: boolean; fileId?: string; fileName: string }) => + default: ({ + open, + fileId, + fileName, + fileSource, + }: { + open: boolean; + fileId?: string; + fileName: string; + fileSource?: string; + }) => open ? ( -
+
{fileName}
) : null, @@ -76,6 +86,7 @@ describe('Citation', () => { metadata: { fileBytes: 2048, fileType: 'application/pdf', + storageType: 'text', }, pageRelevance: { 1: 0.92 }, pages: [1], @@ -107,6 +118,7 @@ describe('Citation', () => { fireEvent.click(fileButton); expect(screen.getByTestId('file-preview-dialog')).toHaveAttribute('data-file-id', 'file-123'); + expect(screen.getByTestId('file-preview-dialog')).toHaveAttribute('data-file-source', 'text'); }); it('keeps standalone web citations as links', () => { diff --git a/client/src/data-provider/Files/queries.ts b/client/src/data-provider/Files/queries.ts index 81f161d6280..da6981571e0 100644 --- a/client/src/data-provider/Files/queries.ts +++ b/client/src/data-provider/Files/queries.ts @@ -3,8 +3,13 @@ import { useQuery, useQueryClient } from '@tanstack/react-query'; import { FileSources, QueryKeys, DynamicQueryKeys, dataService } from 'librechat-data-provider'; import type { QueryObserverResult, UseQueryOptions } from '@tanstack/react-query'; import type t from 'librechat-data-provider'; +import { + addFileToCache, + getDownloadFilename, + registerDownloadFilename, + unregisterDownloadFilename, +} from '~/utils'; import { isEphemeralAgent } from '~/common'; -import { addFileToCache } from '~/utils'; import store from '~/store'; export const useGetFiles = ( @@ -56,6 +61,7 @@ export const useGetFileConfig = ( type FileDownloadOptions = { source?: string | null; direct?: boolean; + purpose?: 'download' | 'preview'; }; export const isDirectDownloadSource = (source?: string | null): boolean => @@ -65,6 +71,7 @@ export const revokeDownloadURL = (url?: string | null): void => { if (!url?.startsWith('blob:')) { return; } + unregisterDownloadFilename(url); window.URL.revokeObjectURL(url); }; @@ -75,7 +82,13 @@ export const useFileDownload = ( ): QueryObserverResult => { const queryClient = useQueryClient(); return useQuery( - [QueryKeys.fileDownload, file_id, options.source ?? '', options.direct ?? true], + [ + QueryKeys.fileDownload, + file_id, + options.source ?? '', + options.direct ?? true, + options.purpose ?? 'download', + ], async () => { if (!userId || !file_id) { console.warn('No user ID provided for file download'); @@ -104,6 +117,10 @@ export const useFileDownload = ( return downloadURL; } + registerDownloadFilename( + downloadURL, + getDownloadFilename(metadata.filename, metadata.file_id, metadata.source), + ); addFileToCache(queryClient, metadata); } catch (e) { console.error('Error parsing file metadata, skipped updating file query cache', e); @@ -126,9 +143,10 @@ export const useFileDownload = ( export const useSharedFileDownload = ( shareId?: string, file_id?: string, + purpose: 'download' | 'preview' = 'download', ): QueryObserverResult => { return useQuery( - [QueryKeys.fileDownload, 'share', shareId ?? '', file_id ?? ''], + [QueryKeys.fileDownload, 'share', shareId ?? '', file_id ?? '', purpose], async () => { if (!shareId || !file_id) { return; diff --git a/client/src/utils/__tests__/downloadFile.test.ts b/client/src/utils/__tests__/downloadFile.test.ts index ed11c3df856..9d64e961ef2 100644 --- a/client/src/utils/__tests__/downloadFile.test.ts +++ b/client/src/utils/__tests__/downloadFile.test.ts @@ -1,4 +1,12 @@ -import { getCodeBlockFilename, isHttpDownloadTarget, triggerDownload } from '../downloadFile'; +import { FileSources } from 'librechat-data-provider'; +import { + getCodeBlockFilename, + getDownloadFilename, + isHttpDownloadTarget, + registerDownloadFilename, + triggerDownload, + unregisterDownloadFilename, +} from '../downloadFile'; describe('downloadFile utilities', () => { let clickSpy: jest.SpyInstance; @@ -68,6 +76,50 @@ describe('downloadFile utilities', () => { jest.advanceTimersByTime(1000); expect(revokeSpy).toHaveBeenCalledWith('blob:https://app.example.com/download-id'); }); + + it('uses registered response metadata to name blob downloads', () => { + const target = 'blob:https://app.example.com/text-download'; + registerDownloadFilename(target, 'report.pdf.txt'); + + triggerDownload(target, 'report.pdf'); + + expect(appendedLink?.download).toBe('report.pdf.txt'); + }); + + it('keeps registered names available for concurrent blob downloads', () => { + const target = 'blob:https://app.example.com/concurrent-download'; + registerDownloadFilename(target, 'report.pdf.txt'); + + triggerDownload(target, 'report.pdf'); + expect(appendedLink?.download).toBe('report.pdf.txt'); + + triggerDownload(target, 'report.pdf'); + expect(appendedLink?.download).toBe('report.pdf.txt'); + }); + + it('clears registered names when blob URLs are released', () => { + const target = 'blob:https://app.example.com/released-download'; + registerDownloadFilename(target, 'report.pdf.txt'); + unregisterDownloadFilename(target); + + triggerDownload(target, 'report.pdf'); + + expect(appendedLink?.download).toBe('report.pdf'); + }); +}); + +describe('getDownloadFilename', () => { + it('adds a text extension for text-source files', () => { + expect(getDownloadFilename('report.pdf', 'file-1', FileSources.text)).toBe('report.pdf.txt'); + }); + + it('recognizes existing text extensions case-insensitively', () => { + expect(getDownloadFilename('NOTES.TXT', 'file-2', FileSources.text)).toBe('NOTES.TXT'); + }); + + it('preserves filenames for other storage sources', () => { + expect(getDownloadFilename('report.pdf', 'file-3', FileSources.local)).toBe('report.pdf'); + }); }); describe('getCodeBlockFilename', () => { diff --git a/client/src/utils/downloadFile.ts b/client/src/utils/downloadFile.ts index ddccc53a2ed..73e5f99e9de 100644 --- a/client/src/utils/downloadFile.ts +++ b/client/src/utils/downloadFile.ts @@ -1,6 +1,32 @@ +import { FileSources } from 'librechat-data-provider'; + +const blobDownloadFilenames = new Map(); + export const isHttpDownloadTarget = (target?: string | null): boolean => /^https?:\/\//i.test(target ?? ''); +export function getDownloadFilename( + fileName: string, + fileId?: string, + fileSource?: string | null, +): string { + const filename = fileName || fileId || 'download'; + if (fileSource !== FileSources.text || filename.toLowerCase().endsWith('.txt')) { + return filename; + } + return `${filename}.txt`; +} + +export function registerDownloadFilename(target: string, filename: string): void { + if (target.startsWith('blob:')) { + blobDownloadFilenames.set(target, filename); + } +} + +export function unregisterDownloadFilename(target: string): void { + blobDownloadFilenames.delete(target); +} + /** * Maps a fenced-block language hint to a file extension. Used to name * downloads of chat code blocks (`code.`). Only languages whose common @@ -63,11 +89,14 @@ export function triggerDownload(target: string, filename: string): void { const isBlob = target.startsWith('blob:'); const link = document.createElement('a'); link.href = target; - link.setAttribute('download', filename); + link.setAttribute('download', blobDownloadFilenames.get(target) ?? filename); document.body.appendChild(link); link.click(); document.body.removeChild(link); if (isBlob) { - setTimeout(() => URL.revokeObjectURL(target), 1000); + setTimeout(() => { + unregisterDownloadFilename(target); + URL.revokeObjectURL(target); + }, 1000); } } diff --git a/packages/data-schemas/src/methods/share.test.ts b/packages/data-schemas/src/methods/share.test.ts index fbd11894456..4f756345939 100644 --- a/packages/data-schemas/src/methods/share.test.ts +++ b/packages/data-schemas/src/methods/share.test.ts @@ -2794,11 +2794,15 @@ describe('Share Methods', () => { expect(result?.updatedAt?.getTime()).toBe(published?.updatedAt?.getTime()); }); - test('does not snapshot transient text-source files', async () => { + test('snapshots database-backed text-source files without embedding their text', async () => { const userId = new mongoose.Types.ObjectId().toString(); const conversationId = `conv_${nanoid()}`; await seedConversation(userId, conversationId); - const textId = await createFile(userId, { source: 'text' }); + const textId = await createFile(userId, { + source: 'text', + filepath: 'mistral_ocr', + text: 'Extracted text', + }); await Message.create({ messageId: `msg_${nanoid()}`, conversationId, @@ -2810,7 +2814,20 @@ describe('Share Methods', () => { const result = await shareMethods.createSharedLink(userId, conversationId); const saved = await SharedLink.findOne({ shareId: result.shareId }).lean(); - expect(saved?.fileSnapshots ?? []).toHaveLength(0); + expect(saved?.fileSnapshots).toHaveLength(1); + expect(saved?.fileSnapshots?.[0]).toMatchObject({ + file_id: textId, + source: 'text', + filepath: 'mistral_ocr', + }); + expect(saved?.fileSnapshots?.[0]).not.toHaveProperty('text'); + + const shared = await shareMethods.getSharedMessages(result.shareId); + expect(shared?.messages[0].files?.[0]).toMatchObject({ + file_id: textId, + source: 'text', + filepath: `/api/share/${result.shareId}/files/${textId}`, + }); }); test('updateSharedLink clears snapshots when snapshotFiles is disabled', async () => { diff --git a/packages/data-schemas/src/methods/share.ts b/packages/data-schemas/src/methods/share.ts index 87d470efc16..73f1640a860 100644 --- a/packages/data-schemas/src/methods/share.ts +++ b/packages/data-schemas/src/methods/share.ts @@ -135,8 +135,8 @@ function sanitizeSharedAttachments(attachments: unknown): t.SharedFile[] | undef * stream with only `storageKey`/`filepath` + the request. Sources requiring * owner-specific credentials (openai/azure assistants, execute_code, vectordb, * OCR/parser pipelines) are skipped — those files degrade to a 404 in the share - * view. `FileSources.text` is intentionally excluded: its `filepath` is a Multer - * temp path that the upload route deletes, so there is nothing durable to stream. + * view. Text-source files are eligible because the share route serves their + * database-backed extracted text instead of the deleted Multer temp path. */ const SNAPSHOT_STREAMABLE_SOURCES = new Set([ FileSources.local, @@ -144,6 +144,7 @@ const SNAPSHOT_STREAMABLE_SOURCES = new Set([ FileSources.cloudfront, FileSources.azure_blob, FileSources.firebase, + FileSources.text, ]); /** Collect `file_id`s from a message's `files`/`attachments` array into `target`. */ @@ -358,11 +359,18 @@ function applyShareFileRoute( file: t.SharedFile, shareId: string, snapshotIds: Set, + textSourceIds?: Set, ): t.SharedFile { const fileId = file.file_id; if (typeof fileId === 'string' && snapshotIds.has(fileId)) { const route = shareFileRoute(shareId, fileId); - const next: t.SharedFile = { ...file, filepath: route }; + const next: t.SharedFile = { + ...file, + filepath: route, + // General storage sources stay private, but `text` is a render semantic: + // clients must preview the database-backed payload as text, not the original MIME. + ...(textSourceIds?.has(fileId) && { source: FileSources.text }), + }; if (file.preview !== undefined) { next.preview = route; } @@ -390,6 +398,7 @@ export function anonymizeSharedContent( newMessageId: string; shareId: string; snapshotIds: Set; + textSourceIds?: Set; includeFiles: boolean; sanitizeUIResourceMarkers?: boolean; }, @@ -420,6 +429,7 @@ export function anonymizeSharedContent( }, params.shareId, params.snapshotIds, + params.textSourceIds, ), ) : undefined; @@ -456,6 +466,7 @@ function anonymizeMessages( newConvoId: string, shareId: string, snapshotIds: Set, + textSourceIds: Set, includeFiles: boolean, anonymizeMessageId: (id: string) => string, anonymizeAssistantId: (id: string) => string, @@ -481,6 +492,7 @@ function anonymizeMessages( }, shareId, snapshotIds, + textSourceIds, ), ) : undefined; @@ -496,6 +508,7 @@ function anonymizeMessages( }, shareId, snapshotIds, + textSourceIds, ), ) : undefined; @@ -517,6 +530,7 @@ function anonymizeMessages( newMessageId, shareId, snapshotIds, + textSourceIds, includeFiles, sanitizeUIResourceMarkers: message.isCreatedByUser !== true, }), @@ -830,6 +844,13 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { const snapshotIds = includeFiles ? new Set((fileSnapshots ?? []).map((snapshot) => snapshot.file_id)) : new Set(); + const textSourceIds = includeFiles + ? new Set( + (fileSnapshots ?? []) + .filter((snapshot) => snapshot.source === FileSources.text) + .map((snapshot) => snapshot.file_id), + ) + : new Set(); const result: t.SharedMessagesResult = { shareId: resolvedShareId, title: share.title, @@ -841,6 +862,7 @@ export function createShareMethods(mongoose: typeof import('mongoose')): { newConvoId, resolvedShareId, snapshotIds, + textSourceIds, includeFiles, anonymizeMessageId, anonymizeAssistantId, From 9f8d71a3c5926438831d05d54881fda6e8b64a36 Mon Sep 17 00:00:00 2001 From: Yorgos K Date: Fri, 21 Aug 2026 20:25:03 +0200 Subject: [PATCH 2/8] =?UTF-8?q?=F0=9F=AA=A2=20fix:=20Preserve=20Response?= =?UTF-8?q?=20Identity=20and=20Branch=20During=20Resumable=20SSE=20Sync=20?= =?UTF-8?q?(#14788)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(client): preserve resumable response identity Fixes #14787 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(client): align resumable sync regressions Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(client): clarify resumable response ownership Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(client): preserve resumed regeneration ordering Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(client): cover missing resumed response row Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(client): preserve resume identity on page reload * fix(client): replace reassigned resume placeholder * fix(client): preserve content during response id handoff * fix(client): limit resume placeholder handoff * fix(client): preserve resume display metadata * fix(client): reconcile resume metadata in one pass * fix(client): reconcile preliminary resume user * fix(client): restore regenerated branch on early abort * test(client): cover external regeneration resume * fix(client): preserve regeneration history on errors * fix(client): replace reused regeneration error ids * fix(client): preserve exact-id regeneration rollback --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Danny Avila --- .../__tests__/request.resumeMetadata.spec.js | 42 +++ api/server/controllers/agents/request.js | 1 + .../SSE/__tests__/useEventHandlers.spec.ts | 59 ++++ .../SSE/__tests__/useResumableSSE.spec.ts | 291 +++++++++++++++++- .../SSE/__tests__/useResumeOnLoad.spec.tsx | 209 ++++++++++++- client/src/hooks/SSE/useEventHandlers.ts | 33 +- client/src/hooks/SSE/useResumableSSE.ts | 187 ++++++++--- client/src/hooks/SSE/useResumeOnLoad.ts | 49 +-- .../api/src/stream/GenerationJobManager.ts | 2 + .../GenerationJobManager.resumeReplay.spec.ts | 16 + .../api/src/stream/__tests__/startup.spec.ts | 2 + .../stream/implementations/RedisJobStore.ts | 1 + .../api/src/stream/interfaces/IJobStore.ts | 4 + packages/api/src/stream/metadata.ts | 3 + packages/api/src/types/stream.ts | 2 + packages/data-provider/src/types/agents.ts | 2 + 16 files changed, 822 insertions(+), 81 deletions(-) diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index cb56912c60f..ca27a838fed 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -1080,6 +1080,48 @@ describe('ResumableAgentController resume metadata', () => { ); }); + it('records regeneration ownership for exact-ID resume reconstruction', async () => { + const conversationId = 'conversation-123'; + const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Regenerate the edited response.', + messageId: 'user-message', + parentMessageId: 'parent-message', + responseMessageId: 'edited-response', + isRegenerate: true, + conversationId, + endpointOption: { + endpoint: 'agents', + modelOptions: { model: 'gpt-4.1' }, + }, + }, + config: {}, + }; + const res = { + headersSent: true, + json: jest.fn(() => { + res.headersSent = true; + }), + status: jest.fn(() => res), + }; + + await AgentController(req, res, jest.fn(), initializeClient, null); + + expect(mockGenerationJobManager.createJob).toHaveBeenCalledWith( + conversationId, + 'user-123', + conversationId, + expect.objectContaining({ + initialMetadata: expect.objectContaining({ + responseMessageId: 'edited-response', + isRegenerate: true, + }), + }), + ); + }); + it('falls back to the model spec preset endpoint when no icon URL is configured', async () => { const conversationId = 'conversation-123'; const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index c6b9892e8fe..1dedc6617d0 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -1050,6 +1050,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit // Persist temporary-chat state so a HITL resume keeps the resumed response // non-persisted instead of trusting the resume request to re-send the flag. isTemporary: req.body?.isTemporary, + ...(isRegenerate && { isRegenerate: true }), ...(scheduleId ? { scheduleId, diff --git a/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts b/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts index 49f5b37f14b..08fb6943678 100644 --- a/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts +++ b/client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts @@ -4,6 +4,7 @@ import { buildCreatedInitialResponse, getExistingConversationAbortMessages, isInitialNewConversationSubmission, + mergeErrorMessages, mergeRegenerateFinalMessages, startedAsNewConversation, } from '~/hooks/SSE/useEventHandlers'; @@ -230,3 +231,61 @@ describe('getExistingConversationAbortMessages', () => { ).toEqual(['user-1']); }); }); + +describe('mergeErrorMessages', () => { + const message = (messageId: string, isCreatedByUser = false) => + ({ + messageId, + conversationId: 'conversation-1', + isCreatedByUser, + text: messageId, + }) as TMessage; + + it('adds the request and error for a normal submission', () => { + const userMessage = message('user-1', true); + const errorMessage = message('assistant-error'); + + expect( + mergeErrorMessages({ + messages: [message('previous-response')], + userMessage, + errorMessage, + }).map(({ messageId }) => messageId), + ).toEqual(['previous-response', 'user-1', 'assistant-error']); + }); + + it('preserves regeneration history without duplicating its user', () => { + const userMessage = message('user-1', true); + const originalResponse = message('assistant-1'); + const laterUser = message('user-2', true); + const laterResponse = message('assistant-2'); + const errorMessage = message('assistant-1_'); + + expect( + mergeErrorMessages({ + messages: [userMessage], + regenerateMessages: [userMessage, originalResponse, laterUser, laterResponse], + userMessage, + errorMessage, + isRegenerate: true, + }).map(({ messageId }) => messageId), + ).toEqual(['user-1', 'assistant-1', 'user-2', 'assistant-2', 'assistant-1_']); + }); + + it('replaces an edited response error that intentionally reuses its id', () => { + const userMessage = message('user-1', true); + const originalResponse = message('assistant-1'); + const errorMessage = { ...originalResponse, text: 'Regeneration failed', error: true }; + + const merged = mergeErrorMessages({ + messages: [userMessage], + regenerateMessages: [userMessage, originalResponse], + userMessage, + errorMessage, + isRegenerate: true, + }); + + expect(merged.map(({ messageId }) => messageId)).toEqual(['user-1', 'assistant-1']); + expect(merged[1]).toEqual(errorMessage); + }); +}); diff --git a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts index 4f09207441c..49fd41fe8e6 100644 --- a/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts +++ b/client/src/hooks/SSE/__tests__/useResumableSSE.spec.ts @@ -2964,14 +2964,19 @@ describe('useResumableSSE', () => { unmount(); }); - /** - * Regenerate: the new run's response id is not in the loaded history yet, so the - * parent-based fallback lands on the answer being REPLACED. Preserving that row's - * content would make the regenerated run's deltas append to the stale answer, so an - * empty snapshot must still clear a row we only matched heuristically. - */ - it('does not preserve content on a row matched only by the parent fallback', async () => { - const submission = buildSubmission(); + it('does not reuse an older response that only shares the user parent', async () => { + const submission = buildSubmission({ + initialResponse: { + messageId: 'resp-1', + conversationId: CONV_ID, + text: '', + isCreatedByUser: false, + sender: 'Custom Assistant', + endpoint: 'azureOpenAI', + iconURL: 'https://example.com/assistant.png', + model: 'gpt-4.1', + }, + }); const chatHelpers = buildChatHelpers(); chatHelpers.getMessages.mockReturnValue([ { @@ -3009,12 +3014,27 @@ describe('useResumableSSE', () => { }); }); - const synced = chatHelpers.setMessages.mock.calls + const syncedMessages = chatHelpers.setMessages.mock.calls .map(([messages]) => messages as TMessage[]) .reverse() - .find((messages) => messages?.some((m) => m.messageId === 'resp-previous')) - ?.find((m) => m.messageId === 'resp-previous'); - expect(synced?.content).toEqual([]); + .find((messages) => messages?.some((m) => m.messageId === 'resp-regenerated')); + expect(syncedMessages?.find((m) => m.messageId === 'resp-previous')?.content).toEqual([ + { type: 'text', text: 'the answer being regenerated' }, + ]); + expect(syncedMessages?.map((message) => message.messageId)).toEqual([ + 'msg-1', + 'resp-previous', + 'resp-regenerated', + ]); + expect(syncedMessages?.find((m) => m.messageId === 'resp-regenerated')).toEqual( + expect.objectContaining({ + content: [], + sender: 'Custom Assistant', + endpoint: 'azureOpenAI', + iconURL: 'https://example.com/assistant.png', + model: 'gpt-4.1', + }), + ); unmount(); }); @@ -4128,3 +4148,250 @@ describe('useResumableSSE', () => { unmount(); }); }); + +describe('useResumableSSE - sync response identity', () => { + beforeEach(() => { + mockSSEInstances.length = 0; + mockSetIsSubmitting.mockClear(); + }); + + const emitSync = async ( + sse: MockSSEInstance, + aggregatedContent: TMessage['content'], + responseMessageId?: string, + sender?: string, + userMessage?: Partial, + ) => { + await act(async () => { + sse._emit('message', { + data: JSON.stringify({ + sync: true, + resumeState: { aggregatedContent, responseMessageId, sender, userMessage }, + }), + }); + }); + }; + + it('updates the submission-owned response when sync omits the response ID', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const activeResponse = { + messageId: 'server-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, activeResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const aggregatedContent: TMessage['content'] = [ + { type: ContentTypes.TEXT, text: { value: 'Recovered answer' } }, + ]; + await emitSync(getLastSSE(), aggregatedContent); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages).toHaveLength(2); + expect(updatedMessages.find((message) => message.messageId === 'server-response-id')).toEqual({ + ...activeResponse, + content: aggregatedContent, + }); + expect( + updatedMessages.find((message) => message.messageId === 'server-user-id_'), + ).toBeUndefined(); + unmount(); + }); + + it('adds resumed sender metadata to an exact persisted response', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const activeResponse = { + messageId: 'server-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, activeResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + await emitSync(getLastSSE(), [], activeResponse.messageId, 'Restored Assistant'); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages[1]).toEqual({ + ...activeResponse, + sender: 'Restored Assistant', + }); + unmount(); + }); + + it('appends a missing submission-owned response after older regeneration siblings', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const olderResponse = { + messageId: 'older-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Earlier answer', + content: [{ type: 'text', text: { value: 'Earlier answer' } }], + isCreatedByUser: false, + } as TMessage; + const activeResponse = { + messageId: 'active-response-id', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: activeResponse }), + isRegenerate: true, + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, olderResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const aggregatedContent: TMessage['content'] = [ + { type: ContentTypes.TEXT, text: { value: 'Regenerated answer' } }, + ]; + await emitSync(getLastSSE(), aggregatedContent); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.find((message) => message.messageId === 'older-response-id')).toEqual( + olderResponse, + ); + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'server-user-id', + 'older-response-id', + 'active-response-id', + ]); + expect(updatedMessages.find((message) => message.messageId === 'active-response-id')).toEqual({ + ...activeResponse, + content: aggregatedContent, + }); + unmount(); + }); + + it('replaces the current-run placeholder without erasing its loaded content', async () => { + const userMessage = { + messageId: 'server-user-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const preliminaryResponse = { + messageId: 'server-user-id_', + parentMessageId: 'server-user-id', + conversationId: CONV_ID, + text: '', + content: [{ type: ContentTypes.TEXT, text: { value: 'Already streaming' } }], + sender: 'Assistant', + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage, initialResponse: preliminaryResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([userMessage, preliminaryResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + await emitSync(getLastSSE(), [], 'assigned-response-id'); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'server-user-id', + 'assigned-response-id', + ]); + expect(updatedMessages[1]).toEqual({ + ...preliminaryResponse, + messageId: 'assigned-response-id', + }); + unmount(); + }); + + it('replaces the current-run user when sync assigns both durable IDs', async () => { + const preliminaryUser = { + messageId: 'client-user-id', + parentMessageId: 'previous-response-id', + conversationId: CONV_ID, + text: 'Hello', + isCreatedByUser: true, + } as TMessage; + const preliminaryResponse = { + messageId: 'client-user-id_', + parentMessageId: preliminaryUser.messageId, + conversationId: CONV_ID, + text: '', + content: [], + isCreatedByUser: false, + } as TMessage; + const submission = { + ...buildSubmission({ userMessage: preliminaryUser, initialResponse: preliminaryResponse }), + resumeStreamId: CONV_ID, + } as TSubmission & { resumeStreamId: string }; + const chatHelpers = buildChatHelpers(); + chatHelpers.getMessages.mockReturnValue([preliminaryUser, preliminaryResponse]); + + const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers)); + await act(async () => { + await Promise.resolve(); + }); + + const assignedUser = { + ...preliminaryUser, + messageId: 'assigned-user-id', + }; + await emitSync(getLastSSE(), [], 'assigned-response-id', undefined, assignedUser); + + const updatedMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0] as TMessage[]; + expect(updatedMessages.map((message) => message.messageId)).toEqual([ + 'assigned-user-id', + 'assigned-response-id', + ]); + expect(updatedMessages[1]?.parentMessageId).toBe('assigned-user-id'); + unmount(); + }); +}); diff --git a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx index 93fb20a0209..1af5beda979 100644 --- a/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx +++ b/client/src/hooks/SSE/__tests__/useResumeOnLoad.spec.tsx @@ -290,6 +290,81 @@ describe('useResumeOnLoad', () => { expect(attached?.resumeGenerationCreatedAt).toBe(4242); }); + it('restores an externally started regeneration after history refreshes', async () => { + const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); + const olderResponse = { + messageId: 'older-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Older response', + isCreatedByUser: false, + } as TMessage; + const newerResponse = { + messageId: 'newer-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Newer response', + isCreatedByUser: false, + } as TMessage; + const observedSubmissions: Array = []; + let messages = [rootUser]; + mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); + + const { rerender, queryClient } = renderUseResumeOnLoad({ + getMessages: () => messages, + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + await act(async () => { + await Promise.resolve(); + }); + + const invalidate = jest.spyOn(queryClient, 'invalidateQueries').mockResolvedValue(undefined); + messages = [rootUser, newerResponse, olderResponse]; + mockUseActiveJobs.mockReturnValue({ + data: { activeJobIds: [CONVERSATION_ID] }, + dataUpdatedAt: 2, + }); + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + createdAt: 4242, + streamId: CONVERSATION_ID, + resumeState: { + aggregatedContent: [{ type: ContentTypes.TEXT, text: 'regenerating' }], + responseMessageId: `${olderResponse.messageId}_`, + userMessage: { + messageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + }, + }, + }, + }); + rerender(); + await act(async () => { + await Promise.resolve(); + }); + + expect(invalidate).toHaveBeenCalledWith({ + queryKey: [QueryKeys.messages, CONVERSATION_ID], + }); + const attached = observedSubmissions[observedSubmissions.length - 1]; + expect(attached?.isRegenerate).toBe(true); + expect(attached?.initialResponse?.messageId).toBe(`${olderResponse.messageId}_`); + expect(attached?.messages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + expect(attached?.regenerateMessages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + }); + it('re-arms once per run rather than on every poll of the active list', async () => { mockUseStreamStatus.mockReturnValue(INACTIVE_STATUS); const { rerender } = renderUseResumeOnLoad({ messages: [] }); @@ -771,6 +846,54 @@ describe('useResumeOnLoad', () => { ]); }); + it('does not claim an older sibling when resume state omits the response ID', async () => { + const observedSubmissions: Array = []; + const userMessage = buildUserMessage(CONVERSATION_ID); + const olderSibling = { + messageId: 'older-sibling-response', + parentMessageId: userMessage.messageId, + conversationId: CONVERSATION_ID, + text: 'Older sibling', + isCreatedByUser: false, + } as TMessage; + + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + streamId: CONVERSATION_ID, + resumeState: { + runSteps: [], + aggregatedContent: [{ type: 'text', text: 'Active branch streaming' }], + conversationId: CONVERSATION_ID, + userMessage: { + messageId: userMessage.messageId, + parentMessageId: userMessage.parentMessageId, + conversationId: CONVERSATION_ID, + text: userMessage.text, + }, + }, + }, + }); + + renderUseResumeOnLoad({ + messages: [userMessage, olderSibling], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + + await act(async () => { + await Promise.resolve(); + }); + + const submission = observedSubmissions[observedSubmissions.length - 1]; + expect(submission?.initialResponse?.messageId).toBe(`${userMessage.messageId}_`); + expect((submission?.messages ?? []).map((message) => message.messageId)).toEqual([ + olderSibling.messageId, + ]); + }); + it('restores the branch that owns a pending OAuth resume user message', async () => { const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); const branchOneResponse = { @@ -834,13 +957,16 @@ describe('useResumeOnLoad', () => { expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(1); }); - it('restores the assistant sibling selected by a pending regenerate response', async () => { + it('restores the regenerate branch without claiming its older response', async () => { const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); const olderResponse = { messageId: 'older-response', parentMessageId: rootUser.messageId, conversationId: CONVERSATION_ID, text: 'Older response', + sender: 'Agent One', + model: 'gpt-5', + iconURL: 'https://example.com/agent-one.png', isCreatedByUser: false, } as TMessage; const newerResponse = { @@ -890,9 +1016,88 @@ describe('useResumeOnLoad', () => { expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(0); const submission = observedSubmissions[observedSubmissions.length - 1]; - expect(submission?.initialResponse?.messageId).toBe(olderResponse.messageId); + expect(submission?.initialResponse?.messageId).toBe(`${olderResponse.messageId}_`); + expect(submission?.initialResponse).toEqual( + expect.objectContaining({ + sender: olderResponse.sender, + model: olderResponse.model, + iconURL: olderResponse.iconURL, + }), + ); + expect(submission?.isRegenerate).toBe(true); expect((submission?.messages ?? []).map((message) => message.messageId)).toEqual([ + rootUser.messageId, newerResponse.messageId, + olderResponse.messageId, + ]); + expect((submission?.regenerateMessages ?? []).map((message) => message.messageId)).toEqual([ + rootUser.messageId, + newerResponse.messageId, + olderResponse.messageId, + ]); + }); + + it('preserves an exact-ID edited regeneration branch for early-abort rollback', async () => { + const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user'); + const editedResponse = { + messageId: 'edited-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Original response before the edit', + isCreatedByUser: false, + } as TMessage; + const siblingResponse = { + messageId: 'sibling-response', + parentMessageId: rootUser.messageId, + conversationId: CONVERSATION_ID, + text: 'Unrelated sibling', + isCreatedByUser: false, + } as TMessage; + const observedSubmissions: Array = []; + + mockUseStreamStatus.mockReturnValue({ + isSuccess: true, + isFetching: false, + data: { + active: true, + status: 'running', + streamId: CONVERSATION_ID, + resumeState: { + runSteps: [], + aggregatedContent: [], + responseMessageId: editedResponse.messageId, + isRegenerate: true, + conversationId: CONVERSATION_ID, + userMessage: { + messageId: rootUser.messageId, + parentMessageId: rootUser.parentMessageId, + conversationId: CONVERSATION_ID, + text: rootUser.text, + }, + }, + }, + }); + + renderUseResumeOnLoad({ + messages: [rootUser, siblingResponse, editedResponse], + onSubmission: (currentSubmission) => observedSubmissions.push(currentSubmission), + }); + + await act(async () => { + await Promise.resolve(); + }); + + const submission = observedSubmissions[observedSubmissions.length - 1]; + expect(submission?.isRegenerate).toBe(true); + expect(submission?.initialResponse?.messageId).toBe(editedResponse.messageId); + expect(submission?.messages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + siblingResponse.messageId, + ]); + expect(submission?.regenerateMessages?.map((message) => message.messageId)).toEqual([ + rootUser.messageId, + siblingResponse.messageId, + editedResponse.messageId, ]); }); diff --git a/client/src/hooks/SSE/useEventHandlers.ts b/client/src/hooks/SSE/useEventHandlers.ts index 4c5b6b0e5f5..570eafbcb57 100644 --- a/client/src/hooks/SSE/useEventHandlers.ts +++ b/client/src/hooks/SSE/useEventHandlers.ts @@ -185,6 +185,35 @@ export const getExistingConversationAbortMessages = ({ return [...sourceMessages]; }; +export const mergeErrorMessages = ({ + messages, + regenerateMessages, + userMessage, + errorMessage, + isRegenerate = false, +}: Pick & { + errorMessage: TMessage; +}): TMessage[] => { + if (isRegenerate) { + const finalMessages: TMessage[] = []; + let replaced = false; + for (const message of regenerateMessages ?? messages) { + if (message.messageId === errorMessage.messageId) { + finalMessages.push(errorMessage); + replaced = true; + } else { + finalMessages.push(message); + } + } + if (!replaced) { + finalMessages.push(errorMessage); + } + return finalMessages; + } + + return [...messages, userMessage, errorMessage]; +}; + export type EventHandlerParams = { isAddedRequest?: boolean; runIndex?: number; @@ -945,14 +974,14 @@ export default function useEventHandlers({ const errorHandler = useCallback( ({ data, submission }: { data?: TResData; submission: EventSubmission }) => { - const { messages, userMessage, initialResponse } = submission; + const { userMessage, initialResponse } = submission; setCompleted((prev) => new Set(prev.add(initialResponse.messageId))); const conversationId = userMessage.conversationId ?? submission.conversation?.conversationId ?? ''; const setErrorMessages = (convoId: string, errorMessage: TMessage) => { - const finalMessages: TMessage[] = [...messages, userMessage, errorMessage]; + const finalMessages = mergeErrorMessages({ ...submission, errorMessage }); setMessages(finalMessages); queryClient.setQueryData([QueryKeys.messages, convoId], finalMessages); }; diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 67ed8957181..613a86e48cd 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -583,18 +583,109 @@ const buildResumeEventSubmission = ( } as EventSubmission; }; +type ResumeMessageIndexes = { + userIndex: number; + responseIndex: number; + preliminaryUserIndex: number; + preliminaryResponseIndex: number; +}; + +const getResumeMessageIndexes = ( + messages: TMessage[], + userMessageId: string, + responseMessageId: string, + preliminaryUserMessageId?: string, + preliminaryResponseMessageId?: string, +): ResumeMessageIndexes => { + let userIndex = -1; + let responseIndex = -1; + let preliminaryUserIndex = -1; + let preliminaryResponseIndex = -1; + const hasPreliminaryResponse = preliminaryResponseMessageId?.endsWith('_') === true; + const eligiblePreliminaryUserId = + hasPreliminaryResponse && preliminaryUserMessageId && preliminaryUserMessageId !== userMessageId + ? preliminaryUserMessageId + : undefined; + const eligiblePreliminaryResponseId = + hasPreliminaryResponse && preliminaryResponseMessageId !== responseMessageId + ? preliminaryResponseMessageId + : undefined; + + for (let index = 0; index < messages.length; index++) { + const messageId = messages[index]?.messageId; + if (userIndex < 0 && messageId === userMessageId) { + userIndex = index; + } + if (responseIndex < 0 && messageId === responseMessageId) { + responseIndex = index; + } + if ( + preliminaryUserIndex < 0 && + eligiblePreliminaryUserId && + messageId === eligiblePreliminaryUserId + ) { + preliminaryUserIndex = index; + } + if ( + preliminaryResponseIndex < 0 && + eligiblePreliminaryResponseId && + messageId === eligiblePreliminaryResponseId + ) { + preliminaryResponseIndex = index; + } + } + + return { userIndex, responseIndex, preliminaryUserIndex, preliminaryResponseIndex }; +}; + const mergeResumeMessages = ( messages: TMessage[], userMessage: TMessage, responseMessage: TMessage, + indexes: ResumeMessageIndexes, ): TMessage[] => { const nextMessages = [...messages]; - const userIndex = nextMessages.findIndex( - (message) => message.messageId === userMessage.messageId, - ); - const responseIndex = nextMessages.findIndex( - (message) => message.messageId === responseMessage.messageId, - ); + let { userIndex, responseIndex, preliminaryResponseIndex } = indexes; + const { preliminaryUserIndex } = indexes; + + if (preliminaryUserIndex >= 0) { + if (userIndex >= 0) { + nextMessages.splice(preliminaryUserIndex, 1); + if (userIndex > preliminaryUserIndex) { + userIndex -= 1; + } + if (responseIndex > preliminaryUserIndex) { + responseIndex -= 1; + } + if (preliminaryResponseIndex > preliminaryUserIndex) { + preliminaryResponseIndex -= 1; + } + } else { + nextMessages[preliminaryUserIndex] = { + ...nextMessages[preliminaryUserIndex], + ...userMessage, + }; + userIndex = preliminaryUserIndex; + } + } + + if (preliminaryResponseIndex >= 0) { + if (responseIndex >= 0) { + nextMessages.splice(preliminaryResponseIndex, 1); + if (userIndex > preliminaryResponseIndex) { + userIndex -= 1; + } + if (responseIndex > preliminaryResponseIndex) { + responseIndex -= 1; + } + } else { + nextMessages[preliminaryResponseIndex] = { + ...nextMessages[preliminaryResponseIndex], + ...responseMessage, + }; + responseIndex = preliminaryResponseIndex; + } + } if (userIndex >= 0) { nextMessages[userIndex] = { ...nextMessages[userIndex], ...userMessage }; @@ -609,9 +700,7 @@ const mergeResumeMessages = ( } if (userIndex >= 0) { - const insertAt = userIndex + 1; - nextMessages.splice(insertAt, 0, responseMessage); - return nextMessages; + return [...nextMessages, responseMessage]; } if (responseIndex >= 0) { @@ -1848,6 +1937,16 @@ export default function useResumableSSE( const runId = v4(); setActiveRunId(runId); + /** Keep the current run's preliminary id long enough to replace that optimistic + * row in place if this snapshot assigns its durable response id. */ + const preliminaryResponseMessageId = currentSubmission.initialResponse?.messageId; + const currentUserMessageId = currentSubmission.userMessage?.messageId; + const preliminaryUserMessageId = + currentUserMessageId && + (currentSubmission.initialResponse?.parentMessageId === currentUserMessageId || + preliminaryResponseMessageId === `${currentUserMessageId}_`) + ? currentUserMessageId + : undefined; const resumeSubmission = buildResumeEventSubmission( currentSubmission, userMessage, @@ -1893,28 +1992,23 @@ export default function useResumableSSE( if (data.resumeState?.aggregatedContent && userMessage?.messageId) { const messages = getMessages() ?? []; const userMsgId = userMessage.messageId; - const serverResponseId = data.resumeState.responseMessageId; const hasResumedContent = data.resumeState.aggregatedContent.length > 0; - - let responseIdx = -1; - /** Only an id match proves the row belongs to THIS generation; the parent-based - * fallback below can land on a prior sibling (e.g. the answer being regenerated). */ - let matchedByResponseId = false; - if (serverResponseId) { - responseIdx = messages.findIndex((m) => m.messageId === serverResponseId); - matchedByResponseId = responseIdx >= 0; - } - if (responseIdx < 0) { - responseIdx = messages.findIndex( - (m) => - !m.isCreatedByUser && - (m.messageId === `${userMsgId}_` || m.parentMessageId === userMsgId), - ); - } + const responseId = resumeSubmission.initialResponse.messageId; + const messageIndexes = getResumeMessageIndexes( + messages, + userMsgId, + responseId, + preliminaryUserMessageId, + preliminaryResponseMessageId, + ); + const responseIdx = + messageIndexes.responseIndex >= 0 + ? messageIndexes.responseIndex + : messageIndexes.preliminaryResponseIndex; logger.log('ResumableSSE', 'SYNC update', { userMsgId, - serverResponseId, + responseId, responseIdx, foundMessageId: responseIdx >= 0 ? messages[responseIdx]?.messageId : null, messagesCount: messages.length, @@ -1924,15 +2018,11 @@ export default function useResumableSSE( if (responseIdx >= 0) { const oldContent = messages[responseIdx]?.content; /** An EMPTY resume snapshot is not authoritative over content we already loaded - * for the SAME response: assigning it would erase that content and leave a bare - * cursor. Restricted to an id match — preserving a fallback-matched row would - * make a regenerated run append to the answer it is replacing — and to a row - * that actually HAS parts, so the array is never swapped for `undefined`. */ + * for the SAME generation-owned response: assigning it would erase that content + * and leave a bare cursor. Require an existing content array so it is never + * swapped for `undefined`. */ const preserveLoadedContent = - !hasResumedContent && - matchedByResponseId && - Array.isArray(oldContent) && - oldContent.length > 0; + !hasResumedContent && Array.isArray(oldContent) && oldContent.length > 0; /** * Replacing the response with `aggregatedContent` drops the * prefix an edited resubmission had retained: the snapshot is @@ -1947,21 +2037,32 @@ export default function useResumableSSE( editPrefixClearedRef.current = true; } const responseMessage = { + ...resumeSubmission.initialResponse, ...messages[responseIdx], + messageId: responseId, + parentMessageId: userMsgId, content: preserveLoadedContent ? oldContent : data.resumeState.aggregatedContent, + sender: messages[responseIdx]?.sender ?? resumeSubmission.initialResponse.sender, iconURL: preferDefinedString( messages[responseIdx]?.iconURL, - data.resumeState.iconURL, + resumeSubmission.initialResponse.iconURL ?? data.resumeState.iconURL, + ), + model: preferDefinedString( + messages[responseIdx]?.model, + resumeSubmission.initialResponse.model ?? data.resumeState.model, ), - model: preferDefinedString(messages[responseIdx]?.model, data.resumeState.model), } as TMessage; - const updated = mergeResumeMessages(messages, userMessage, responseMessage); + const updated = mergeResumeMessages( + messages, + userMessage, + responseMessage, + messageIndexes, + ); logger.log('ResumableSSE', 'SYNC updating message', { messageId: responseMessage.messageId, oldContentLength: Array.isArray(oldContent) ? oldContent.length : 0, newContentLength: data.resumeState.aggregatedContent?.length, preservedExistingContent: preserveLoadedContent, - matchedByResponseId, }); setMessages(updated); resetContentHandler(); @@ -1975,18 +2076,14 @@ export default function useResumableSSE( * only in the matched branch left this path adding an offset * to indices that were already absolute. */ editPrefixClearedRef.current = true; - const responseId = serverResponseId ?? `${userMsgId}_`; const newMessage = { + ...resumeSubmission.initialResponse, messageId: responseId, parentMessageId: userMsgId, - conversationId: currentSubmission.conversation?.conversationId ?? '', - text: '', content: data.resumeState.aggregatedContent, isCreatedByUser: false, - iconURL: data.resumeState.iconURL, - model: data.resumeState.model, } as TMessage; - setMessages(mergeResumeMessages(messages, userMessage, newMessage)); + setMessages(mergeResumeMessages(messages, userMessage, newMessage, messageIndexes)); resetContentHandler(); syncStepMessage(newMessage); } diff --git a/client/src/hooks/SSE/useResumeOnLoad.ts b/client/src/hooks/SSE/useResumeOnLoad.ts index 9ff950cc735..f0b0f7653fe 100644 --- a/client/src/hooks/SSE/useResumeOnLoad.ts +++ b/client/src/hooks/SSE/useResumeOnLoad.ts @@ -128,15 +128,27 @@ function buildSubmissionFromResumeState( (m) => m.isCreatedByUser && m.messageId === userMessageData?.messageId, ); - // Try to find existing response message in the messages array (from database). - // Regeneration can expose the in-flight placeholder id with trailing underscores - // while the persisted sibling uses the unpadded id. Prefer both exact identities - // before falling back to the shared parent, where several branch siblings can match. + // A trailing underscore distinguishes an in-flight regeneration from the persisted + // response it replaces. Only the exact response id proves generation ownership. + const existingResponseMessage = messages.find( + (m) => !m.isCreatedByUser && m.messageId === responseMessageId, + ); + // The persisted row may seed display metadata, but never identity or deduplication. const unpaddedResponseMessageId = responseMessageId.replace(/_+$/, ''); - const existingResponseMessage = - messages.find((m) => !m.isCreatedByUser && m.messageId === responseMessageId) ?? - messages.find((m) => !m.isCreatedByUser && m.messageId === unpaddedResponseMessageId) ?? - messages.find((m) => !m.isCreatedByUser && m.parentMessageId === userMessageData?.messageId); + const persistedRegenerationResponse = + unpaddedResponseMessageId !== responseMessageId + ? messages.find((m) => !m.isCreatedByUser && m.messageId === unpaddedResponseMessageId) + : undefined; + const responseMetadataMessage = existingResponseMessage ?? persistedRegenerationResponse; + const isRegenerateResume = + resumeState.isRegenerate === true || persistedRegenerationResponse != null; + let regenerateMessages: TMessage[] | undefined; + if (isRegenerateResume) { + regenerateMessages = + unpaddedResponseMessageId === responseMessageId + ? [...messages] + : messages.filter((message) => message.messageId !== responseMessageId); + } // Create or use existing user message const userMessage: TMessage = @@ -169,9 +181,9 @@ function buildSubmissionFromResumeState( content: (resumeState.aggregatedContent as TMessage['content']) ?? [], isCreatedByUser: false, role: 'assistant', - sender: existingResponseMessage?.sender ?? resumeState.sender, - model: preferDefinedString(existingResponseMessage?.model, resumeState.model), - iconURL: preferDefinedString(existingResponseMessage?.iconURL, resumeState.iconURL), + sender: responseMetadataMessage?.sender ?? resumeState.sender, + model: preferDefinedString(responseMetadataMessage?.model, resumeState.model), + iconURL: preferDefinedString(responseMetadataMessage?.iconURL, resumeState.iconURL), } as TMessage; // Re-paused turn: seed the approval / ask-user controls straight onto the @@ -186,17 +198,13 @@ function buildSubmissionFromResumeState( endpoint: null, } as TConversation; - // On reload, `messages` is the full DB array, which already holds the paused user - // row and the partial (unfinished) assistant row under the same ids that - // `userMessage` / `initialResponse` (and the resume final event's request/response - // messages) re-supply. Strip them so createdHandler/finalHandler — which build - // `[...messages, requestMessage, responseMessage]` — don't append a duplicate pair. - const pausedResponseIdUnpadded = initialResponse.messageId.replace(/_+$/, ''); + // Non-regenerate resumes strip the persisted request/response pair before handlers + // re-supply it. A regeneration keeps the original branch for early-abort rollback; + // explicit resume metadata covers edited regenerations that reuse the exact response id. const dedupedMessages = messages.filter( (m) => - m.messageId !== userMessage.messageId && m.messageId !== initialResponse.messageId && - m.messageId !== pausedResponseIdUnpadded, + (isRegenerateResume || m.messageId !== userMessage.messageId), ); return { @@ -204,7 +212,8 @@ function buildSubmissionFromResumeState( userMessage, initialResponse, conversation, - isRegenerate: false, + isRegenerate: isRegenerateResume, + ...(regenerateMessages && { regenerateMessages }), isTemporary: false, endpointOption: {}, // Signal to useResumableSSE to subscribe to existing stream instead of starting new diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 3a89dfb2d98..25febf31ab5 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -2528,6 +2528,7 @@ class GenerationJobManagerClass { generationProtocolVersion: jobData.generationProtocolVersion, userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, + isRegenerate: jobData.isRegenerate, sender: jobData.sender, endpoint: jobData.endpoint, iconURL: jobData.iconURL, @@ -6928,6 +6929,7 @@ class GenerationJobManagerClass { aggregatedContent, userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, + isRegenerate: jobData.isRegenerate, conversationId: jobData.conversationId, sender: jobData.sender, iconURL: jobData.iconURL, diff --git a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts index d252af0d828..dee6a9e117d 100644 --- a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts +++ b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts @@ -67,6 +67,22 @@ describe('GenerationJobManager resume replay events', () => { manager = undefined; }); + test('projects regeneration ownership into resume state', async () => { + manager = createInMemoryManager(); + const streamId = `regenerate-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId, { + initialMetadata: { + responseMessageId: 'edited-response', + isRegenerate: true, + }, + }); + + await expect(manager.getResumeState(streamId)).resolves.toMatchObject({ + responseMessageId: 'edited-response', + isRegenerate: true, + }); + }); + test('includes OAuth run step and delta replay events in resume state', async () => { manager = createInMemoryManager(); const streamId = `oauth-delta-resume-${Date.now()}`; diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index 7298dd40265..e114c8d3c5a 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -99,6 +99,7 @@ describe('GenerationJobManager startup telemetry', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + isRegenerate: true, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', @@ -129,6 +130,7 @@ describe('GenerationJobManager startup telemetry', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + isRegenerate: true, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index ea46b4a4e1b..b73a4d7af98 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -4547,6 +4547,7 @@ export class RedisJobStore implements IJobStoreV2 { recoveredSteerId: data.recoveredSteerId || undefined, userMessage: data.userMessage ? JSON.parse(data.userMessage) : undefined, responseMessageId: data.responseMessageId || undefined, + isRegenerate: data.isRegenerate != null ? data.isRegenerate === '1' : undefined, createdEventEmitted: data.createdEventEmitted === '1', sender: data.sender || undefined, syncSent: data.syncSent === '1', diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index a03e28298d1..e086b517600 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -75,6 +75,9 @@ export interface SerializableJobData { /** Response message ID for reconnection */ responseMessageId?: string; + /** Whether this generation replaces an existing assistant branch. */ + isRegenerate?: boolean; + /** * Whether this run has activity labels enabled (per-endpoint * `activityLabel: true`). Set once at run start so the resume path can @@ -294,6 +297,7 @@ export type JobMetadataPatch = Partial< Pick< SerializableJobData, | 'responseMessageId' + | 'isRegenerate' | 'sender' | 'conversationId' | 'userMessage' diff --git a/packages/api/src/stream/metadata.ts b/packages/api/src/stream/metadata.ts index b334412267f..da8c108b34a 100644 --- a/packages/api/src/stream/metadata.ts +++ b/packages/api/src/stream/metadata.ts @@ -6,6 +6,9 @@ export function sanitizeJobMetadata(metadata: Partial): J if (metadata.responseMessageId) { patch.responseMessageId = metadata.responseMessageId; } + if (metadata.isRegenerate !== undefined) { + patch.isRegenerate = metadata.isRegenerate; + } if (metadata.sender) { patch.sender = metadata.sender; } diff --git a/packages/api/src/types/stream.ts b/packages/api/src/types/stream.ts index b6727ebfbfa..fca5d8e9438 100644 --- a/packages/api/src/types/stream.ts +++ b/packages/api/src/types/stream.ts @@ -17,6 +17,8 @@ export interface GenerationJobMetadata { userMessage?: Agents.UserMessageMeta; /** Response message ID for tracking */ responseMessageId?: string; + /** Whether this generation replaces an existing assistant branch. */ + isRegenerate?: boolean; /** Sender label for the response (e.g., "GPT-4.1", "Claude") */ sender?: string; /** Endpoint identifier for abort handling */ diff --git a/packages/data-provider/src/types/agents.ts b/packages/data-provider/src/types/agents.ts index ae8499c61f7..35d2e39c9f8 100644 --- a/packages/data-provider/src/types/agents.ts +++ b/packages/data-provider/src/types/agents.ts @@ -263,6 +263,8 @@ export namespace Agents { aggregatedContent?: MessageContentComplex[]; userMessage?: UserMessageMeta; responseMessageId?: string; + /** True when the live generation replaces an existing assistant branch. */ + isRegenerate?: boolean; conversationId?: string; sender?: string; iconURL?: string; From f02ce63d571c7b250b6bef6a51f173adad4fa1e9 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:30:29 -0700 Subject: [PATCH 3/8] =?UTF-8?q?=E2=9C=82=EF=B8=8F=20fix:=20Strip=20Redunda?= =?UTF-8?q?nt=20Server-Name=20Prefixes=20from=20MCP=20Tool=20Keys=20(#1473?= =?UTF-8?q?2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ✂️ fix: Strip Redundant Server-Name Prefixes from MCP Tool Keys MCP servers that prefix every tool with their own name produce model-facing keys that embed the server twice once the _mcp_ suffix is appended, pushing long tool names past provider 64-character function-name limits. Tool keys now drop a leading _ prefix (case-insensitive, skipped when a sibling tool already owns the stripped name). The original upstream name is recorded as serverToolName on the cached definition and is always what tool calls send to the server, and runtime lookups also try the stripped spelling of persisted pre-strip keys so existing agents keep resolving. * 🩹 fix: Keep Stripped MCP Tool Keys Provider-Safe and Collision-Free Assistant writers submit catalog entries verbatim, so the internal serverToolName mapping is now removed from provider-facing definitions before they reach create/update payloads. Prefix stripping is collision-guarded over the resulting name set rather than raw siblings only, which also covers case-variant prefixed pairs under the case-insensitive match. Assistant payload healing now rewrites a pre-strip persisted key to the stripped catalog key when that key actually exists in the loaded definitions, and legacy agent references keep their persisted spelling as the runtime instance name so per-tool options stay applied while the upstream call still uses the matched entry's raw name. * 🧷 fix: Harden Stripped MCP Tool Keys Against Heal, Collision, and Cache Edges The pre-strip heal now resolves the key boundary against both raw and normalized server spellings, mapping back to the raw name for the shadow and membership guards, so keys persisted after server-name normalization heal too. Collision detection iterates to a fixpoint so a fallback to a raw name cannot silently collide with another sibling's stripped result, and a stripped remainder equal to a synthetic marker (wildcard or server pin) is never produced. MCP catalog cache slices are versioned so replicas that predate serverToolName never read stripped entries during a rolling deploy; stale slices expire on their own. * 🔎 fix: Resolve Pre-Strip Keys in Event-Driven Definitions and Reinspect Persisted Catalogs The event-driven definitions loader now tries the stripped spelling of a persisted key when the exact lookup misses, keeping the persisted name so it matches the runtime instance, which stops legacy agents from failing initialization with expected tools unavailable. The registry storage schema version is bumped so followers rebuild persisted toolFunctions instead of republishing pre-strip definitions into the versioned catalog namespace. The assistants heal also fails closed when a normalized-suffix reference lands on a contested server-name slot, since rewriting persisted data must not bind an ambiguous reference to the tie-break winner. * 🛰️ fix: Reserve the Synthetic OAuth Name and Heal User-Owned Server Keys A stripped remainder equal to oauth would make the client stream handlers treat a real tool call as a synthetic authentication prompt, so it joins the reserved remainders alongside the wildcard and pin markers. The assistants heal now audits the FULL accessible server set on every run instead of operator config names only, since assistants reference user-owned servers whose catalogs the definitions loader already resolves; an unavailable audit still skips healing entirely. * 🧬 fix: Verify Upstream Identity for Legacy Keys and Reserve Sibling Raw Names Stripped results now reserve every sibling's raw name even when that sibling itself strips, so a stripped key can never shadow another tool's pre-rollout persisted references within the same snapshot. Every legacy fallback (runtime lookup, event-driven definitions, assistants heal) accepts a stripped-spelling match only when the entry's recorded serverToolName proves the same upstream tool, so a stale key for a removed tool degrades to unavailable instead of calling a different sibling. To keep that identity visible to the heal, assistant tool definitions retain serverToolName and the controllers sanitize entries through toProviderToolDefinition at the provider submission boundary instead. The agent editor migrates pre-strip persisted ids the same identity-verified way, with the upstream name exposed on the MCP tools payload. * 🧭 fix: Heal Wildcard Tool Options and Reserve the OAuth Namespace Wildcard-expanded catalogs rename stripped tools without any agent.tools entry to preserve the spelling, so buildToolClassification now aliases persisted pre-strip tool_options keys onto the current instance names in place, identity-gated on the definition's recorded upstream name and never overriding an explicit entry. Both loading modes flow through it: instances carry mcpServerToolName from createToolInstance and event-driven definitions thread serverToolName from the catalog. stripServerNamePrefix also reserves the entire oauth namespace rather than the exact name, since the client stream handlers classify every oauth-prefixed key as a synthetic authentication call. * 🛡️ fix: Derive the Full Reserved Namespace and Heal Approval Policies The reservation guard now covers every namespace consumers classify by prefix: the wildcard and server-pin markers alongside oauth, plus the server-scoped mcp_ pluginKey namespace that pre-strip keys could never enter. Stripping also never produces a key whose isActionTool classification differs from the raw key's, since a server whose normalized name contains _action_ would otherwise see a real MCP tool routed down the OpenAPI action path past MCP authorization. Admin toolApproval globs written against upstream tool naming keep applying: pattern lists are healed at run wiring with the current names of tools whose pre-strip spelling matches, list-level so deny, ask, and allow precedence is unchanged and a non-matching deny can no longer fail open. The MCP tools wire type also declares serverToolName end to end. * 🪪 fix: Alias Both Key Spellings for Approval Policies and Hook Matchers Identity aliases are now collected once at tool classification, in both directions: a stripped instance aliases its pre-strip spelling and a legacy-named instance aliases its current catalog spelling, with the current name recorded on legacy matches by the runtime lookup and the event-driven definitions loader alike. The aliases ride the agent config through both loading modes, so approval pattern healing applies to deny rules written against either spelling, closing the bypass where a rule targeting the current name missed an unedited agent's legacy instance. Programmatic approval hook matchers get the same treatment: each hook is additionally registered under an anchored exact-name pattern for tools whose other spelling its regex matches, keeping the admin's matcher semantics intact while argument, user, and tenant specific deny or ask decisions keep executing for renamed tools. * 🔁 fix: Alias Tool Options in Both Spelling Directions Options aliasing now consumes the same bidirectional alias pairs as policy healing and hook matchers, so options the editor migrated to the current catalog spelling still reach a legacy-named instance retained by an unedited agent.tools entry. The previous serverToolName-only derivation skipped exactly that case since the legacy key equals the instance name there. * 🤝 fix: Reserve the Agent Handoff Namespace Before Stripping The client renders any lc_transfer_to_ prefixed call as an agent handoff and the background and intent passes exclude such names, so a stripped remainder inside that namespace would misclassify a real upstream tool. It joins the mcp_ pluginKey namespace as a bare-prefix reservation, which pre-strip keys could never enter. * ⚡ fix: Reuse the Loader's Server Snapshot and Index the Editor Catalog getAssistantToolDefinitions now returns the accessible-server snapshot from the same merged registry read that resolved the catalogs, and the heal consumes it instead of repeating the app-config and registry round trips on the assistant write path; without a snapshot the heal still fetches and fails closed as before. The agent editor's id migration uses a memoized tool_id map, so the per-key form heal does constant-time lookups instead of scanning the catalog per option. --------- Co-authored-by: Danny Avila --- api/server/controllers/assistants/v1.js | 36 +++- api/server/controllers/assistants/v2.js | 35 +++- api/server/controllers/mcp.js | 4 + .../__tests__/getCachedTools.lock.spec.js | 4 +- .../Config/__tests__/getCachedTools.spec.js | 6 +- api/server/services/MCP.js | 192 ++++++++++++++---- api/server/services/MCP.spec.js | 140 +++++++++++++ api/server/services/ToolService.js | 9 +- api/server/services/__tests__/MCP.spec.js | 147 +++++++++++++- .../Tools/ItemDialog/sections/McpSection.tsx | 41 +++- packages/api/src/agents/hitl/policy.spec.ts | 87 ++++++++ packages/api/src/agents/hitl/policy.ts | 88 ++++++++ packages/api/src/agents/hitl/runtime.ts | 23 ++- packages/api/src/agents/initialize.ts | 7 + packages/api/src/agents/run.ts | 24 ++- packages/api/src/mcp/assistants.spec.ts | 52 ++++- packages/api/src/mcp/assistants.ts | 45 +++- packages/api/src/mcp/catalog/store.ts | 15 +- .../src/mcp/registry/MCPServerInspector.ts | 10 +- .../src/mcp/registry/MCPServersInitializer.ts | 8 +- .../__tests__/MCPServerInspector.test.ts | 27 +++ packages/api/src/mcp/tools.spec.ts | 43 ++++ packages/api/src/mcp/tools.ts | 17 +- packages/api/src/mcp/types/index.ts | 3 + packages/api/src/mcp/utils.ts | 2 + packages/api/src/tools/classification.spec.ts | 107 ++++++++++ packages/api/src/tools/classification.ts | 99 ++++++++- packages/api/src/tools/definitions.spec.ts | 66 ++++++ packages/api/src/tools/definitions.ts | 53 ++++- packages/data-provider/src/config.ts | 115 +++++++++++ packages/data-provider/src/schemas.ts | 3 + .../data-provider/src/splitMCPToolKey.spec.ts | 126 ++++++++++++ packages/data-provider/src/types/queries.ts | 3 + 33 files changed, 1537 insertions(+), 100 deletions(-) diff --git a/api/server/controllers/assistants/v1.js b/api/server/controllers/assistants/v1.js index 926ab7db4dc..b27e5530b1d 100644 --- a/api/server/controllers/assistants/v1.js +++ b/api/server/controllers/assistants/v1.js @@ -7,7 +7,11 @@ const validateAuthor = require('~/server/middleware/assistants/validateAuthor'); const { getStrategyFunctions } = require('~/server/services/Files/strategies'); const { deleteAssistantActions } = require('~/server/services/ActionService'); const { getOpenAIClient, fetchAssistants } = require('./helpers'); -const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP'); +const { + healMcpToolNames, + getAssistantToolDefinitions, + toProviderToolDefinition, +} = require('~/server/services/MCP'); const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools'); /** @@ -30,8 +34,16 @@ const createAssistant = async (req, res) => { delete assistantData.conversation_starters; delete assistantData.append_current_datetime; - const toolDefinitions = await getAssistantToolDefinitions({ req, tools }); - const healedTools = await healMcpToolNames({ req, tools, toolDefinitions }); + const { toolDefinitions, accessibleServerNames } = await getAssistantToolDefinitions({ + req, + tools, + }); + const healedTools = await healMcpToolNames({ + req, + tools, + toolDefinitions, + accessibleServerNames, + }); assistantData.tools = healedTools .map((tool) => { @@ -59,7 +71,8 @@ const createAssistant = async (req, res) => { return toolDef; }) .filter((tool) => tool) - .flat(); + .flat() + .map(toProviderToolDefinition); let azureModelIdentifier = null; if (openai.locals?.azureOptions) { @@ -145,8 +158,16 @@ const patchAssistant = async (req, res) => { ...updateData } = req.body; - const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools }); - const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions }); + const { toolDefinitions, accessibleServerNames } = await getAssistantToolDefinitions({ + req, + tools: updateData.tools, + }); + const healedTools = await healMcpToolNames({ + req, + tools: updateData.tools, + toolDefinitions, + accessibleServerNames, + }); updateData.tools = healedTools .map((tool) => { @@ -174,7 +195,8 @@ const patchAssistant = async (req, res) => { return toolDef; }) .filter((tool) => tool) - .flat(); + .flat() + .map(toProviderToolDefinition); if (openai.locals?.azureOptions && updateData.model) { updateData.model = openai.locals.azureOptions.azureOpenAIApiDeploymentName; diff --git a/api/server/controllers/assistants/v2.js b/api/server/controllers/assistants/v2.js index a436ed611db..ec0a3f9309d 100644 --- a/api/server/controllers/assistants/v2.js +++ b/api/server/controllers/assistants/v2.js @@ -2,7 +2,11 @@ const { logger } = require('@librechat/data-schemas'); const { ToolCallTypes } = require('librechat-data-provider'); const validateAuthor = require('~/server/middleware/assistants/validateAuthor'); const { validateAndUpdateTool } = require('~/server/services/ActionService'); -const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP'); +const { + healMcpToolNames, + getAssistantToolDefinitions, + toProviderToolDefinition, +} = require('~/server/services/MCP'); const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools'); const { updateAssistantDoc } = require('~/models'); const { getOpenAIClient } = require('./helpers'); @@ -28,8 +32,16 @@ const createAssistant = async (req, res) => { delete assistantData.conversation_starters; delete assistantData.append_current_datetime; - const toolDefinitions = await getAssistantToolDefinitions({ req, tools }); - const healedTools = await healMcpToolNames({ req, tools, toolDefinitions }); + const { toolDefinitions, accessibleServerNames } = await getAssistantToolDefinitions({ + req, + tools, + }); + const healedTools = await healMcpToolNames({ + req, + tools, + toolDefinitions, + accessibleServerNames, + }); assistantData.tools = healedTools .map((tool) => { @@ -57,7 +69,8 @@ const createAssistant = async (req, res) => { return toolDef; }) .filter((tool) => tool) - .flat(); + .flat() + .map(toProviderToolDefinition); let azureModelIdentifier = null; if (openai.locals?.azureOptions) { @@ -134,8 +147,16 @@ const updateAssistant = async ({ req, openai, assistant_id, updateData }) => { } let hasFileSearch = false; - const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools }); - const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions }); + const { toolDefinitions, accessibleServerNames } = await getAssistantToolDefinitions({ + req, + tools: updateData.tools, + }); + const healedTools = await healMcpToolNames({ + req, + tools: updateData.tools, + toolDefinitions, + accessibleServerNames, + }); for (const tool of healedTools) { /** Agents-runtime-only tools (e.g. ask_user_question) cannot execute on * the assistants runtime — drop them even when posted directly, since @@ -201,7 +222,7 @@ const updateAssistant = async ({ req, openai, assistant_id, updateData }) => { }; } - updateData.tools = tools; + updateData.tools = tools.map(toProviderToolDefinition); if (openai.locals?.azureOptions && updateData.model) { updateData.model = openai.locals.azureOptions.azureOpenAIApiDeploymentName; diff --git a/api/server/controllers/mcp.js b/api/server/controllers/mcp.js index 1eb1b6eb3d0..19ea8d90b56 100644 --- a/api/server/controllers/mcp.js +++ b/api/server/controllers/mcp.js @@ -302,6 +302,10 @@ const getMCPTools = async (req, res) => { name: toolName, pluginKey: toolKey, description: toolData.function.description || '', + /** Upstream identity for keys that stripped a redundant + * server-name prefix — the agent editor migrates legacy + * persisted ids only when this proves the same tool. */ + ...(toolData.serverToolName != null && { serverToolName: toolData.serverToolName }), }); } } diff --git a/api/server/services/Config/__tests__/getCachedTools.lock.spec.js b/api/server/services/Config/__tests__/getCachedTools.lock.spec.js index 6302d20cadc..50e99ecfd0e 100644 --- a/api/server/services/Config/__tests__/getCachedTools.lock.spec.js +++ b/api/server/services/Config/__tests__/getCachedTools.lock.spec.js @@ -248,7 +248,7 @@ describe('global tool cache write lock', () => { expect.objectContaining({ keys: [ `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`, - `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`, + `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:v2:config-current`, ], arguments: [ 'generation-current', @@ -342,7 +342,7 @@ describe('global tool cache write lock', () => { `tools:mcp:write-fence:{user-1:server-1}`, `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-legacy-fence:{user-1:server-1}`, `${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`, - `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`, + `${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:v2:config-current`, ], }), ); diff --git a/api/server/services/Config/__tests__/getCachedTools.spec.js b/api/server/services/Config/__tests__/getCachedTools.spec.js index 6d3947392f1..dd5e231f023 100644 --- a/api/server/services/Config/__tests__/getCachedTools.spec.js +++ b/api/server/services/Config/__tests__/getCachedTools.spec.js @@ -29,10 +29,10 @@ describe('MCP tool cache', () => { it('uses collision-safe configuration-addressed keys', () => { expect(ToolCacheKeys.MCP_APP_SERVER('server:name', 'config/a')).toBe( - 'tools:mcp:app:server%3Aname:config%2Fa', + 'tools:mcp:app:v2:server%3Aname:config%2Fa', ); expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).toBe( - 'tools:mcp:user:{tenant%3Auser:server%3Aname}:config%2Fa', + 'tools:mcp:user:{tenant%3Auser:server%3Aname}:v2:config%2Fa', ); expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).not.toBe( ToolCacheKeys.MCP_SERVER('tenant', 'user:server:name', 'config/a'), @@ -49,7 +49,7 @@ describe('MCP tool cache', () => { }); it('keeps the legacy user key available for non-generation callers', () => { - expect(ToolCacheKeys.MCP_SERVER('user123', 'github')).toBe('tools:mcp:user123:github'); + expect(ToolCacheKeys.MCP_SERVER('user123', 'github')).toBe('tools:mcp:v2:user123:github'); }); it('gets and sets static global tools without touching MCP slices', async () => { diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 0c351996faf..15bbecc025a 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -10,9 +10,12 @@ const { splitMCPToolKey, normalizeServerName, normalizeMCPToolKey, + stripServerNamePrefix, + stripServerNamePrefixes, buildServerNameAliases, findShadowedServerNames, getAssistantToolDefinitions: loadAssistantToolDefinitions, + toProviderToolDefinition, resolveMCPServerContext, normalizeJsonSchema, GenerationJobManager, @@ -216,8 +219,10 @@ async function resolveMcpServerContext(req) { */ /** * Names of every MCP server the user can reach (operator config + user DB), - * for the legacy-key heal's collision detection in `initializeAgent`. Only - * consulted when a configured server name needs normalization. + * for legacy-key healing: collision detection in `initializeAgent` (consulted + * when a configured server name needs normalization) and the assistants heal + * in `healMcpToolNames` (always, since assistants reference user-owned + * servers too). * @param {string} [userId] * @param {string} [role] * @returns {Promise} @@ -251,7 +256,7 @@ async function getAccessibleMcpServerNames(userId, role) { * @param {Record} params.toolDefinitions * @returns {Promise>} */ -async function healMcpToolNames({ req, tools, toolDefinitions }) { +async function healMcpToolNames({ req, tools, toolDefinitions, accessibleServerNames }) { const list = tools ?? []; const needsHeal = list.some( (tool) => @@ -262,21 +267,36 @@ async function healMcpToolNames({ req, tools, toolDefinitions }) { if (!needsHeal) { return list; } - const rawServerNames = await resolveMcpConfigNames(req); /** Cross-tier shadowing (DB `foo` vs operator `foo!`) is invisible to * operator names alone — the shadow set must come from the FULL - * accessible audit. Every rewrite candidate here is normalization- - * sensitive by construction, so an incomplete audit skips healing - * entirely (the raw key stays raw and fails closed). */ - const audit = await resolveCollisionAuditNames({ - rawServerNames, - userId: req.user?.id, - role: req.user?.role, - }); - if (!audit.complete) { - return list; + * accessible audit: assistants reference user-owned servers too (the + * definitions loader resolves them), so their pre-strip keys must heal + * against the same catalog. Callers holding the loader's snapshot pass + * it to avoid repeating the app-config and registry reads on the write + * path; without one, the audit is fetched here, and when it cannot + * complete healing is skipped entirely (the raw key stays raw and fails + * closed). */ + let auditNames = accessibleServerNames; + if (auditNames == null) { + const rawServerNames = await resolveMcpConfigNames(req); + try { + const accessible = await getAccessibleMcpServerNames(req.user?.id, req.user?.role); + auditNames = [...new Set([...accessible, ...rawServerNames])]; + } catch (error) { + logger.warn( + '[healMcpToolNames] Accessible-server audit unavailable; skipping legacy-key healing:', + error, + ); + return list; + } } - const shadowed = findShadowedServerNames(audit.names); + const shadowed = findShadowedServerNames(auditNames); + /** A pre-strip key persisted AFTER server-name normalization carries the + * NORMALIZED suffix, which the raw config names cannot match — the + * boundary must resolve against both spellings and map back to the raw + * name for the shadow and membership guards. */ + const serverNameAliases = buildServerNameAliases(auditNames); + const boundaryNames = [...new Set([...auditNames, ...serverNameAliases.keys()])]; const seen = new Set(); const healedList = []; for (const tool of list) { @@ -286,15 +306,47 @@ async function healMcpToolNames({ req, tools, toolDefinitions }) { tool.includes(Constants.mcp_delimiter) && toolDefinitions[tool] == null ) { - const [, parsedServerName] = splitMCPToolKey(tool, rawServerNames); - if ( - parsedServerName != null && - rawServerNames.includes(parsedServerName) && - !shadowed.has(parsedServerName) - ) { - const healed = normalizeMCPToolKey(tool, rawServerNames); + const [, parsedServerName] = splitMCPToolKey(tool, boundaryNames); + let rawServerName; + if (parsedServerName != null && auditNames.includes(parsedServerName)) { + rawServerName = parsedServerName; + } else if (parsedServerName != null) { + const aliased = serverNameAliases.get(parsedServerName); + /** A normalized spelling on a CONTESTED slot is ambiguous between the + * tie-break winner and its shadowed rivals — rewriting persisted + * data must fail closed here, mirroring the raw-spelling shadow + * guard, rather than bind the reference to the winner. */ + const contested = + aliased != null && + auditNames.some( + (name) => name !== aliased && normalizeServerName(name) === parsedServerName, + ); + rawServerName = contested ? undefined : aliased; + } + if (rawServerName != null && !shadowed.has(rawServerName)) { + const healed = normalizeMCPToolKey(tool, auditNames); if (toolDefinitions[healed] != null) { healedTool = healed; + } else { + /** Catalog keys built after redundant-prefix stripping no longer + * match a pre-strip persisted key — without this second candidate + * the exact-lookup below silently drops the tool from the + * assistant. The rewrite only lands when the stripped key actually + * exists in the loaded definitions, so an unstripped catalog + * (collision guard kept the raw name) never heals into a phantom. */ + const keyServerName = normalizeServerName(rawServerName); + const [healedToolName] = splitMCPToolKey(healed, [keyServerName]); + const strippedName = stripServerNamePrefix(healedToolName, keyServerName); + const strippedKey = `${strippedName}${Constants.mcp_delimiter}${keyServerName}`; + /** Rewrite only when the stripped entry PROVES the same upstream + * identity — a stale key for a removed tool must not be healed + * onto a different sibling that kept its raw name. */ + if ( + strippedName !== healedToolName && + toolDefinitions[strippedKey]?.serverToolName === healedToolName + ) { + healedTool = strippedKey; + } } } } @@ -807,6 +859,11 @@ async function createMCPTools({ } const serverTools = []; + const keyServerName = normalizeServerName(serverName); + const keyToolNames = stripServerNamePrefixes( + result.tools.map((tool) => tool.name), + keyServerName, + ); for (const tool of result.tools) { const toolInstance = await createMCPTool({ res, @@ -821,7 +878,7 @@ async function createMCPTools({ serverName, /** Model-facing key: matches the normalized `availableTools` keys and * the instance name `createToolInstance` will assign. */ - toolKey: `${tool.name}${Constants.mcp_delimiter}${normalizeServerName(serverName)}`, + toolKey: `${keyToolNames.get(tool.name) ?? tool.name}${Constants.mcp_delimiter}${keyServerName}`, requestBody, requestScopedConnections, config: serverConfig, @@ -936,18 +993,49 @@ async function createMCPTool({ /** Legacy keys persisted pre-normalization (assistants, direct tool * calls) carry the RAW server name, while `availableTools` is keyed by - * the canonical normalized key — look up both spellings. */ + * the canonical normalized key — look up both spellings. Keys are also + * built after redundant server-name-prefix stripping now, so a persisted + * pre-strip key (`acme_foo_mcp_acme`) must additionally try + * its stripped spelling or the tool degrades to an unavailable stub. */ + const keyServerName = serverName != null ? normalizeServerName(serverName) : undefined; const canonicalToolKey = - serverName != null - ? `${toolName}${Constants.mcp_delimiter}${normalizeServerName(serverName)}` - : toolKey; - const findToolDefinition = (tools) => - tools?.[toolKey]?.function ?? - (canonicalToolKey !== toolKey ? tools?.[canonicalToolKey]?.function : undefined); - - /** @type {LCTool | undefined} */ - let toolDefinition = findToolDefinition(availableTools); - if (!toolDefinition) { + keyServerName != null ? `${toolName}${Constants.mcp_delimiter}${keyServerName}` : toolKey; + const strippedToolName = + keyServerName != null ? stripServerNamePrefix(toolName, keyServerName) : toolName; + const strippedToolKey = + strippedToolName !== toolName + ? `${strippedToolName}${Constants.mcp_delimiter}${keyServerName}` + : null; + const candidateToolKeys = [toolKey]; + if (canonicalToolKey !== toolKey) { + candidateToolKeys.push(canonicalToolKey); + } + if (strippedToolKey != null && !candidateToolKeys.includes(strippedToolKey)) { + candidateToolKeys.push(strippedToolKey); + } + let matchedToolKey = toolKey; + const findToolEntry = (tools) => { + for (const key of candidateToolKeys) { + const entry = tools?.[key]; + if (!entry?.function) { + continue; + } + /** The stripped-spelling candidate is only a legacy match when the + * entry PROVES the same upstream identity — without this, a stale + * reference to a removed tool could strip onto a DIFFERENT sibling + * that kept its raw name and silently call the wrong tool. */ + if (key === strippedToolKey && entry.serverToolName !== toolName) { + continue; + } + matchedToolKey = key; + return entry; + } + return undefined; + }; + + /** @type {LCFunctionTool | undefined} */ + let toolEntry = findToolEntry(availableTools); + if (!toolEntry) { const cachedAt = useMissingToolCache ? missingToolCache.get(toolKey) : undefined; if (cachedAt && Date.now() - cachedAt < MISSING_TOOL_TTL_MS) { logger.debug( @@ -976,15 +1064,15 @@ async function createMCPTool({ if (result?.availableTools) { onAvailableTools?.(result.availableTools); } - toolDefinition = findToolDefinition(result?.availableTools); + toolEntry = findToolEntry(result?.availableTools); - if (!toolDefinition && useMissingToolCache) { + if (!toolEntry && useMissingToolCache) { missingToolCache.set(toolKey, Date.now()); evictStale(missingToolCache, MISSING_TOOL_TTL_MS); } } - if (!toolDefinition) { + if (!toolEntry) { logger.warn( `[MCP][${serverName}][${toolName}] Tool definition not found, returning unavailable stub.`, ); @@ -998,10 +1086,20 @@ async function createMCPTool({ requestBody, requestScopedConnections, provider, + /** A legacy pre-strip key that resolves to the stripped entry KEEPS its + * persisted spelling as the instance name: `agent.tools` entries and + * `tool_options` keys reference that spelling, and renaming the instance + * would silently detach those per-tool settings. The upstream call name + * still comes from the MATCHED entry — its recorded raw name, or the + * matched key's own tool half when the entry was never stripped. */ toolName, + serverToolName: + toolEntry.serverToolName ?? + (matchedToolKey === strippedToolKey ? strippedToolName : toolName), + currentToolName: matchedToolKey === strippedToolKey ? strippedToolName : undefined, serverName, serverConfig, - toolDefinition, + toolDefinition: toolEntry['function'], streamId, jobCreatedAt, }); @@ -1014,6 +1112,8 @@ function createToolInstance({ requestBody: capturedRequestBody, requestScopedConnections: capturedRequestScopedConnections, toolName, + serverToolName = toolName, + currentToolName, serverName, serverConfig: capturedServerConfig, toolDefinition, @@ -1091,7 +1191,9 @@ function createToolInstance({ const result = await mcpManager.callTool({ serverName, serverConfig: capturedServerConfig, - toolName, + /** The upstream server never sees stripped names — a key that dropped + * a redundant server-name prefix calls the ORIGINAL tool. */ + toolName: serverToolName, provider, toolArguments, options: { @@ -1170,6 +1272,17 @@ function createToolInstance({ }); toolInstance.mcp = true; toolInstance.mcpRawServerName = serverName; + if (serverToolName !== toolName) { + /** Upstream identity for stripped keys — lets the options aliasing in + * `buildToolClassification` heal legacy `tool_options` spellings. */ + toolInstance.mcpServerToolName = serverToolName; + } + if (currentToolName != null && currentToolName !== toolName) { + /** Current catalog spelling for a LEGACY-named instance, so approval + * policies and hook matchers written against the current name still + * reach it (see `collectMCPToolAliases`). */ + toolInstance.mcpCurrentToolName = currentToolName; + } // Ephemeral request-scoped servers (runtime body placeholders) tear their // connection down at request end, so they must never be backgrounded. A // missing/stale config means the server's lifetime is unknowable, so fail @@ -1430,6 +1543,7 @@ async function getServerConnectionStatus( module.exports = { createMCPTool, createMCPTools, + toProviderToolDefinition, createMCPPermissionContext, userCanUseMCPServers, getMCPSetupData, diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index 976e6f81e82..c73fa2968b5 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -1830,6 +1830,146 @@ describe('User parameter passing tests', () => { expect(mockReinitMCPServer).not.toHaveBeenCalled(); }); + it('rejects a stripped-spelling entry without matching upstream identity', async () => { + /** A stale key for a removed tool must degrade to the unavailable stub, + * not resolve onto a DIFFERENT sibling whose key coincides with the + * stripped spelling. */ + const mockUser = { id: 'stale-identity-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + mockReinitMCPServer.mockResolvedValue(null); + + const staleKey = `acme_acme_foo${D}acme`; + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + toolKey: staleKey, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`acme_foo${D}acme`]: { + function: { + name: `acme_foo${D}acme`, + description: 'Different tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + expect(mockReinitMCPServer).toHaveBeenCalled(); + expect(mcpTool.description).toBe( + "This tool's MCP server is temporarily unavailable. Please try again shortly.", + ); + }); + + it('sends the raw upstream tool name when the key stripped a redundant server-name prefix', async () => { + const mockUser = { id: 'stripped-prefix-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + const callTool = jest.fn().mockResolvedValue(['ok', null]); + mockGetMCPManager.mockReturnValue({ callTool }); + + const strippedKey = `trace_top_time_consuming_operations${D}acme`; + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + toolKey: strippedKey, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [strippedKey]: { + serverToolName: 'acme_trace_top_time_consuming_operations', + function: { + name: strippedKey, + description: 'Trace', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + await mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ); + + expect(mcpTool.name).toBe(strippedKey); + expect(callTool).toHaveBeenCalledWith( + expect.objectContaining({ + serverName: 'acme', + toolName: 'acme_trace_top_time_consuming_operations', + }), + ); + }); + + it('resolves a legacy pre-strip tool key to the stripped definition without reinit', async () => { + const mockUser = { id: 'legacy-prefix-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + const callTool = jest.fn().mockResolvedValue(['ok', null]); + mockGetMCPManager.mockReturnValue({ callTool }); + + const strippedKey = `trace_top_time_consuming_operations${D}acme`; + const legacyKey = `acme_trace_top_time_consuming_operations${D}acme`; + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + toolKey: legacyKey, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [strippedKey]: { + serverToolName: 'acme_trace_top_time_consuming_operations', + function: { + name: strippedKey, + description: 'Trace', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + expect(mockReinitMCPServer).not.toHaveBeenCalled(); + + await mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ); + + /** The persisted spelling stays the instance name so `agent.tools` and + * `tool_options` keyed by it keep applying; only the upstream call + * uses the recorded raw name. */ + expect(mcpTool.name).toBe(legacyKey); + expect(callTool).toHaveBeenCalledWith( + expect.objectContaining({ + serverName: 'acme', + toolName: 'acme_trace_top_time_consuming_operations', + }), + ); + }); + it('should reject tool execution when user lacks MCP server use permission', async () => { const mockUser = { id: 'mcp-denied-user', role: 'USER' }; const mockRes = { write: jest.fn(), flush: jest.fn() }; diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index f03b6e8933b..81aae5aef93 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -1021,7 +1021,7 @@ async function loadToolDefinitionsWrapper({ return definitions; }; - let { toolDefinitions, toolRegistry, hasDeferredTools, mcpResolution } = + let { toolDefinitions, toolRegistry, hasDeferredTools, mcpToolAliases, mcpResolution } = await loadToolDefinitions( { userId: req.user.id, @@ -1134,6 +1134,7 @@ async function loadToolDefinitionsWrapper({ toolDefinitions = reloadResult.toolDefinitions; toolRegistry = reloadResult.toolRegistry; hasDeferredTools = reloadResult.hasDeferredTools; + mcpToolAliases = reloadResult.mcpToolAliases; mcpResolution = reloadResult.mcpResolution; } } @@ -1242,6 +1243,7 @@ async function loadToolDefinitionsWrapper({ dynamicToolContextMap, toolDefinitions, hasDeferredTools, + mcpToolAliases, actionsEnabled, primedCodeFiles, }; @@ -1434,7 +1436,7 @@ async function loadAgentTools({ /** Build tool registry from MCP tools and create PTC/tool search tools if configured */ const deferredToolsEnabled = checkCapability(AgentCapabilities.deferred_tools); const programmaticToolsEnabled = enabledCapabilities.has(AgentCapabilities.programmatic_tools); - const { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools } = + const { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools, mcpToolAliases } = await buildToolClassification({ loadedTools, userId: req.user.id, @@ -1504,6 +1506,7 @@ async function loadAgentTools({ dynamicToolContextMap, toolDefinitions, hasDeferredTools, + mcpToolAliases, actionsEnabled, tools: agentTools, primedCodeFiles, @@ -1523,6 +1526,7 @@ async function loadAgentTools({ dynamicToolContextMap, toolDefinitions, hasDeferredTools, + mcpToolAliases, actionsEnabled, tools: agentTools, primedCodeFiles, @@ -1653,6 +1657,7 @@ async function loadAgentTools({ userMCPAuthMap, toolDefinitions, hasDeferredTools, + mcpToolAliases, actionsEnabled, tools: agentTools, primedCodeFiles, diff --git a/api/server/services/__tests__/MCP.spec.js b/api/server/services/__tests__/MCP.spec.js index aa4ee0c0809..0fae7a3058d 100644 --- a/api/server/services/__tests__/MCP.spec.js +++ b/api/server/services/__tests__/MCP.spec.js @@ -115,8 +115,11 @@ describe('getAssistantToolDefinitions', () => { }); expect(definitions).toEqual({ - code_interpreter: { type: 'code_interpreter' }, - [toolKey]: mcpDefinition, + toolDefinitions: { + code_interpreter: { type: 'code_interpreter' }, + [toolKey]: mcpDefinition, + }, + accessibleServerNames: ['app-server'], }); expect(getMCPServerTools).toHaveBeenCalledWith('u1', 'app-server', serverConfig); }); @@ -135,7 +138,8 @@ describe('getAssistantToolDefinitions', () => { require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot }); await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({ - [toolKey]: mcpDefinition, + toolDefinitions: { [toolKey]: mcpDefinition }, + accessibleServerNames: ['app-server'], }); expect(cacheMCPServerTools).toHaveBeenCalledWith({ userId: 'u1', @@ -159,7 +163,8 @@ describe('getAssistantToolDefinitions', () => { reinitMCPServer.mockResolvedValue({ availableTools: { [toolKey]: mcpDefinition } }); await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({ - [toolKey]: mcpDefinition, + toolDefinitions: { [toolKey]: mcpDefinition }, + accessibleServerNames: ['app-server'], }); expect(reinitMCPServer).toHaveBeenCalledWith({ user: req.user, @@ -398,6 +403,140 @@ describe('healMcpToolNames', () => { expect(healed).toEqual([`search${Constants.mcp_delimiter}foo!`]); }); + it('heals a pre-strip prefixed key to the stripped catalog key', async () => { + /** Catalog keys drop a redundant leading server-name prefix now; an + * assistant saved before that resubmits the prefixed key and the exact + * lookup would silently drop the tool. */ + getAppConfig.mockResolvedValue({ mcpConfig: { acme: {} } }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ acme: {} }); + const strippedKey = `search${Constants.mcp_delimiter}acme`; + const toolDefinitions = { + [strippedKey]: { type: 'function', serverToolName: 'acme_search' }, + }; + + const healed = await healMcpToolNames({ + req, + tools: [`acme_search${Constants.mcp_delimiter}acme`], + toolDefinitions, + }); + + expect(healed).toEqual([strippedKey]); + }); + + it('heals a pre-strip key whose server suffix is already normalized', async () => { + /** Keys persisted after server-name normalization carry the NORMALIZED + * suffix, which the raw config names cannot match — the strip heal must + * resolve the boundary against both spellings. */ + getAppConfig.mockResolvedValue({ mcpConfig: { 'My Server': {} } }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ 'My Server': {} }); + const strippedKey = `search${Constants.mcp_delimiter}My_Server`; + const toolDefinitions = { + [strippedKey]: { type: 'function', serverToolName: 'my_server_search' }, + }; + + const healed = await healMcpToolNames({ + req, + tools: [`my_server_search${Constants.mcp_delimiter}My_Server`], + toolDefinitions, + }); + + expect(healed).toEqual([strippedKey]); + }); + + it('reuses a provided accessible-server snapshot without re-reading config', async () => { + /** The controllers pass the definitions loader's snapshot so the write + * path does not repeat the app-config and registry round trips. */ + const strippedKey = `search${Constants.mcp_delimiter}acme`; + const toolDefinitions = { + [strippedKey]: { type: 'function', serverToolName: 'acme_search' }, + }; + + const healed = await healMcpToolNames({ + req, + tools: [`acme_search${Constants.mcp_delimiter}acme`], + toolDefinitions, + accessibleServerNames: ['acme'], + }); + + expect(healed).toEqual([strippedKey]); + expect(getAppConfig).not.toHaveBeenCalled(); + expect(mockRegistry.getAllServerConfigs).not.toHaveBeenCalled(); + }); + + it('heals a pre-strip key for a USER-OWNED server absent from the operator config', async () => { + /** Assistants reference user DB servers too — the definitions loader + * resolves them, so the heal's audit must include them or the legacy + * key stays unhealed and the controllers drop the tool on edit. */ + getAppConfig.mockResolvedValue({ mcpConfig: {} }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ acme: {} }); + const strippedKey = `search${Constants.mcp_delimiter}acme`; + const toolDefinitions = { + [strippedKey]: { type: 'function', serverToolName: 'acme_search' }, + }; + + const healed = await healMcpToolNames({ + req, + tools: [`acme_search${Constants.mcp_delimiter}acme`], + toolDefinitions, + }); + + expect(healed).toEqual([strippedKey]); + }); + + it('does not heal a stale key onto a sibling that lacks matching upstream identity', async () => { + /** With `acme_acme_foo` removed upstream while `acme_foo` kept its raw + * name, the stale key's stripped spelling exists but belongs to a + * DIFFERENT tool — the identity check must reject the rewrite. */ + getAppConfig.mockResolvedValue({ mcpConfig: { acme: {} } }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ acme: {} }); + const staleKey = `acme_acme_foo${Constants.mcp_delimiter}acme`; + const toolDefinitions = { + [`acme_foo${Constants.mcp_delimiter}acme`]: { type: 'function' }, + [`foo${Constants.mcp_delimiter}acme`]: { type: 'function', serverToolName: 'acme_foo' }, + }; + + const healed = await healMcpToolNames({ req, tools: [staleKey], toolDefinitions }); + + expect(healed).toEqual([staleKey]); + }); + + it('fails closed on a normalized-suffix key whose slot is CONTESTED', async () => { + /** `My Server` and `My_Server!` both normalize to `My_Server`, so a + * normalized-suffix reference is ambiguous between them — rewriting + * persisted data must not bind it to the tie-break winner. */ + getAppConfig.mockResolvedValue({ mcpConfig: { 'My Server': {}, 'My_Server!': {} } }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ 'My Server': {}, 'My_Server!': {} }); + const legacyKey = `my_server_search${Constants.mcp_delimiter}My_Server`; + const toolDefinitions = { [`search${Constants.mcp_delimiter}My_Server`]: { type: 'function' } }; + + const healed = await healMcpToolNames({ req, tools: [legacyKey], toolDefinitions }); + + expect(healed).toEqual([legacyKey]); + }); + + it('keeps a prefixed key whose stripped spelling is not in the loaded definitions', async () => { + /** When the catalog kept the raw name (bare-sibling collision), the + * prefixed key IS canonical and must not be rewritten into a key owned + * by the bare tool. */ + getAppConfig.mockResolvedValue({ mcpConfig: { acme: {} } }); + mockRegistry.ensureConfigServers.mockResolvedValue({}); + mockRegistry.getAllServerConfigs.mockResolvedValue({ acme: {} }); + const prefixedKey = `acme_search${Constants.mcp_delimiter}acme`; + const toolDefinitions = { + [prefixedKey]: { type: 'function' }, + [`search${Constants.mcp_delimiter}acme`]: { type: 'function' }, + }; + + const healed = await healMcpToolNames({ req, tools: [prefixedKey], toolDefinitions }); + + expect(healed).toEqual([prefixedKey]); + }); + it('skips the config read entirely when every delimiter-bearing name resolves', async () => { const key = `search${Constants.mcp_delimiter}srv`; const healed = await healMcpToolNames({ diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx index d03ed996303..95248599d6e 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx @@ -8,6 +8,7 @@ import { splitMCPToolKey, normalizeServerName, buildServerNameAliases, + stripServerNamePrefix, } from 'librechat-data-provider'; import type { MouseEvent } from 'react'; import type { TranslationKeys } from '~/hooks/useLocalize'; @@ -136,7 +137,7 @@ export default function McpSection({ item }: Props) { * runtime heal keeps them active, and per-tool updates could never replace * the legacy entry. Tokens and other servers' entries pass through. */ - const toCurrentToolId = useCallback( + const toNormalizedToolId = useCallback( (entry: string): string => { const normalizedName = normalizeServerName(serverName); if ( @@ -172,6 +173,44 @@ export default function McpSection({ item }: Props) { [serverName, serverToken, serverAllToken, mcpServersMap], ); + /** + * Second migration stage: catalog keys drop a redundant leading server-name + * prefix, so a pre-strip persisted id would show its tool unchecked and a + * per-tool toggle could silently drop it from the selection. The rewrite is + * identity-verified — it only lands when the stripped catalog entry records + * this exact raw name as its upstream tool — so a stale id for a removed + * tool can never migrate onto a different sibling. + */ + /** Constant-time lookups for the migration below — the form heal calls it + * per persisted key, so linear catalog scans go O(options × tools). */ + const toolsById = useMemo(() => new Map(tools.map((tool) => [tool.tool_id, tool])), [tools]); + + const toStrippedToolId = useCallback( + (entry: string): string => { + if (entry === serverToken || entry === serverAllToken || toolsById.has(entry)) { + return entry; + } + const normalizedName = normalizeServerName(serverName); + const [toolPart, parsed] = splitMCPToolKey(entry, [normalizedName]); + if (parsed !== normalizedName) { + return entry; + } + const strippedPart = stripServerNamePrefix(toolPart, normalizedName); + if (strippedPart === toolPart) { + return entry; + } + const strippedId = `${strippedPart}${Constants.mcp_delimiter}${normalizedName}`; + const target = toolsById.get(strippedId); + return target?.metadata.serverToolName === toolPart ? strippedId : entry; + }, + [serverName, serverToken, serverAllToken, toolsById], + ); + + const toCurrentToolId = useCallback( + (entry: string): string => toStrippedToolId(toNormalizedToolId(entry)), + [toNormalizedToolId, toStrippedToolId], + ); + const isServerSelection = useCallback( (token: string): boolean => { const allServerNames = Array.from(new Set([...mcpServersMap.keys(), serverName])); diff --git a/packages/api/src/agents/hitl/policy.spec.ts b/packages/api/src/agents/hitl/policy.spec.ts index 1a43846d546..e95f3b034ac 100644 --- a/packages/api/src/agents/hitl/policy.spec.ts +++ b/packages/api/src/agents/hitl/policy.spec.ts @@ -2,6 +2,9 @@ import type { Agents, TToolApprovalPolicy } from 'librechat-data-provider'; import { resolveToolApprovalPolicy, isHITLEnabled, + healToolApprovalPolicy, + collectAliasMatcherNames, + buildAliasMatcherPattern, mapToolApprovalPolicy, buildToolApprovalPayload, buildAskUserQuestionPayload, @@ -773,3 +776,87 @@ describe('exemptAskUserQuestionFromApproval', () => { expect(exemptAskUserQuestionFromApproval(undefined, NAME)).toBeUndefined(); }); }); + +describe('healToolApprovalPolicy', () => { + const aliases = [ + { name: 'delete_thing_mcp_acme', aliasName: 'acme_delete_thing_mcp_acme' }, + { name: 'search_mcp_acme', aliasName: 'acme_search_mcp_acme' }, + ]; + + it('appends current names to lists whose patterns match only the legacy spelling', () => { + /** Admin YAML written against upstream naming must keep applying — a + * non-matching deny fails OPEN. */ + const healed = healToolApprovalPolicy( + { enabled: true, deny: ['acme_delete_*'], ask: ['acme_search_mcp_acme'] }, + aliases, + ); + + expect(healed?.deny).toEqual(['acme_delete_*', 'delete_thing_mcp_acme']); + expect(healed?.ask).toEqual(['acme_search_mcp_acme', 'search_mcp_acme']); + }); + + it('heals list-level so allow semantics are preserved, not tightened', () => { + const healed = healToolApprovalPolicy({ enabled: true, allow: ['acme_search_*'] }, aliases); + + expect(healed?.allow).toEqual(['acme_search_*', 'search_mcp_acme']); + }); + + it('skips names the list already matches and leaves non-matching lists untouched', () => { + const healed = healToolApprovalPolicy( + { enabled: true, deny: ['*_mcp_acme'], allow: ['unrelated_tool'] }, + aliases, + ); + + expect(healed?.deny).toEqual(['*_mcp_acme']); + expect(healed?.allow).toEqual(['unrelated_tool']); + }); + + it('passes through without aliases or policy', () => { + expect(healToolApprovalPolicy(undefined, aliases)).toBeUndefined(); + const policy: TToolApprovalPolicy = { enabled: true, deny: ['x'] }; + expect(healToolApprovalPolicy(policy, [])).toBe(policy); + }); +}); + +describe('healToolApprovalPolicy reverse direction', () => { + it('appends a legacy-named instance when the pattern targets the current catalog name', () => { + /** An unedited agent retains the pre-strip instance name — a deny written + * against the current catalog name must still reach it. */ + const aliases = [{ name: 'acme_search_mcp_acme', aliasName: 'search_mcp_acme' }]; + const healed = healToolApprovalPolicy( + { enabled: true, mode: 'bypass', deny: ['search_mcp_acme'] }, + aliases, + ); + + expect(healed?.deny).toEqual(['search_mcp_acme', 'acme_search_mcp_acme']); + }); +}); + +describe('collectAliasMatcherNames', () => { + const aliases = [ + { name: 'search_mcp_acme', aliasName: 'acme_search_mcp_acme' }, + { name: 'acme_list_mcp_acme', aliasName: 'list_mcp_acme' }, + ]; + + it('returns names whose alias matches the regex while the name does not', () => { + expect(collectAliasMatcherNames('^acme_search_mcp_acme$', aliases)).toEqual([ + 'search_mcp_acme', + ]); + expect(collectAliasMatcherNames('^list_mcp_acme$', aliases)).toEqual(['acme_list_mcp_acme']); + }); + + it('skips names the matcher already matches and invalid patterns', () => { + expect(collectAliasMatcherNames('_mcp_acme$', aliases)).toEqual([]); + expect(collectAliasMatcherNames('(unclosed', aliases)).toEqual([]); + expect(collectAliasMatcherNames(undefined, aliases)).toEqual([]); + }); + + it('builds an anchored exact-name pattern with escaped names', () => { + const pattern = buildAliasMatcherPattern(['a.b_mcp_acme', 'c_mcp_acme']); + const regex = new RegExp(pattern); + expect(regex.test('a.b_mcp_acme')).toBe(true); + expect(regex.test('axb_mcp_acme')).toBe(false); + expect(regex.test('c_mcp_acme')).toBe(true); + expect(regex.test('xc_mcp_acme')).toBe(false); + }); +}); diff --git a/packages/api/src/agents/hitl/policy.ts b/packages/api/src/agents/hitl/policy.ts index 535ce37d72b..d3350ec0632 100644 --- a/packages/api/src/agents/hitl/policy.ts +++ b/packages/api/src/agents/hitl/policy.ts @@ -2,6 +2,7 @@ import { randomUUID, createHash } from 'crypto'; import { openAIBaseSchema, googleBaseSchema, anthropicBaseSchema } from 'librechat-data-provider'; import type { Agents, TToolApprovalPolicy } from 'librechat-data-provider'; import type { ToolPolicyConfig } from '@librechat/agents'; +import type { MCPToolAlias } from '~/tools/classification'; /** * Default decisions offered to the user for a paused tool call. @@ -86,6 +87,93 @@ export function isHITLEnabled(policy: TToolApprovalPolicy | undefined): boolean * defaults apply). The `enabled` field is LibreChat-only and stripped here — * it's consumed separately via {@link isHITLEnabled} to gate the SDK opt-out. */ +/** Anchored-glob matcher mirroring the SDK's `createToolPolicyHook` semantics exactly. */ +function globToRegex(pattern: string): RegExp { + const escaped = pattern.replace(/[.+?^${}()|[\]\\]/g, '\\$&'); + return new RegExp('^' + escaped.replace(/\*/g, '.*') + '$'); +} + +/** + * Extends each `toolApproval` pattern list with the names of tools whose + * OTHER spelling matches, so admin YAML keeps applying when a tool's key + * spelling changed in either direction: patterns written against pre-strip + * upstream naming reach the stripped instances (a non-matching `deny` would + * otherwise FAIL OPEN), and patterns written against the current catalog + * naming reach legacy-named instances retained by unedited agents. Healing + * is list-level (literal names appended, patterns never rewritten), so + * `deny`/`ask`/`allow` precedence semantics are unchanged, and a name + * already matched by its own list is skipped. + */ +export function healToolApprovalPolicy( + policy: TToolApprovalPolicy | undefined, + aliases: readonly MCPToolAlias[], +): TToolApprovalPolicy | undefined { + if (!policy || aliases.length === 0) { + return policy; + } + const healList = (patterns: string[] | undefined): string[] | undefined => { + if (!patterns || patterns.length === 0) { + return patterns; + } + const regexes = patterns.map(globToRegex); + const appended: string[] = []; + for (const { name, aliasName } of aliases) { + if (name === aliasName || regexes.some((regex) => regex.test(name))) { + continue; + } + if (regexes.some((regex) => regex.test(aliasName))) { + appended.push(name); + } + } + return appended.length > 0 ? [...patterns, ...appended] : patterns; + }; + return { + ...policy, + allow: healList(policy.allow), + deny: healList(policy.deny), + ask: healList(policy.ask), + }; +} + +function escapeRegExp(value: string): string { + return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); +} + +/** + * Names whose OTHER spelling matches a programmatic hook's regex matcher + * while their own name does not — the hook must also fire for these or its + * argument-, user-, or tenant-specific deny/ask decisions are silently + * skipped for renamed tools. Mirrors the SDK's unanchored `new RegExp(pattern)` + * matcher semantics; an invalid pattern matches nothing there, so it aliases + * nothing here. + */ +export function collectAliasMatcherNames( + matcher: string | undefined, + aliases: readonly MCPToolAlias[], +): string[] { + if (!matcher || aliases.length === 0) { + return []; + } + let regex: RegExp; + try { + regex = new RegExp(matcher); + } catch { + return []; + } + const names: string[] = []; + for (const { name, aliasName } of aliases) { + if (name !== aliasName && !regex.test(name) && regex.test(aliasName)) { + names.push(name); + } + } + return names; +} + +/** Anchored exact-name pattern for the alias-matched names of one hook matcher. */ +export function buildAliasMatcherPattern(names: readonly string[]): string { + return `^(?:${names.map(escapeRegExp).join('|')})$`; +} + export function mapToolApprovalPolicy( policy: TToolApprovalPolicy | undefined, ): ToolPolicyConfig | undefined { diff --git a/packages/api/src/agents/hitl/runtime.ts b/packages/api/src/agents/hitl/runtime.ts index b7c7404ad8a..5d96ae05a38 100644 --- a/packages/api/src/agents/hitl/runtime.ts +++ b/packages/api/src/agents/hitl/runtime.ts @@ -1,7 +1,13 @@ import { HookRegistry, createToolPolicyHook } from '@librechat/agents'; import type { TToolApprovalPolicy } from 'librechat-data-provider'; +import type { MCPToolAlias } from '~/tools/classification'; import type { ToolApprovalHookContext } from './hooks'; -import { isHITLEnabled, mapToolApprovalPolicy } from './policy'; +import { + isHITLEnabled, + mapToolApprovalPolicy, + collectAliasMatcherNames, + buildAliasMatcherPattern, +} from './policy'; import { buildToolApprovalHooks } from './hooks'; /** @@ -33,6 +39,7 @@ export interface HITLRunWiring { export function buildHITLRunWiring( policy: TToolApprovalPolicy | undefined, context: ToolApprovalHookContext = {}, + mcpToolAliases: readonly MCPToolAlias[] = [], ): HITLRunWiring | undefined { if (!isHITLEnabled(policy)) { return undefined; @@ -52,6 +59,20 @@ export function buildHITLRunWiring( 'PreToolUse', matcher ? { pattern: matcher, hooks: [hook] } : { hooks: [hook] }, ); + /** A matcher written against a tool's OTHER key spelling (pre-strip or + * current) would silently never fire for the renamed instance, skipping + * its argument/user/tenant-specific deny or ask. The SAME hook is + * registered again under an exact-name pattern for those aliased names + * — a separate entry keeps the admin's regex semantics and the SDK's + * pattern-length cap intact, and the name sets are disjoint so the hook + * never fires twice for one call. */ + const aliasNames = matcher ? collectAliasMatcherNames(matcher, mcpToolAliases) : []; + if (aliasNames.length > 0) { + registry.register('PreToolUse', { + pattern: buildAliasMatcherPattern(aliasNames), + hooks: [hook], + }); + } } return { humanInTheLoop: { enabled: true }, hooks: registry }; diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 7e49a30a913..9c429832d25 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -40,6 +40,7 @@ import type { } from '~/types'; import type { LCAvailableTools, RequestScopedMCPConnectionStore } from '../mcp/types'; import type { TFilterFilesByAgentAccess } from './resources'; +import type { MCPToolAlias } from '~/tools/classification'; import { injectSkillCatalog, resolveManualSkills, @@ -279,6 +280,8 @@ export type InitializedAgent = Agent & { requestScopedConnections?: RequestScopedMCPConnectionStore; /** Serializable tool definitions for event-driven execution */ toolDefinitions?: LCTool[]; + /** Both-direction identity aliases for MCP tools whose key spelling changed */ + mcpToolAliases?: MCPToolAlias[]; /** Precomputed flag indicating if any tools have defer_loading enabled (for efficient runtime checks) */ hasDeferredTools?: boolean; /** @@ -444,6 +447,7 @@ export interface InitializeAgentParams { /** Serializable tool definitions for event-driven mode */ toolDefinitions?: LCTool[]; hasDeferredTools?: boolean; + mcpToolAliases?: MCPToolAlias[]; actionsEnabled?: boolean; /** * Pre-uploaded code-env file refs for the agent's @@ -1106,6 +1110,7 @@ export async function initializeAgent( mcpAvailableTools, requestScopedConnections, hasDeferredTools, + mcpToolAliases, actionsEnabled, tools: structuredTools, primedCodeFiles, @@ -1119,6 +1124,7 @@ export async function initializeAgent( requestScopedConnections: undefined, toolDefinitions: [], hasDeferredTools: false, + mcpToolAliases: [], actionsEnabled: undefined, primedCodeFiles: undefined, }; @@ -1557,6 +1563,7 @@ export async function initializeAgent( userMCPAuthMap, toolDefinitions, hasDeferredTools, + mcpToolAliases, backgroundToolNames, intentToolNames, actionsEnabled, diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index e198b13095f..8e7b9400849 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -41,6 +41,7 @@ import type { BaseMessage } from '@librechat/agents/langchain/messages'; import type { AppConfig, IUser } from '@librechat/data-schemas'; import type { ToolInputValidationError } from '~/agents/toolValidation'; import type { ResolvedAlwaysApplySkill } from '~/agents/skills'; +import type { MCPToolAlias } from '~/tools/classification'; import type { SubagentUsageEvent } from '~/agents/usage'; import type * as t from '~/types'; import { @@ -49,6 +50,11 @@ import { stripBackgroundFromToolRegistry, stripBackgroundFromToolDefinitions, } from '~/agents/background'; +import { + resolveToolApprovalPolicy, + healToolApprovalPolicy, + exemptAskUserQuestionFromApproval, +} from '~/agents/hitl/policy'; import { ASK_USER_QUESTION_TOOL_NAME, createAskUserQuestionTool, @@ -57,7 +63,6 @@ import { createSubagentWakeupHandleHook, usesSubagentCompletionWakeups, } from '~/agents/subagentDelivery'; -import { resolveToolApprovalPolicy, exemptAskUserQuestionFromApproval } from '~/agents/hitl/policy'; import { applyCustomHandoffPromptKeyCompatibility } from '~/agents/handoffPromptKeyCompatibility'; import { stripIntentFromToolRegistry, stripIntentFromToolDefinitions } from '~/agents/intent'; import { isSteeringSupported, isSteerPreemptSupported } from '~/agents/steering/runtime'; @@ -378,6 +383,8 @@ type RunAgent = Omit & { toolDefinitions?: LCTool[]; /** Precomputed flag indicating if any tools have defer_loading enabled */ hasDeferredTools?: boolean; + /** Both-direction identity aliases for MCP tools whose key spelling changed */ + mcpToolAliases?: MCPToolAlias[]; /** Names of tools injected with the `run_in_background` param (excluded from eager execution). */ backgroundToolNames?: string[]; /** Names of tools with the host-injected `intent` param (stripped from self-spawn inputs). */ @@ -1694,18 +1701,29 @@ export async function createRun({ // would pause with no approval surface or resume endpoint, and the route would emit a // normal final response / `[DONE]` with the tool call dangling. Only AgentClient (chat + // resume) passes `hitlCapable`; without it the run is identical to the no-HITL path. + /** Both-direction key-spelling aliases collected at tool classification — + * identical in instance and event-driven loading modes. */ + const mcpToolAliases = agents.flatMap((agent) => agent.mcpToolAliases ?? []); const hitl = hitlCapable ? buildHITLRunWiring( // The ask tool is exempt from the approval prompt (unless explicitly // listed by the admin) — approving the right to ask a question is a - // pure double-pause; the tool has no side effects to gate. - exemptAskUserQuestionFromApproval(toolApprovalPolicy, ASK_USER_QUESTION_TOOL_NAME), + // pure double-pause; the tool has no side effects to gate. Pattern + // lists are healed against the tools' other key spellings first, so + // admin globs written for pre-strip upstream names keep applying (a + // non-matching deny would fail OPEN), and rules written against + // current catalog names reach legacy-named instances. + exemptAskUserQuestionFromApproval( + healToolApprovalPolicy(toolApprovalPolicy, mcpToolAliases), + ASK_USER_QUESTION_TOOL_NAME, + ), { userId: user?.id, conversationId: requestBody?.conversationId, tenantId: tenantId ?? user?.tenantId, appConfig, }, + mcpToolAliases, ) : undefined; /** diff --git a/packages/api/src/mcp/assistants.spec.ts b/packages/api/src/mcp/assistants.spec.ts index a3269114040..38eafdd2047 100644 --- a/packages/api/src/mcp/assistants.spec.ts +++ b/packages/api/src/mcp/assistants.spec.ts @@ -1,7 +1,7 @@ import { Constants } from 'librechat-data-provider'; import type { LCAvailableTools, ParsedServerConfig } from './types'; import type { AssistantToolDefinitionsDeps } from './assistants'; -import { getAssistantToolDefinitions } from './assistants'; +import { getAssistantToolDefinitions, toProviderToolDefinition } from './assistants'; const serverConfig: ParsedServerConfig = { type: 'streamable-http', @@ -48,12 +48,47 @@ describe('getAssistantToolDefinitions', () => { const deps = createDeps(); await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual({ - ...params.staticTools, - ...catalog, + toolDefinitions: { ...params.staticTools, ...catalog }, + accessibleServerNames: ['app-server'], }); expect(deps.getMCPServerTools).toHaveBeenCalledWith('user-1', 'app-server', serverConfig); }); + it('retains serverToolName for the heal; toProviderToolDefinition strips it at submission', async () => { + /** The heal verifies legacy rewrites against the recorded upstream + * identity, so the loader keeps the field; assistant writers submit + * entries verbatim, so the controllers sanitize each entry through + * toProviderToolDefinition before the provider sees it. */ + const strippedKey = `search${Constants.mcp_delimiter}app-server`; + const strippedCatalog: LCAvailableTools = { + [strippedKey]: { + type: 'function', + serverToolName: 'app-server_search', + ['function']: { + name: strippedKey, + description: '', + parameters: { type: 'object', properties: {} }, + }, + }, + }; + const deps = createDeps({ getMCPServerTools: jest.fn().mockResolvedValue(strippedCatalog) }); + + const { toolDefinitions } = await getAssistantToolDefinitions(params, deps); + + expect(toolDefinitions[strippedKey]?.serverToolName).toBe('app-server_search'); + + const sanitized = toProviderToolDefinition(toolDefinitions[strippedKey]); + expect(sanitized).toEqual({ + type: 'function', + ['function']: strippedCatalog[strippedKey]['function'], + }); + expect(sanitized).not.toHaveProperty('serverToolName'); + expect(toProviderToolDefinition('code_interpreter')).toBe('code_interpreter'); + expect(toProviderToolDefinition(params.staticTools.code_interpreter)).toBe( + params.staticTools.code_interpreter, + ); + }); + it('reconnects a user server when neither cache nor local snapshot has a catalog', async () => { const recoveredCatalog = { ...catalog }; const recoverServerTools = jest.fn().mockResolvedValue(recoveredCatalog); @@ -64,8 +99,8 @@ describe('getAssistantToolDefinitions', () => { }); await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual({ - ...params.staticTools, - ...recoveredCatalog, + toolDefinitions: { ...params.staticTools, ...recoveredCatalog }, + accessibleServerNames: ['app-server'], }); expect(recoverServerTools).toHaveBeenCalledWith('app-server', serverConfig); }); @@ -118,7 +153,10 @@ describe('getAssistantToolDefinitions', () => { cacheMCPServerTools, }); - await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual(params.staticTools); + await expect(getAssistantToolDefinitions(params, deps)).resolves.toEqual({ + toolDefinitions: params.staticTools, + accessibleServerNames: ['app-server'], + }); expect(cacheMCPServerTools).toHaveBeenCalledWith({ userId: 'user-1', serverName: 'app-server', @@ -192,7 +230,7 @@ describe('getAssistantToolDefinitions', () => { }, deps, ), - ).resolves.toBe(staticTools); + ).resolves.toEqual({ toolDefinitions: staticTools }); expect(deps.ensureConfigServers).not.toHaveBeenCalled(); expect(deps.getMCPServerTools).not.toHaveBeenCalled(); }); diff --git a/packages/api/src/mcp/assistants.ts b/packages/api/src/mcp/assistants.ts index 6b0160ce4ac..78906210b45 100644 --- a/packages/api/src/mcp/assistants.ts +++ b/packages/api/src/mcp/assistants.ts @@ -6,7 +6,7 @@ import { splitMCPToolKey, } from 'librechat-data-provider'; import type { MCPOptions } from 'librechat-data-provider'; -import type { LCAvailableTools, ParsedServerConfig } from '~/mcp/types'; +import type { LCAvailableTools, LCFunctionTool, ParsedServerConfig } from '~/mcp/types'; import { createConcurrencyLimiter } from '~/utils/promise'; import { findShadowedServerNames } from '~/mcp/utils'; @@ -154,11 +154,22 @@ async function loadServerCatalog( throw new Error(`MCP tool definitions unavailable for assistant server "${serverName}"`); } +export interface AssistantToolDefinitionsResult { + toolDefinitions: LCAvailableTools; + /** + * Every server name the principal can reach, from the same merged registry + * read that resolved the catalogs — the legacy-key heal reuses it instead + * of repeating the app-config and registry round trips on the write path. + * `undefined` when the payload references no MCP tools (nothing to heal). + */ + accessibleServerNames?: string[]; +} + /** Loads the static catalog with the configuration-addressed MCP slices referenced by an assistant. */ export async function getAssistantToolDefinitions( params: AssistantToolDefinitionsParams, deps: AssistantToolDefinitionsDeps, -): Promise { +): Promise { const mcpToolNames = params.tools?.filter( (tool): tool is string => @@ -166,7 +177,7 @@ export async function getAssistantToolDefinitions( ) ?? []; const userId = params.user?.id; if (mcpToolNames.length === 0 || !userId) { - return params.staticTools; + return { toolDefinitions: params.staticTools }; } const configs = await resolveAssistantMcpConfigs( @@ -182,5 +193,31 @@ export async function getAssistantToolDefinitions( (serverName) => loadServerCatalog(userId, serverName, configs[serverName], deps, recover), ), ); - return Object.assign({}, params.staticTools, ...serverCatalogs); + /** Entries keep `serverToolName` here: the assistants heal verifies legacy + * key rewrites against that upstream identity. The controllers sanitize + * through {@link toProviderToolDefinition} at the submission boundary. */ + return { + toolDefinitions: Object.assign({}, params.staticTools, ...serverCatalogs), + accessibleServerNames: [ + ...new Set([...Object.keys(configs), ...Object.keys(params.mcpConfig)]), + ], + }; +} + +/** + * Assistant writers submit tool entries VERBATIM as provider tool definitions + * (`assistantData.tools` in the v1/v2 controllers), and providers reject + * unknown fields — the internal `serverToolName` mapping must never leave the + * catalog. Strings and entries without the mapping pass through by reference; + * the cached catalog keeps the mapping for the runtime call path. + */ +export function toProviderToolDefinition(tool: T): T | LCFunctionTool { + if (tool == null || typeof tool !== 'object') { + return tool; + } + const entry = tool as Partial; + if (entry.serverToolName == null || entry.type !== 'function' || entry['function'] == null) { + return tool; + } + return { type: entry.type, ['function']: entry['function'] }; } diff --git a/packages/api/src/mcp/catalog/store.ts b/packages/api/src/mcp/catalog/store.ts index 96cae3c75c5..574a71f4ad1 100644 --- a/packages/api/src/mcp/catalog/store.ts +++ b/packages/api/src/mcp/catalog/store.ts @@ -116,14 +116,23 @@ redis.call('DEL', KEYS[2]) return 1 `; +/** + * Catalog entries can carry `serverToolName` (redundant server-name prefix + * stripping): an older replica reading a stripped entry ignores the mapping + * and calls the stripped key segment upstream. Versioning the MCP catalog + * slices keeps mixed-version replicas on their own representation during a + * rolling deploy; stale slices simply expire. + */ +const CATALOG_VERSION = 'v2'; + export const ToolCacheKeys = { GLOBAL: 'tools:global', MCP_APP_SERVER: (serverName: string, configGeneration: string): string => - `tools:mcp:app:${encodeURIComponent(serverName)}:${encodeURIComponent(configGeneration)}`, + `tools:mcp:app:${CATALOG_VERSION}:${encodeURIComponent(serverName)}:${encodeURIComponent(configGeneration)}`, MCP_SERVER: (userId: string, serverName: string, configGeneration?: string): string => configGeneration - ? `tools:mcp:user:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}:${encodeURIComponent(configGeneration)}` - : `tools:mcp:${userId}:${serverName}`, + ? `tools:mcp:user:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}:${CATALOG_VERSION}:${encodeURIComponent(configGeneration)}` + : `tools:mcp:${CATALOG_VERSION}:${userId}:${serverName}`, MCP_SERVER_GENERATION: (userId: string, serverName: string): string => `tools:metadata:mcp:user-generation:{${encodeURIComponent(userId)}:${encodeURIComponent(serverName)}}`, MCP_SERVER_LEGACY_FENCE: (userId: string, serverName: string): string => diff --git a/packages/api/src/mcp/registry/MCPServerInspector.ts b/packages/api/src/mcp/registry/MCPServerInspector.ts index b0c1b436e0b..fe05d41e6cf 100644 --- a/packages/api/src/mcp/registry/MCPServerInspector.ts +++ b/packages/api/src/mcp/registry/MCPServerInspector.ts @@ -1,5 +1,5 @@ import { logger } from '@librechat/data-schemas'; -import { Constants, normalizeServerName } from 'librechat-data-provider'; +import { Constants, normalizeServerName, stripServerNamePrefixes } from 'librechat-data-provider'; import type { JsonSchemaType } from '@librechat/data-schemas'; import type { MCPConnection } from '~/mcp/connection'; import type * as t from '~/mcp/types'; @@ -188,10 +188,16 @@ export class MCPServerInspector { /** Model-facing key: must match the runtime instance name, which embeds * the normalized server name (see `createToolInstance` in MCP.js). */ const keyServerName = normalizeServerName(serverName); + const keyToolNames = stripServerNamePrefixes( + tools.map((tool) => tool.name), + keyServerName, + ); tools.forEach((tool) => { - const name = `${tool.name}${Constants.mcp_delimiter}${keyServerName}`; + const keyToolName = keyToolNames.get(tool.name) ?? tool.name; + const name = `${keyToolName}${Constants.mcp_delimiter}${keyServerName}`; toolFunctions[name] = { type: 'function', + ...(keyToolName !== tool.name && { serverToolName: tool.name }), ['function']: { name, description: tool.description, diff --git a/packages/api/src/mcp/registry/MCPServersInitializer.ts b/packages/api/src/mcp/registry/MCPServersInitializer.ts index 4cdd148383a..143d14c699b 100644 --- a/packages/api/src/mcp/registry/MCPServersInitializer.ts +++ b/packages/api/src/mcp/registry/MCPServersInitializer.ts @@ -22,9 +22,13 @@ const DEFAULT_FOLLOWER_RETRY_MS = 3000; * followers short-circuit on the stale status and never re-tag entries written * by the previous version. Bumped to 3 so cached entries whose `serverInstructions` still holds * inspector-fetched text are rewritten with the declaration preserved and the text moved to - * `resolvedInstructions`. + * `resolvedInstructions`. Bumped to 4 so persisted `toolFunctions` are rebuilt + * with redundant server-name prefixes stripped and `serverToolName` recorded — + * otherwise a follower accepts the previous deployment's config hash and + * republishes pre-strip definitions into the current catalog namespace + * indefinitely. */ -const REGISTRY_STORAGE_SCHEMA_VERSION = 3; +const REGISTRY_STORAGE_SCHEMA_VERSION = 4; const parseDurationMs = ( value: string | undefined, diff --git a/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts b/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts index 52990606991..c77b1b32e7d 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPServerInspector.test.ts @@ -597,6 +597,33 @@ describe('MCPServerInspector', () => { expect(result[key]['function'].name).toBe(key); }); + it('strips a redundant server-name prefix from keys and records the raw name', async () => { + mockConnection.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ + complete: true, + tools: [ + { + name: 'acme_trace_top_time_consuming_operations', + description: 'Trace', + inputSchema: { type: 'object', properties: {} }, + }, + { + name: 'list_services', + description: 'List', + inputSchema: { type: 'object', properties: {} }, + }, + ], + }); + + const { tools: result } = await MCPServerInspector.getToolCatalog('acme', mockConnection); + + const strippedKey = 'trace_top_time_consuming_operations_mcp_acme'; + const plainKey = 'list_services_mcp_acme'; + expect(Object.keys(result).sort()).toEqual([plainKey, strippedKey].sort()); + expect(result[strippedKey]['function'].name).toBe(strippedKey); + expect(result[strippedKey].serverToolName).toBe('acme_trace_top_time_consuming_operations'); + expect(result[plainKey].serverToolName).toBeUndefined(); + }); + it('rejects an incomplete snapshot before it can replace cached tools', async () => { mockConnection.fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({ tools: [{ name: 'partial', inputSchema: { type: 'object' } }], diff --git a/packages/api/src/mcp/tools.spec.ts b/packages/api/src/mcp/tools.spec.ts index 03726be831a..56fad5f6661 100644 --- a/packages/api/src/mcp/tools.spec.ts +++ b/packages/api/src/mcp/tools.spec.ts @@ -495,6 +495,49 @@ describe('createMCPToolCacheService', () => { }); }); + it('strips a redundant server-name prefix from keys and records the raw name', async () => { + /** `acme_trace..._mcp_acme` carries the server twice and can push the + * model-facing name past provider function-name limits (64). */ + const deps = createMockDeps(); + const tools: MCPToolInput[] = [ + { name: 'acme_trace_top_time_consuming_operations', description: 'Trace' }, + { name: 'list_services', description: 'List' }, + ]; + const result = await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'acme', + tools, + }); + + const strippedKey = toolName('trace_top_time_consuming_operations', 'acme'); + const plainKey = toolName('list_services', 'acme'); + expect(Object.keys(result ?? {}).sort()).toEqual([plainKey, strippedKey].sort()); + expect(result?.[strippedKey]?.['function'].name).toBe(strippedKey); + expect(result?.[strippedKey]?.serverToolName).toBe( + 'acme_trace_top_time_consuming_operations', + ); + expect(result?.[plainKey]?.serverToolName).toBeUndefined(); + }); + + it('keeps the prefixed key when stripping would collide with a sibling tool', async () => { + const deps = createMockDeps(); + const tools: MCPToolInput[] = [ + { name: 'search', description: 'Plain' }, + { name: 'acme_search', description: 'Prefixed' }, + ]; + const result = await createMCPToolCacheService(deps).updateMCPServerTools({ + userId: 'u1', + serverName: 'acme', + tools, + }); + + const plainKey = toolName('search', 'acme'); + const prefixedKey = toolName('acme_search', 'acme'); + expect(Object.keys(result ?? {}).sort()).toEqual([prefixedKey, plainKey].sort()); + expect(result?.[plainKey]?.serverToolName).toBeUndefined(); + expect(result?.[prefixedKey]?.serverToolName).toBeUndefined(); + }); + it('builds request-scoped tools without caching them', async () => { const deps = createMockDeps({ getServerConfig: jest.fn().mockResolvedValue(requestScopedConfig), diff --git a/packages/api/src/mcp/tools.ts b/packages/api/src/mcp/tools.ts index b1b801e1ca8..d043d2cad1f 100644 --- a/packages/api/src/mcp/tools.ts +++ b/packages/api/src/mcp/tools.ts @@ -1,5 +1,10 @@ import { logger } from '@librechat/data-schemas'; -import { Constants, buildServerNameAliases, normalizeServerName } from 'librechat-data-provider'; +import { + Constants, + buildServerNameAliases, + normalizeServerName, + stripServerNamePrefixes, +} from 'librechat-data-provider'; import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { JsonSchemaType } from '@librechat/agents'; import type { LCAvailableTools, LCFunctionTool, ParsedServerConfig } from './types'; @@ -234,8 +239,13 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS * `normalizeServerName(serverName)`. The cache STORE itself stays keyed * by the raw config name. */ const keyServerName = normalizeServerName(serverName); + const keyToolNames = stripServerNamePrefixes( + tools.map((tool) => tool.name), + keyServerName, + ); for (const tool of tools) { - const name = `${tool.name}${mcpDelimiter}${keyServerName}`; + const keyToolName = keyToolNames.get(tool.name) ?? tool.name; + const name = `${keyToolName}${mcpDelimiter}${keyServerName}`; const entry: LCFunctionTool = { type: 'function', ['function']: { @@ -246,6 +256,9 @@ export function createMCPToolCacheService(deps: MCPToolCacheDeps): MCPToolCacheS : ({ type: 'object', properties: {} } as JsonSchemaType), }, }; + if (keyToolName !== tool.name) { + entry.serverToolName = tool.name; + } serverTools[name] = entry; } diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index 4585ac1bd60..ae59ed4af4f 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -49,6 +49,9 @@ export interface MCPResource { export interface LCFunctionTool { type: 'function'; ['function']: LCTool; + /** Raw upstream tool name when the model-facing key stripped a redundant + * server-name prefix — tool calls must send THIS name to the server. */ + serverToolName?: string; } export type LCAvailableTools = Record; diff --git a/packages/api/src/mcp/utils.ts b/packages/api/src/mcp/utils.ts index a715a1df09e..0abbe460347 100644 --- a/packages/api/src/mcp/utils.ts +++ b/packages/api/src/mcp/utils.ts @@ -631,4 +631,6 @@ export { normalizeServerName, normalizeMCPToolKey, buildServerNameAliases, + stripServerNamePrefix, + stripServerNamePrefixes, } from 'librechat-data-provider'; diff --git a/packages/api/src/tools/classification.spec.ts b/packages/api/src/tools/classification.spec.ts index b543d2bad19..642bbc1e4ac 100644 --- a/packages/api/src/tools/classification.spec.ts +++ b/packages/api/src/tools/classification.spec.ts @@ -4,8 +4,10 @@ import type { GenericTool } from '@librechat/agents'; import type { LCToolRegistry } from './classification'; import { buildToolRegistryFromAgentOptions, + aliasMCPToolOptions, agentHasProgrammaticTools, buildToolClassification, + collectMCPToolAliases, getServerNameFromTool, agentHasDeferredTools, } from './classification'; @@ -28,6 +30,111 @@ describe('classification.ts', () => { }); }); + describe('collectMCPToolAliases', () => { + it('collects both alias directions from definitions', () => { + const defs = [ + { name: 'search_mcp_acme', serverName: 'acme', serverToolName: 'acme_search' }, + { + name: 'acme_list_mcp_acme', + serverName: 'acme', + serverToolName: 'acme_list', + currentToolName: 'list', + }, + { name: 'plain_mcp_acme', serverName: 'acme' }, + ]; + + expect(collectMCPToolAliases(defs)).toEqual([ + { name: 'search_mcp_acme', aliasName: 'acme_search_mcp_acme' }, + { name: 'acme_list_mcp_acme', aliasName: 'list_mcp_acme' }, + ]); + }); + + it('normalizes the server name when reconstructing alias keys', () => { + const defs = [ + { + name: 'search_mcp_My_Server', + serverName: 'My Server', + serverToolName: 'my_server_search', + }, + ]; + + expect(collectMCPToolAliases(defs)).toEqual([ + { name: 'search_mcp_My_Server', aliasName: 'my_server_search_mcp_My_Server' }, + ]); + }); + }); + + describe('aliasMCPToolOptions', () => { + it('aliases pre-strip option keys onto the current instance name, identity-gated', () => { + /** Wildcard-expanded catalogs rename stripped tools without any + * `agent.tools` entry to preserve the spelling — persisted defer, + * programmatic, background, and intent settings must follow. */ + const defs = [ + { + name: 'search_mcp_acme', + serverName: 'acme', + serverToolName: 'acme_search', + }, + { name: 'list_items_mcp_acme', serverName: 'acme' }, + ]; + const agentToolOptions: AgentToolOptions = { + acme_search_mcp_acme: { defer_loading: true }, + }; + + aliasMCPToolOptions(collectMCPToolAliases(defs), agentToolOptions); + + expect(agentToolOptions['search_mcp_acme']).toEqual({ defer_loading: true }); + const registry = buildToolRegistryFromAgentOptions(defs, agentToolOptions); + expect(registry.get('search_mcp_acme')?.defer_loading).toBe(true); + }); + + it('aliases current-keyed options back onto a legacy-named instance', () => { + /** The editor migrates `tool_options` keys to the current catalog + * spelling, while an unedited `agent.tools` entry keeps the legacy + * instance name — options must follow the reverse direction too. */ + const defs = [ + { + name: 'acme_search_mcp_acme', + serverName: 'acme', + serverToolName: 'acme_search', + currentToolName: 'search', + }, + ]; + const agentToolOptions: AgentToolOptions = { + search_mcp_acme: { defer_loading: true }, + }; + + aliasMCPToolOptions(collectMCPToolAliases(defs), agentToolOptions); + + expect(agentToolOptions['acme_search_mcp_acme']).toEqual({ defer_loading: true }); + const registry = buildToolRegistryFromAgentOptions(defs, agentToolOptions); + expect(registry.get('acme_search_mcp_acme')?.defer_loading).toBe(true); + }); + + it('never overrides an explicit entry under the instance name', () => { + const defs = [{ name: 'search_mcp_acme', serverName: 'acme', serverToolName: 'acme_search' }]; + const agentToolOptions: AgentToolOptions = { + search_mcp_acme: { defer_loading: false }, + acme_search_mcp_acme: { defer_loading: true }, + }; + + aliasMCPToolOptions(collectMCPToolAliases(defs), agentToolOptions); + + expect(agentToolOptions['search_mcp_acme']).toEqual({ defer_loading: false }); + }); + + it('does nothing without recorded upstream identity', () => { + const defs = [{ name: 'search_mcp_acme', serverName: 'acme' }]; + const agentToolOptions: AgentToolOptions = { + acme_search_mcp_acme: { defer_loading: true }, + }; + + aliasMCPToolOptions(collectMCPToolAliases(defs), agentToolOptions); + + expect(agentToolOptions['search_mcp_acme']).toBeUndefined(); + }); + }); + describe('buildToolRegistryFromAgentOptions', () => { it('should use agent tool options for defer_loading', () => { const tools = [ diff --git a/packages/api/src/tools/classification.ts b/packages/api/src/tools/classification.ts index 2ede26ffcb4..be83fa08ef6 100644 --- a/packages/api/src/tools/classification.ts +++ b/packages/api/src/tools/classification.ts @@ -6,7 +6,7 @@ */ import { logger } from '@librechat/data-schemas'; -import { Constants } from 'librechat-data-provider'; +import { Constants, normalizeServerName } from 'librechat-data-provider'; import { Providers, createToolSearch, @@ -33,6 +33,47 @@ export interface ToolDefinition { parameters?: JsonSchemaType; /** MCP server name extracted from tool name */ serverName?: string; + /** Raw upstream tool name when the model-facing key stripped a redundant server-name prefix */ + serverToolName?: string; + /** Current catalog tool name when a LEGACY persisted key kept its pre-strip spelling */ + currentToolName?: string; +} + +/** An MCP tool name plus its OTHER spelling (legacy for stripped instances, current for legacy-named ones). */ +export interface MCPToolAlias { + name: string; + aliasName: string; +} + +/** + * Collects both directions of identity aliases from MCP tool definitions, so + * approval policies and hook matchers written against EITHER spelling keep + * applying: a stripped instance aliases its pre-strip name, and a + * legacy-named instance (persisted key retained) aliases its current catalog + * name. Works in both loading modes because both funnel their definitions + * through {@link buildToolClassification}. + */ +export function collectMCPToolAliases(mcpToolDefs: ToolDefinition[]): MCPToolAlias[] { + const aliases: MCPToolAlias[] = []; + for (const def of mcpToolDefs) { + if (!def.serverName) { + continue; + } + const keySuffix = `${Constants.mcp_delimiter}${normalizeServerName(def.serverName)}`; + if (def.serverToolName) { + const aliasName = `${def.serverToolName}${keySuffix}`; + if (aliasName !== def.name) { + aliases.push({ name: def.name, aliasName }); + } + } + if (def.currentToolName) { + const aliasName = `${def.currentToolName}${keySuffix}`; + if (aliasName !== def.name) { + aliases.push({ name: def.name, aliasName }); + } + } + } + return aliases; } /** @@ -49,6 +90,31 @@ export function getServerNameFromTool(toolName: string): string | undefined { return undefined; } +/** + * Aliases persisted `tool_options` keys onto the instance names IN PLACE, so + * every downstream reader of `agent.tool_options` (the registry build for + * defer/programmatic, the background and intent passes) sees the healed keys + * in BOTH loading modes and BOTH spelling directions: options keyed by a + * pre-strip spelling follow a renamed (wildcard-expanded) instance, and + * options the editor migrated to the CURRENT catalog spelling still reach a + * legacy-named instance an unedited `agent.tools` entry retained. + * Identity-gated through {@link collectMCPToolAliases}, and an explicit + * entry under the instance's own name always wins. + */ +export function aliasMCPToolOptions( + aliases: readonly MCPToolAlias[], + agentToolOptions?: AgentToolOptions, +): void { + if (!agentToolOptions || Object.keys(agentToolOptions).length === 0) { + return; + } + for (const { name, aliasName } of aliases) { + if (agentToolOptions[name] == null && agentToolOptions[aliasName] != null) { + agentToolOptions[name] = agentToolOptions[aliasName]; + } + } +} + /** * Builds a tool registry from agent-level tool_options. * @@ -104,6 +170,10 @@ interface MCPToolInstance { mcpJsonSchema?: JsonSchemaType; /** Server this tool came from, carried from resolution instead of re-parsed */ mcpRawServerName?: string; + /** Raw upstream tool name when the instance name stripped a redundant server-name prefix */ + mcpServerToolName?: string; + /** Current catalog tool name when a legacy persisted key kept its pre-strip spelling */ + mcpCurrentToolName?: string; } /** @@ -129,6 +199,14 @@ export function extractMCPToolDefinition(tool: MCPToolInstance): ToolDefinition def.serverName = serverName; } + if (tool.mcpServerToolName) { + def.serverToolName = tool.mcpServerToolName; + } + + if (tool.mcpCurrentToolName) { + def.currentToolName = tool.mcpCurrentToolName; + } + return def; } @@ -214,6 +292,8 @@ export interface BuildToolClassificationResult { additionalTools: GenericTool[]; /** Whether any tools have defer_loading enabled (precomputed for efficiency) */ hasDeferredTools: boolean; + /** Both-direction identity aliases for MCP tools whose key spelling changed (see {@link collectMCPToolAliases}) */ + mcpToolAliases: MCPToolAlias[]; } /** @@ -282,10 +362,13 @@ export async function buildToolClassification( toolDefinitions: [], toolRegistry: undefined, hasDeferredTools: false, + mcpToolAliases: [], }; } const mcpToolDefs = mcpTools.map(extractMCPToolDefinition); + const mcpToolAliases = collectMCPToolAliases(mcpToolDefs); + aliasMCPToolOptions(mcpToolAliases, agentToolOptions); const toolRegistry: LCToolRegistry = buildToolRegistry(mcpToolDefs, agentToolOptions); /** Clean up temporary mcpJsonSchema property from tools now that registry is populated */ @@ -318,7 +401,13 @@ export async function buildToolClassification( logger.debug( `[buildToolClassification] Agent ${agentId} has no programmatic or deferred tools, skipping PTC/ToolSearch`, ); - return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools: false }; + return { + toolRegistry, + toolDefinitions, + additionalTools, + hasDeferredTools: false, + mcpToolAliases, + }; } /** Tool search uses local mode (no API key needed) */ @@ -357,7 +446,7 @@ export async function buildToolClassification( } if (!hasProgrammaticTools) { - return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools }; + return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools, mcpToolAliases }; } /** In definitions-only mode, add PTC definition without creating the tool instance */ @@ -374,7 +463,7 @@ export async function buildToolClassification( logger.debug( `[buildToolClassification] PTC definition added for agent ${agentId} (definitions only)`, ); - return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools }; + return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools, mcpToolAliases }; } try { @@ -408,5 +497,5 @@ export async function buildToolClassification( logger.error('[buildToolClassification] Error creating PTC tool:', error); } - return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools }; + return { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools, mcpToolAliases }; } diff --git a/packages/api/src/tools/definitions.spec.ts b/packages/api/src/tools/definitions.spec.ts index 5df9b16d988..9e7fc239e19 100644 --- a/packages/api/src/tools/definitions.spec.ts +++ b/packages/api/src/tools/definitions.spec.ts @@ -575,6 +575,72 @@ describe('definitions.ts', () => { expect(getItemDef?.description).toBe('Get a specific item'); }); + it('resolves a pre-strip persisted key against the stripped catalog, keeping the persisted name', async () => { + /** Catalog keys drop a redundant leading server-name prefix; an agent + * saved before that must still resolve, and the definition keeps the + * persisted spelling so it matches the runtime instance name. */ + const mockServerTools = { + search_mcp_acme: { + serverToolName: 'acme_search', + function: { + name: 'search_mcp_acme', + description: 'Search things', + parameters: { type: 'object', properties: {} }, + }, + }, + }; + + mockGetOrFetchMCPServerTools.mockResolvedValue(mockServerTools); + + const params: LoadToolDefinitionsParams = { + userId: 'user-123', + agentId: 'agent-123', + tools: ['acme_search_mcp_acme'], + }; + + const deps: LoadToolDefinitionsDeps = { + getOrFetchMCPServerTools: mockGetOrFetchMCPServerTools, + isBuiltInTool: mockIsBuiltInTool, + }; + + const result = await loadToolDefinitions(params, deps); + + expect(result.toolDefinitions).toHaveLength(1); + expect(result.toolDefinitions[0]?.name).toBe('acme_search_mcp_acme'); + expect(result.toolDefinitions[0]?.description).toBe('Search things'); + }); + + it('rejects a stripped-spelling match without matching upstream identity', async () => { + /** A stale key for a removed tool must not resolve onto a DIFFERENT + * sibling whose key merely coincides with the stripped spelling. */ + const mockServerTools = { + acme_foo_mcp_acme: { + function: { + name: 'acme_foo_mcp_acme', + description: 'Different tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }; + + mockGetOrFetchMCPServerTools.mockResolvedValue(mockServerTools); + + const params: LoadToolDefinitionsParams = { + userId: 'user-123', + agentId: 'agent-123', + tools: ['acme_acme_foo_mcp_acme'], + }; + + const deps: LoadToolDefinitionsDeps = { + getOrFetchMCPServerTools: mockGetOrFetchMCPServerTools, + isBuiltInTool: mockIsBuiltInTool, + }; + + const result = await loadToolDefinitions(params, deps); + + expect(result.toolDefinitions).toHaveLength(0); + }); + it('union-flattens MCP tool schemas for Google, but preserves unions otherwise', async () => { const mockServerTools = { issue_write_mcp_github: { diff --git a/packages/api/src/tools/definitions.ts b/packages/api/src/tools/definitions.ts index 8904bf9fcbb..d4b5967b1d6 100644 --- a/packages/api/src/tools/definitions.ts +++ b/packages/api/src/tools/definitions.ts @@ -10,11 +10,13 @@ import { Constants, isActionTool, splitMCPToolKey, + normalizeServerName, + stripServerNamePrefix, buildServerNameAliases, } from 'librechat-data-provider'; import type { LCToolRegistry, JsonSchemaType, LCTool, GenericTool } from '@librechat/agents'; import type { AgentToolOptions } from 'librechat-data-provider'; -import type { ToolDefinition } from './classification'; +import type { MCPToolAlias, ToolDefinition } from './classification'; import { resolveJsonSchemaRefs, normalizeJsonSchema, sanitizeGeminiSchema } from '~/mcp/zod'; import { buildToolClassification } from './classification'; import { getToolDefinition } from './registry/definitions'; @@ -27,6 +29,7 @@ export interface MCPServerTool { description?: string; parameters?: JsonSchemaType; }; + serverToolName?: string; } export type MCPServerTools = Record; @@ -92,6 +95,8 @@ export interface LoadToolDefinitionsResult { toolDefinitions: (ToolDefinition | LCTool)[]; toolRegistry: LCToolRegistry; hasDeferredTools: boolean; + /** Both-direction identity aliases for MCP tools whose key spelling changed */ + mcpToolAliases: MCPToolAlias[]; mcpResolution: { expectedToolCount: number; resolvedToolCount: number; @@ -145,6 +150,7 @@ export async function loadToolDefinitions( toolDefinitions: [], toolRegistry: new Map(), hasDeferredTools: false, + mcpToolAliases: [], mcpResolution: { expectedToolCount: 0, resolvedToolCount: 0 }, }; @@ -247,9 +253,38 @@ export async function loadToolDefinitions( continue; } + /** Catalog keys are built after redundant server-name-prefix stripping — + * a pre-strip persisted key (`acme_search_mcp_acme`) must also try its + * stripped spelling or the agent fails initialization with its expected + * tools "unavailable". The definition keeps the PERSISTED name so it + * matches the runtime instance `createMCPTool` builds for the same key, + * and the stripped entry is accepted only when its recorded raw name + * PROVES the same upstream identity. */ + const findToolMatch = ( + tools: Record, + ): { def: MCPServerTool; currentToolName?: string } | undefined => { + const direct = tools[toolName]; + if (direct?.function) { + return { def: direct }; + } + const keyServerName = normalizeServerName(serverName); + const [toolPart] = splitMCPToolKey(toolName, [parsed]); + const strippedPart = stripServerNamePrefix(toolPart, keyServerName); + if (strippedPart === toolPart) { + return undefined; + } + const entry = tools[`${strippedPart}${Constants.mcp_delimiter}${keyServerName}`]; + /** `currentToolName` records the catalog spelling so approval policies + * and hook matchers written against it still reach this legacy-named + * definition (see `collectMCPToolAliases`). */ + return entry?.serverToolName === toolPart + ? { def: entry, currentToolName: strippedPart } + : undefined; + }; + const selectedToolMissing = isMCPAllPlaceholder(toolName) ? Object.keys(serverTools).length === 0 - : !serverTools[toolName]?.function; + : !findToolMatch(serverTools)?.def.function; if (selectedToolMissing && refreshMCPServerTools && !refreshedServerNames.has(serverName)) { refreshedServerNames.add(serverName); const refreshedTools = await refreshMCPServerTools(userId, serverName); @@ -267,6 +302,7 @@ export async function loadToolDefinitions( description: toolDef.function.description || undefined, parameters: buildMcpParameters(toolDef.function.parameters), serverName, + serverToolName: toolDef.serverToolName, }); resolvedMCPToolCount++; } @@ -274,13 +310,15 @@ export async function loadToolDefinitions( continue; } - const toolDef = serverTools[toolName]; - if (toolDef?.function) { + const toolMatch = findToolMatch(serverTools); + if (toolMatch?.def.function) { mcpToolDefs.push({ name: toolName, - description: toolDef.function.description || undefined, - parameters: buildMcpParameters(toolDef.function.parameters), + description: toolMatch.def.function.description || undefined, + parameters: buildMcpParameters(toolMatch.def.function.parameters), serverName, + serverToolName: toolMatch.def.serverToolName, + currentToolName: toolMatch.currentToolName, }); resolvedMCPToolCount++; } @@ -301,6 +339,8 @@ export async function loadToolDefinitions( mcp: true as const, mcpJsonSchema: def.parameters, mcpRawServerName: def.serverName, + mcpServerToolName: def.serverToolName, + mcpCurrentToolName: def.currentToolName, })) as unknown as GenericTool[]; const classificationResult = await buildToolClassification({ @@ -350,6 +390,7 @@ export async function loadToolDefinitions( toolDefinitions: allDefinitions, toolRegistry, hasDeferredTools, + mcpToolAliases: classificationResult.mcpToolAliases, mcpResolution: { expectedToolCount: expectedMCPToolCount, resolvedToolCount: resolvedMCPToolCount, diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index c7661f39768..7ef02184bc2 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -12,6 +12,7 @@ import { ComponentTypes, SettingTypes, OptionTypes } from './generate'; import { MAX_SUBAGENTS, MAX_SUBAGENTS_CEILING } from './limits'; import { STATEFUL_CODE_ENVIRONMENTS } from './stateful-code'; import { specsConfigSchema, TSpecsConfig } from './models'; +import { isActionTool } from './types/assistants'; import { REFILL_INTERVAL_UNITS } from './balance'; import { fileConfigSchema } from './file-config'; import { apiBaseUrl } from './api-endpoints'; @@ -3177,6 +3178,120 @@ export function normalizeMCPToolKey(toolKey: string, rawServerNames: readonly st return `${toolKey.slice(0, toolKey.length - matched.length)}${normalized}`; } +/** + * Strips a redundant leading server-name prefix from a raw upstream tool name + * before it is embedded into a model-facing key, so the key doesn't carry the + * server twice (`acme_trace_..._mcp_acme`) and push long tool names + * past provider function-name limits (64 chars). The match is case-insensitive + * because display-cased server names ("Acme") conventionally prefix their + * tools in lowercase. Ingestion that strips must record the original name + * (`serverToolName` on the cached definition) — tool calls send THAT name back + * to the server, never the stripped one. Catalog producers must not call this + * directly: only {@link stripServerNamePrefixes} sees the whole sibling set and + * can keep colliding results apart. + */ +export function stripServerNamePrefix(toolName: string, normalizedServerName: string): string { + const prefixLength = normalizedServerName.length + 1; + if (toolName.length <= prefixLength) { + return toolName; + } + const prefix = toolName.slice(0, prefixLength).toLowerCase(); + if (prefix !== `${normalizedServerName.toLowerCase()}_`) { + return toolName; + } + const stripped = toolName.slice(prefixLength); + if (isReservedMCPToolName(stripped)) { + return toolName; + } + /** `isActionTool` classifies keys by the RELATIVE position of `_action_` + * and `_mcp_`; stripping moves the first `_mcp_` earlier, so a server + * whose normalized name contains `_action_` could see a real MCP tool + * reclassified as an OpenAPI action (bypassing MCP authorization). Never + * produce a key whose classification differs from the raw key's. */ + const keySuffix = `${Constants.mcp_delimiter}${normalizedServerName}`; + if (isActionTool(`${stripped}${keySuffix}`) !== isActionTool(`${toolName}${keySuffix}`)) { + return toolName; + } + return stripped; +} + +/** + * Synthetic markers consumed by prefix (`isMCPAllPlaceholder`, the server-pin + * skip, the client's OAuth stream classification), so each reserves BOTH its + * exact name and its `${marker}${mcp_delimiter}` namespace: a stripped + * remainder inside any of them would turn a real upstream tool into the + * server-wide wildcard, the UI pin placeholder, or a synthetic OAuth call. + */ +const RESERVED_MCP_TOOL_MARKERS: readonly string[] = [ + `${Constants.mcp_all}`, + `${Constants.mcp_server}`, + 'oauth', +]; + +function isReservedMCPToolName(toolName: string): boolean { + /** `mcp_` opens the server-scoped pluginKey namespace (`mcp_${serverName}`), + * and `lc_transfer_to_` opens the agent-handoff namespace (the client + * renders such calls as handoffs; the background and intent passes exclude + * them) — pre-strip tool keys could never enter either, since they always + * began with the server name itself. */ + if ( + toolName.startsWith(`${Constants.mcp_prefix}`) || + toolName.startsWith(`${Constants.LC_TRANSFER_TO_}`) + ) { + return true; + } + return RESERVED_MCP_TOOL_MARKERS.some( + (marker) => toolName === marker || toolName.startsWith(`${marker}${Constants.mcp_delimiter}`), + ); +} + +/** + * Maps every raw tool name in a server's catalog to its model-facing name, + * stripping redundant server-name prefixes collision-free: when two names + * yield the same result — a bare `foo` next to `_foo`, or the + * case-variant pair `_Foo` / `_Foo` under the case-insensitive + * prefix match — every collider keeps its raw name, so two distinct upstream + * tools can never collapse onto one key. Unprefixed names count against the + * result set through their identity mapping, which is what makes the bare-name + * case fall out of the same counter. + */ +export function stripServerNamePrefixes( + toolNames: readonly string[], + normalizedServerName: string, +): Map { + const rawNames = new Set(toolNames); + const finalNames = new Map( + toolNames.map((name) => { + const stripped = stripServerNamePrefix(name, normalizedServerName); + /** Every sibling's RAW name is reserved even when that sibling itself + * strips away: keys persisted BEFORE stripping embed raw names, so a + * stripped result landing on another sibling's raw name would route + * that sibling's legacy references to the wrong upstream tool. */ + return [name, stripped !== name && rawNames.has(stripped) ? name : stripped]; + }), + ); + /** Reverting a collider to its raw name can itself collide with ANOTHER + * sibling's stripped result (`foo` / `acme_foo` / `acme_acme_foo`), so the + * guard iterates to a fixpoint. Each pass converts at least one stripped + * result back to its unique raw name, so it terminates within the catalog + * size. */ + let changed = true; + while (changed) { + changed = false; + const counts = new Map(); + finalNames.forEach((result) => { + counts.set(result, (counts.get(result) ?? 0) + 1); + }); + finalNames.forEach((result, raw) => { + if (result !== raw && (counts.get(result) ?? 0) > 1) { + finalNames.set(raw, raw); + changed = true; + } + }); + } + return finalNames; +} + export function splitMCPToolKey( toolKey: string, knownServerNames?: readonly string[], diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index ba40501687a..68f435a948e 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -765,6 +765,9 @@ export const tPluginSchema = z.object({ chatMenu: z.boolean().optional(), isButton: z.boolean().optional(), toolkit: z.boolean().optional(), + /** Raw upstream tool name when the model-facing key stripped a redundant + * server-name prefix — proves upstream identity for legacy id migration. */ + serverToolName: z.string().optional(), }); export type TPlugin = z.infer; diff --git a/packages/data-provider/src/splitMCPToolKey.spec.ts b/packages/data-provider/src/splitMCPToolKey.spec.ts index 1b05ac404d2..e39d55afc92 100644 --- a/packages/data-provider/src/splitMCPToolKey.spec.ts +++ b/packages/data-provider/src/splitMCPToolKey.spec.ts @@ -4,6 +4,8 @@ import { splitToolCallName, normalizeMCPToolKey, buildServerNameAliases, + stripServerNamePrefix, + stripServerNamePrefixes, } from './config'; describe('splitMCPToolKey', () => { @@ -208,3 +210,127 @@ describe('splitToolCallName oauth precedence', () => { ]); }); }); + +describe('stripServerNamePrefix', () => { + it('strips a leading server-name prefix from the tool name', () => { + expect(stripServerNamePrefix('acme_trace_top_time_consuming_operations', 'acme')).toBe( + 'trace_top_time_consuming_operations', + ); + }); + + it('matches the prefix case-insensitively', () => { + /** Display-cased server names ("Acme") conventionally prefix their + * tools in lowercase — the redundancy is the same either way. */ + expect(stripServerNamePrefix('acme_list_services', 'Acme')).toBe('list_services'); + }); + + it('returns the name unchanged when the prefix does not match', () => { + expect(stripServerNamePrefix('github_create_issue', 'acme')).toBe('github_create_issue'); + }); + + it('requires the underscore separator, not a bare substring match', () => { + expect(stripServerNamePrefix('acmecorp_tool', 'acme')).toBe('acmecorp_tool'); + }); + + it('keeps a name that is exactly the server name or would strip to empty', () => { + expect(stripServerNamePrefix('acme', 'acme')).toBe('acme'); + expect(stripServerNamePrefix('acme_', 'acme')).toBe('acme_'); + }); +}); + +describe('stripServerNamePrefixes', () => { + it('maps every raw name to its stripped model-facing name', () => { + const map = stripServerNamePrefixes(['acme_search', 'list_services'], 'acme'); + expect(map.get('acme_search')).toBe('search'); + expect(map.get('list_services')).toBe('list_services'); + }); + + it('keeps the prefixed name when stripping would collide with a bare sibling', () => { + /** A server exposing BOTH `search` and `acme_search` must keep two + * distinct keys — stripping would collapse them into one. */ + const map = stripServerNamePrefixes(['search', 'acme_search'], 'acme'); + expect(map.get('search')).toBe('search'); + expect(map.get('acme_search')).toBe('acme_search'); + }); + + it('keeps both raw names when case-variant prefixed siblings strip to the same result', () => { + /** The prefix match is case-insensitive, so `acme_Foo` and `Acme_Foo` are + * distinct upstream tools with the SAME stripped remainder — both must + * fall back to their raw names or one silently overwrites the other. */ + const map = stripServerNamePrefixes(['acme_Foo', 'Acme_Foo'], 'acme'); + expect(map.get('acme_Foo')).toBe('acme_Foo'); + expect(map.get('Acme_Foo')).toBe('Acme_Foo'); + }); + + it('collisions do not suppress stripping of unrelated siblings', () => { + const map = stripServerNamePrefixes(['search', 'acme_search', 'acme_trace'], 'acme'); + expect(map.get('acme_trace')).toBe('trace'); + }); + + it('reserves every sibling raw name, even when that sibling itself strips', () => { + /** Keys persisted BEFORE stripping embed raw names: if `acme_acme_foo` + * stripped to `acme_foo`, a pre-rollout reference to the REAL `acme_foo` + * would exact-match the wrong tool in the same snapshot. */ + const map = stripServerNamePrefixes(['acme_foo', 'acme_acme_foo'], 'acme'); + expect(map.get('acme_foo')).toBe('foo'); + expect(map.get('acme_acme_foo')).toBe('acme_acme_foo'); + }); + + it('resolves secondary collisions introduced by a fallback to a raw name', () => { + /** `acme_foo` falls back to raw because of the bare `foo`, which then + * collides with `acme_acme_foo`'s stripped result — the guard must + * iterate until no two final names coincide. */ + const map = stripServerNamePrefixes(['foo', 'acme_foo', 'acme_acme_foo'], 'acme'); + expect(map.get('foo')).toBe('foo'); + expect(map.get('acme_foo')).toBe('acme_foo'); + expect(map.get('acme_acme_foo')).toBe('acme_acme_foo'); + expect(new Set(map.values()).size).toBe(3); + }); + + it('never strips a remainder that equals a synthetic MCP marker', () => { + /** `sys__all__sys` keys expand to every server tool, `sys__server__sys` + * keys are skipped as UI placeholders, and `oauth${mcp_delimiter}` names + * get OAuth-only handling in the client stream handlers — a real + * upstream tool must not be renamed onto any of them. */ + expect(stripServerNamePrefix(`acme_${Constants.mcp_all}`, 'acme')).toBe( + `acme_${Constants.mcp_all}`, + ); + expect(stripServerNamePrefix(`acme_${Constants.mcp_server}`, 'acme')).toBe( + `acme_${Constants.mcp_server}`, + ); + expect(stripServerNamePrefix('acme_oauth', 'acme')).toBe('acme_oauth'); + /** Each marker is consumed by PREFIX (`isMCPAllPlaceholder`, the + * server-pin skip, the client's OAuth classification), so the whole + * `${marker}${mcp_delimiter}` namespace stays raw, not just the exact + * name. */ + expect(stripServerNamePrefix(`acme_oauth${Constants.mcp_delimiter}reset`, 'acme')).toBe( + `acme_oauth${Constants.mcp_delimiter}reset`, + ); + expect( + stripServerNamePrefix(`acme_${Constants.mcp_all}${Constants.mcp_delimiter}reset`, 'acme'), + ).toBe(`acme_${Constants.mcp_all}${Constants.mcp_delimiter}reset`); + expect( + stripServerNamePrefix(`acme_${Constants.mcp_server}${Constants.mcp_delimiter}reset`, 'acme'), + ).toBe(`acme_${Constants.mcp_server}${Constants.mcp_delimiter}reset`); + /** `mcp_` opens the server-scoped pluginKey namespace and + * `lc_transfer_to_` the agent-handoff namespace — pre-strip tool keys + * could never enter either. */ + expect(stripServerNamePrefix('acme_mcp_status', 'acme')).toBe('acme_mcp_status'); + expect(stripServerNamePrefix('acme_lc_transfer_to_status', 'acme')).toBe( + 'acme_lc_transfer_to_status', + ); + }); + + it('never flips isActionTool classification for the produced key', () => { + /** `isActionTool` compares the FIRST `_action_` and `_mcp_` positions; + * stripping moves `_mcp_` earlier, so a server whose normalized name + * contains `_action_` (e.g. "svc action v1") would see a real MCP tool + * reclassified as an OpenAPI action and bypass MCP authorization. */ + expect(stripServerNamePrefix('svc_action_v1_report', 'svc_action_v1')).toBe( + 'svc_action_v1_report', + ); + /** A remainder containing `_action_` in the tool half does not flip and + * still strips. */ + expect(stripServerNamePrefix('acme_do_action_thing', 'acme')).toBe('do_action_thing'); + }); +}); diff --git a/packages/data-provider/src/types/queries.ts b/packages/data-provider/src/types/queries.ts index 8d95a94a914..2a481d3df55 100644 --- a/packages/data-provider/src/types/queries.ts +++ b/packages/data-provider/src/types/queries.ts @@ -130,6 +130,9 @@ export type MCPTool = { name: string; pluginKey: string; description: string; + /** Raw upstream tool name when the model-facing key stripped a redundant + * server-name prefix — gates the agent editor's legacy id migration. */ + serverToolName?: string; }; export type MCPServer = { From 8ae94afa919cfa3fe739ea2f5b4f97d6e963b65f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 21 Aug 2026 16:33:01 -0400 Subject: [PATCH 4/8] =?UTF-8?q?=F0=9F=AA=A1=20fix:=20Thread=20Parent=20Mes?= =?UTF-8?q?sage=20ID=20Through=20MCP=20Request-Scoped=20Bodies=20(#15095)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: Unify MCP request-scoped headers * fix: address request-scoped MCP review findings * test: preserve request scope on status errors * fix: treat authorized on-demand MCP servers as ready * refactor: separate MCP readiness from connection state * fix: preserve on-demand MCP readiness labels * test: satisfy OpenAI conversation ownership guard * fix: keep MCP action predicates boolean * fix: close deferred MCP request context gaps * fix: preserve on-demand MCP configuration actions * fix: fail closed on unavailable MCP parent context * test: complete MCP connecting-state mocks * fix: preserve missing MCP parent on continuations * fix: align native MCP request identities * fix: preserve edited MCP parent identity * test: use scoped Agent initializer fixture * test: expose MCP request body helper * fix: preserve MCP turn identity across resume * style: sort stream metadata imports * fix: carry normalized MCP identity to execution --- CONTEXT.md | 1 + api/app/clients/BaseClient.js | 12 +- api/app/clients/specs/BaseClient.test.js | 15 +++ api/app/clients/tools/util/handleTools.js | 2 +- .../agents/__tests__/openai.spec.js | 85 +++++++++++- .../request.partialDisconnect.spec.js | 5 + .../__tests__/request.resumeMetadata.spec.js | 126 +++++++++++++++++- .../agents/__tests__/responses.unit.spec.js | 21 ++- .../agents/__tests__/resume.spec.js | 30 ++++- api/server/controllers/agents/client.js | 25 ++-- api/server/controllers/agents/openai.js | 27 ++-- api/server/controllers/agents/request.js | 72 +++++++--- api/server/controllers/agents/responses.js | 28 ++-- api/server/controllers/agents/resume.js | 8 ++ api/server/routes/__tests__/mcp.spec.js | 27 ++++ api/server/routes/mcp.js | 6 + .../services/Endpoints/agents/addedConvo.js | 3 + .../services/Endpoints/agents/initialize.js | 11 ++ .../Endpoints/agents/initialize.spec.js | 28 +++- api/server/services/MCP.js | 22 ++- api/server/services/MCP.spec.js | 55 ++++++++ api/server/services/ToolService.js | 30 ++++- .../services/__tests__/ToolService.spec.js | 32 +++++ client/src/Providers/AgentPanelContext.tsx | 17 ++- client/src/common/types.ts | 2 + .../MCP/MCPServerStatusIcon.spec.tsx | 71 ++++++++++ .../components/MCP/MCPServerStatusIcon.tsx | 17 ++- .../MCP/ServerInitializationSection.spec.tsx | 78 +++++++++++ .../MCP/ServerInitializationSection.tsx | 46 ++++--- .../src/components/MCP/mcpServerUtils.spec.ts | 66 +++++++++ client/src/components/MCP/mcpServerUtils.ts | 55 +++++++- .../ItemDialog/__tests__/McpSection.spec.tsx | 33 ++++- .../Tools/ItemDialog/sections/McpSection.tsx | 68 +++++----- .../Agents/Tools/ToolsMarketplaceDialog.tsx | 4 +- .../__tests__/ToolsMarketplaceDialog.spec.tsx | 6 +- .../MCPBuilder/MCPCardActions.spec.tsx | 55 ++++++++ .../SidePanel/MCPBuilder/MCPCardActions.tsx | 9 +- .../SidePanel/MCPBuilder/MCPServerCard.tsx | 7 +- .../MCPBuilder/MCPStatusBadge.spec.tsx | 43 ++++++ .../SidePanel/MCPBuilder/MCPStatusBadge.tsx | 25 +++- client/src/hooks/MCP/useMCPServerManager.ts | 19 ++- client/src/locales/en/translation.json | 1 + .../src/agents/__tests__/initialize.test.ts | 26 ++++ packages/api/src/agents/discovery.spec.ts | 34 +++++ packages/api/src/agents/discovery.ts | 3 + packages/api/src/agents/handlers.ts | 5 +- packages/api/src/agents/initialize.ts | 18 ++- .../api/src/agents/openai/service.spec.ts | 84 ++++++++++++ packages/api/src/agents/openai/service.ts | 23 +++- .../api/src/mcp/__tests__/request.test.ts | 51 ++++++- .../mcp/__tests__/scope.integration.test.ts | 14 +- packages/api/src/mcp/__tests__/utils.test.ts | 20 +-- packages/api/src/mcp/request.ts | 31 ++++- packages/api/src/mcp/types/index.ts | 3 + packages/api/src/mcp/utils.ts | 50 +++---- .../api/src/stream/GenerationJobManager.ts | 1 + .../stream/__tests__/RedisJobStore.spec.ts | 15 +++ .../api/src/stream/__tests__/startup.spec.ts | 10 ++ .../stream/implementations/RedisJobStore.ts | 1 + .../api/src/stream/interfaces/IJobStore.ts | 4 + packages/api/src/stream/metadata.ts | 3 + packages/api/src/types/stream.ts | 4 + packages/data-provider/src/types/queries.ts | 6 + 63 files changed, 1500 insertions(+), 199 deletions(-) create mode 100644 client/src/components/MCP/MCPServerStatusIcon.spec.tsx create mode 100644 client/src/components/MCP/ServerInitializationSection.spec.tsx create mode 100644 client/src/components/MCP/mcpServerUtils.spec.ts create mode 100644 client/src/components/SidePanel/MCPBuilder/MCPStatusBadge.spec.tsx diff --git a/CONTEXT.md b/CONTEXT.md index 276b1f8ecab..3d27d23cae8 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,6 +1,7 @@ # Domain language - **Agent run envelope**: the versioned, JSON-safe request contract created after ingress authentication and protocol validation but before agent, provider, tool, or MCP initialization. It carries only the validated protocol payload and the minimum trusted principal identifiers. The execution host rehydrates all runtime state from those identifiers. +- **MCP runtime request body**: trusted chat identifiers supplied only while an MCP server handles an agent request. It enables request-scoped header placeholders without retaining user-specific request data on a shared server definition. - **Subagent thread**: a durable, view-only child conversation owned by one parent conversation and subagent identity. A parent agent may continue it by stable `threadId`; each continuation uses a fresh execution lease restored from the canonical child transcript. It is not an ordinary human-writable chat. - **Live subagent task owner**: the one API process holding a detached child execution, its abort controller, and its bounded control queue. Redis may route trusted poll/control envelopes to that owner, but it does not migrate or persist the executor; Mongo persists only the logical child thread and its continuation fence. - **Subagent completion wakeup**: a durable internal `continue` trigger pre-registered before detached child execution so a process crash cannot lose the wakeup. Delivery defers until the child's terminal transcript is persisted, targets the initiating agent and exact parent response branch, carries task metadata rather than child output, waits for the parent generation to settle, and starts the parent turn that collects the result through the existing task store. diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 00e3fc372f6..edfab4d2866 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -347,16 +347,22 @@ class BaseClient { const conversationId = requestConvoId ?? crypto.randomUUID(); const parentMessageId = opts.parentMessageId ?? Constants.NO_PARENT; const userMessageId = - overrideUserMessageId ?? opts.overrideParentMessageId ?? crypto.randomUUID(); - let responseMessageId = opts.responseMessageId ?? crypto.randomUUID(); + opts.preallocatedUserMessageId ?? + overrideUserMessageId ?? + opts.overrideParentMessageId ?? + crypto.randomUUID(); + let responseMessageId = + opts.responseMessageId ?? opts.preallocatedResponseMessageId ?? crypto.randomUUID(); let head = isEdited ? responseMessageId : parentMessageId; this.currentMessages = (await this.loadHistory(conversationId, head)) ?? []; this.conversationId = conversationId; if (isEdited && !isContinued) { - responseMessageId = crypto.randomUUID(); + responseMessageId = opts.preallocatedResponseMessageId ?? crypto.randomUUID(); head = responseMessageId; this.currentMessages[this.currentMessages.length - 1].messageId = head; + } else if (opts.preallocatedResponseMessageId != null) { + responseMessageId = opts.preallocatedResponseMessageId; } if (opts.isRegenerate && responseMessageId.endsWith('_')) { diff --git a/api/app/clients/specs/BaseClient.test.js b/api/app/clients/specs/BaseClient.test.js index b3b900e5c7f..6e78d3431fc 100644 --- a/api/app/clients/specs/BaseClient.test.js +++ b/api/app/clients/specs/BaseClient.test.js @@ -691,6 +691,21 @@ describe('BaseClient', () => { ); }); + it('honors response and user message IDs preallocated before initialization', async () => { + TestClient = initializeFakeClient(apiKey, options, messageHistory); + + const result = await TestClient.handleStartMethods('request-scoped MCP', { + conversationId, + parentMessageId: '3', + preallocatedUserMessageId: 'preallocated-user', + preallocatedResponseMessageId: 'preallocated-response', + }); + + expect(result.userMessage.messageId).toBe('preallocated-user'); + expect(result.responseMessageId).toBe('preallocated-response'); + expect(TestClient.responseMessageId).toBe('preallocated-response'); + }); + it('applies edited reasoning content from its typed payload before regeneration', async () => { const responseMessageId = 'response-with-reasoning'; const newHistory = [ diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 5a85c0d2242..e0f354e3672 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -592,7 +592,7 @@ const loadTools = async ({ user: safeUser, userMCPAuthMap, configServers, - requestBody: options.req?.body, + requestBody: options.requestBody ?? options.req?.body, requestScopedConnections, res: options.res, streamId: options.req?._resumableStreamId || null, diff --git a/api/server/controllers/agents/__tests__/openai.spec.js b/api/server/controllers/agents/__tests__/openai.spec.js index 389d58d8eb6..85c76d488bd 100644 --- a/api/server/controllers/agents/__tests__/openai.spec.js +++ b/api/server/controllers/agents/__tests__/openai.spec.js @@ -125,6 +125,13 @@ jest.mock('@librechat/api', () => ({ buildInitialToolSessions: jest.fn().mockReturnValue(mockInitialSessions), AgentRunEnvelopeError: MockAgentRunEnvelopeError, createAgentRunEnvelope: (...args) => mockCreateAgentRunEnvelope(...args), + createMCPRuntimeRequestBody: ({ messageId, conversationId, parentMessageId }) => ({ + messageId, + conversationId, + ...(parentMessageId !== undefined && { + parentMessageId: parentMessageId ?? '00000000-0000-0000-0000-000000000000', + }), + }), scopeSkillIds: jest.fn().mockImplementation((ids) => ids), resolveAgentScopedSkillIds: jest .fn() @@ -499,7 +506,10 @@ describe('OpenAIChatCompletionController', () => { const toolExecuteOptions = createToolExecuteHandler.mock.calls.at(-1)[0]; await toolExecuteOptions.loadTools(['file_search'], 'agent-123'); expect(loadToolsForExecution).toHaveBeenLastCalledWith( - expect.objectContaining({ agentResourceType: ResourceType.REMOTE_AGENT }), + expect.objectContaining({ + agentResourceType: ResourceType.REMOTE_AGENT, + requestBody: initializeParams.requestBody, + }), ); }); @@ -681,6 +691,79 @@ describe('OpenAIChatCompletionController', () => { }); describe('recursionLimit resolution', () => { + it('threads the OpenAI parent message id through both MCP execution bodies', async () => { + const { validateRequest, createRun, initializeAgent } = require('@librechat/api'); + const { getConvo } = require('~/models'); + validateRequest.mockReturnValueOnce({ + request: { + model: 'agent-123', + messages: [], + stream: false, + conversation_id: 'conversation-123', + parent_message_id: 'parent-123', + }, + }); + getConvo.mockResolvedValueOnce({ conversationId: 'conversation-123', user: 'user-123' }); + + await OpenAIChatCompletionController(req, res); + + expect(initializeAgent).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: { + messageId: 'chatcmpl-mock-nanoid-123', + conversationId: 'conversation-123', + parentMessageId: 'parent-123', + }, + }), + expect.anything(), + ); + expect(createRun).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: { + messageId: 'chatcmpl-mock-nanoid-123', + conversationId: 'conversation-123', + parentMessageId: 'parent-123', + }, + }), + ); + expect(mockProcessStream).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + configurable: expect.objectContaining({ + requestBody: { + messageId: 'chatcmpl-mock-nanoid-123', + conversationId: 'conversation-123', + parentMessageId: 'parent-123', + }, + }), + }), + expect.anything(), + ); + }); + + it('does not synthesize an MCP parent for a continuation that omits it', async () => { + const { validateRequest, initializeAgent } = require('@librechat/api'); + const { getConvo } = require('~/models'); + validateRequest.mockReturnValueOnce({ + request: { + model: 'agent-123', + messages: [], + stream: false, + conversation_id: 'conversation-123', + }, + }); + getConvo.mockResolvedValueOnce({ conversationId: 'conversation-123', user: 'user-123' }); + + await OpenAIChatCompletionController(req, res); + + const requestBody = initializeAgent.mock.calls.at(-1)[0].requestBody; + expect(requestBody).toEqual({ + messageId: 'chatcmpl-mock-nanoid-123', + conversationId: 'conversation-123', + }); + expect(requestBody).not.toHaveProperty('parentMessageId'); + }); + it('should pass resolveRecursionLimit result to processStream config', async () => { const { resolveRecursionLimit } = require('@librechat/api'); resolveRecursionLimit.mockReturnValueOnce(75); diff --git a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js index 7428c2287f1..f0bd8e2d03d 100644 --- a/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js +++ b/api/server/controllers/agents/__tests__/request.partialDisconnect.spec.js @@ -83,6 +83,11 @@ jest.mock('@librechat/api', () => ({ getAgentStartupTelemetry: jest.fn(() => undefined), acceptAgentStartupTelemetry: jest.fn(), isUnpersistedPreliminaryParent: jest.fn(async () => false), + createMCPRuntimeRequestBody: ({ messageId, conversationId, parentMessageId }) => ({ + messageId, + conversationId, + parentMessageId, + }), })); jest.mock('~/server/cleanup', () => ({ diff --git a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js index ca27a838fed..3668c564d3e 100644 --- a/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js +++ b/api/server/controllers/agents/__tests__/request.resumeMetadata.spec.js @@ -195,6 +195,11 @@ jest.mock('@librechat/api', () => ({ return messages.length === 0; }, deleteAgentCheckpoint: (...args) => mockDeleteAgentCheckpoint(...args), + createMCPRuntimeRequestBody: ({ messageId, conversationId, parentMessageId }) => ({ + messageId, + conversationId, + parentMessageId, + }), })); jest.mock('~/server/cleanup', () => ({ @@ -383,6 +388,35 @@ describe('ResumableAgentController resume metadata', () => { }, ); + it.each(['overrideUserMessageId', 'overrideConvoId'])( + 'rejects a non-string %s before admission', + async (field) => { + const req = { + user: { id: 'user-123' }, + body: { + text: 'Invalid override identity', + messageId: 'user-message', + clientRequestId: 'override-request', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + [field]: { malformed: true }, + }, + config: {}, + }; + const res = { json: jest.fn(), status: jest.fn(() => res) }; + + await AgentController(req, res, jest.fn(), jest.fn(), null); + + expect(res.status).toHaveBeenCalledWith(400); + expect(res.json).toHaveBeenCalledWith( + expect.objectContaining({ code: 'INVALID_OVERRIDE_ID' }), + ); + expect(mockGenerationJobManager.claimGeneration).not.toHaveBeenCalled(); + expect(mockGenerationJobManager.createJob).not.toHaveBeenCalled(); + expect(mockCheckAndIncrementPendingRequest).not.toHaveBeenCalled(); + }, + ); + it.each([ ['empty recovery id', { clientRequestId: 'steer-recovery:' }], ['regenerate', { isRegenerate: true }], @@ -695,9 +729,14 @@ describe('ResumableAgentController resume metadata', () => { preemptCapable: true, agent_id: undefined, isTemporary: true, - responseMessageId: 'follow-up-user_', + responseMessageId: expect.stringMatching(/^[0-9a-f-]{36}$/), + mcpRequestBody: { + messageId: expect.stringMatching(/^[0-9a-f-]{36}$/), + conversationId, + parentMessageId: expect.stringMatching(/^[0-9a-f-]{36}$/), + }, userMessage: { - messageId: 'follow-up-user', + messageId: expect.stringMatching(/^[0-9a-f-]{36}$/), parentMessageId: 'original-response', conversationId, text: 'Check Google Workspace availability.', @@ -1023,6 +1062,89 @@ describe('ResumableAgentController resume metadata', () => { ); }); + it('preallocates response-scoped MCP identities before native Agent initialization', async () => { + const initializeClient = jest.fn().mockRejectedValue(new Error('stop after MCP discovery')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Use request-scoped headers.', + messageId: 'incoming-client-message', + parentMessageId: 'previous-response', + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + + await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); + + expect(initializeClient).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: { + messageId: expect.stringMatching(/^[0-9a-f-]{36}$/), + conversationId: 'conversation-123', + parentMessageId: expect.stringMatching(/^[0-9a-f-]{36}$/), + }, + }), + ); + const [{ requestBody }] = initializeClient.mock.calls[0]; + const jobOptions = mockGenerationJobManager.createJob.mock.calls[0][3]; + expect(jobOptions.initialMetadata.responseMessageId).toBe(requestBody.messageId); + expect(jobOptions.initialMetadata.userMessage.messageId).toBe(requestBody.parentMessageId); + expect(jobOptions.initialMetadata.mcpRequestBody).toBe(requestBody); + expect(requestBody.messageId).not.toBe(req.body.messageId); + }); + + it('uses the effective overridden conversation in the MCP request body', async () => { + const initializeClient = jest.fn().mockRejectedValue(new Error('stop after MCP discovery')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Continue in the overridden conversation.', + messageId: 'incoming-client-message', + parentMessageId: 'previous-response', + conversationId: 'source-conversation', + overrideConvoId: 'overridden-conversation__0', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + + await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); + + const [{ requestBody }] = initializeClient.mock.calls[0]; + const jobOptions = mockGenerationJobManager.createJob.mock.calls[0][3]; + expect(requestBody.conversationId).toBe('overridden-conversation'); + expect(jobOptions.initialMetadata.mcpRequestBody).toBe(requestBody); + }); + + it('preallocates the replacement response as the MCP parent for edited content', async () => { + const initializeClient = jest.fn().mockRejectedValue(new Error('stop after MCP discovery')); + const req = { + user: { id: 'user-123' }, + body: { + text: 'Edited response text.', + messageId: 'existing-user-message', + responseMessageId: 'existing-response-message', + parentMessageId: 'previous-response', + overrideParentMessageId: 'existing-user-message', + editedContent: { index: 0, type: 'text', text: 'Edited response text.' }, + conversationId: 'conversation-123', + endpointOption: { endpoint: 'agents', modelOptions: { model: 'gpt-4.1' } }, + }, + config: {}, + }; + + await AgentController(req, createResumableResponse(), jest.fn(), initializeClient, null); + + const [{ requestBody }] = initializeClient.mock.calls[0]; + const jobOptions = mockGenerationJobManager.createJob.mock.calls[0][3]; + expect(requestBody.messageId).toMatch(/^[0-9a-f-]{36}$/); + expect(requestBody.parentMessageId).toBe(requestBody.messageId); + expect(requestBody.messageId).not.toBe('existing-response-message'); + expect(jobOptions.initialMetadata.mcpRequestBody).toBe(requestBody); + }); + it('stores model spec icon fallbacks and agent ids in early resume metadata', async () => { const conversationId = 'conversation-123'; const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading')); diff --git a/api/server/controllers/agents/__tests__/responses.unit.spec.js b/api/server/controllers/agents/__tests__/responses.unit.spec.js index f3500996d6f..b3059be876c 100644 --- a/api/server/controllers/agents/__tests__/responses.unit.spec.js +++ b/api/server/controllers/agents/__tests__/responses.unit.spec.js @@ -122,6 +122,13 @@ jest.mock('@librechat/api', () => ({ buildToolSet: jest.fn().mockReturnValue(new Set()), AgentRunEnvelopeError: MockAgentRunEnvelopeError, createAgentRunEnvelope: (...args) => mockCreateAgentRunEnvelope(...args), + createMCPRuntimeRequestBody: ({ messageId, conversationId, parentMessageId }) => ({ + messageId, + conversationId, + ...(parentMessageId !== undefined && { + parentMessageId: parentMessageId ?? '00000000-0000-0000-0000-000000000000', + }), + }), buildAgentScopedContext: (...args) => mockBuildAgentScopedContext(...args), buildInlineMemoryContext: (...args) => mockBuildInlineMemoryContext(...args), buildAgentContextAttachmentsByAgentId: (...args) => @@ -547,6 +554,15 @@ describe('createResponse controller', () => { expect(mockCreateAgentRunEnvelope.mock.invocationCallOrder[0]).toBeLessThan( initializeAgent.mock.invocationCallOrder[0], ); + expect(initializeAgent).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: { + messageId: 'resp_mock-123', + conversationId: expect.any(String), + }, + }), + expect.anything(), + ); expect(req.body).not.toBe(requestBody); expect(req.body).toEqual(requestBody); expect(JSON.stringify(mockCreateAgentRunEnvelope.mock.results[0].value)).not.toContain( @@ -734,7 +750,10 @@ describe('createResponse controller', () => { const toolExecuteOptions = createToolExecuteHandler.mock.calls.at(-1)[0]; await toolExecuteOptions.loadTools(['file_search'], 'agent-123'); expect(loadToolsForExecution).toHaveBeenLastCalledWith( - expect.objectContaining({ agentResourceType: ResourceType.REMOTE_AGENT }), + expect.objectContaining({ + agentResourceType: ResourceType.REMOTE_AGENT, + requestBody: initializeParams.requestBody, + }), ); }); }); diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index 8fa406d7d6f..71c06a26346 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -103,6 +103,11 @@ jest.mock('@librechat/api', () => ({ decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args), checkAndIncrementPendingRequest: (...args) => mockCheckAndIncrementPendingRequest(...args), isSteerPreemptSupported: jest.fn(() => true), + createMCPRuntimeRequestBody: ({ messageId, conversationId, parentMessageId }) => ({ + messageId, + conversationId, + parentMessageId, + }), })); jest.mock('~/models', () => ({ @@ -292,7 +297,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { }); mockAddTitle = jest.fn().mockResolvedValue(undefined); - mockInitializeClient = jest.fn(async ({ req, checkpointNamespace }) => { + mockInitializeClient = jest.fn(async ({ req, checkpointNamespace, requestBody }) => { // Capture the request state the controller seeds BEFORE reconstruction. capturedInit = { parentMessageId: req.body.parentMessageId, @@ -301,6 +306,7 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { conversationCreatedAt: req.conversationCreatedAt, timezone: req.body.timezone, checkpointNamespace, + requestBody, }; return { client: makeClient(), userMCPAuthMap: { server1: { token: 't' } } }; }); @@ -1195,6 +1201,11 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { // initializeAgent scopes thread files off req.body.parentMessageId, seeded // from the paused user message's parent before initializeClient runs. expect(capturedInit.parentMessageId).toBe(THREAD_PARENT_ID); + expect(capturedInit.requestBody).toEqual({ + messageId: RESPONSE_MSG_ID, + conversationId: CONVO_ID, + parentMessageId: USER_MSG_ID, + }); expect(mockInitializeClient).toHaveBeenCalledTimes(1); const client = await mockInitializeClient.mock.results[0].value.then((r) => r.client); @@ -1206,6 +1217,23 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { ); }); + it('reuses the persisted MCP identity for edited and overridden turns', async () => { + const persistedMCPRequestBody = { + messageId: RESPONSE_MSG_ID, + conversationId: 'overridden-conversation', + parentMessageId: RESPONSE_MSG_ID, + }; + mockGenerationJobManager.getJob.mockResolvedValue( + makeToolApprovalJob({ metadata: { mcpRequestBody: persistedMCPRequestBody } }), + ); + + await post(approveBody()); + await settled; + await flush(); + + expect(capturedInit.requestBody).toBe(persistedMCPRequestBody); + }); + it('reuses the persisted generation checkpoint namespace and keeps legacy fallback explicit', async () => { mockGenerationJobManager.getJob.mockResolvedValue( makeToolApprovalJob({ metadata: { checkpointNamespace: 'generation-1000' } }), diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index bb9201ab899..db5c9f5f614 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -60,6 +60,7 @@ const { createActivityLabelWiring, createActivityPhaseWiring, createReasoningLabelHostWiring, + createMCPRuntimeRequestBody, generateReasoningLabelRevision, getLabelUsageSequenceSeed, createAssistantPhaseStampingHandlers, @@ -2952,11 +2953,13 @@ class AgentClient extends BaseClient { last_agent_index: this.agentConfigs?.size ?? 0, user_id: this.user ?? this.options.req.user?.id, hide_sequential_outputs: this.options.agent.hide_sequential_outputs, - requestBody: { - messageId: this.responseMessageId, - conversationId: this.conversationId, - parentMessageId: this.parentMessageId, - }, + requestBody: + this.options.mcpRequestBody ?? + createMCPRuntimeRequestBody({ + messageId: this.responseMessageId, + conversationId: this.conversationId, + parentMessageId: this.parentMessageId, + }), user: createSafeUser(this.options.req.user), }, recursionLimit: resolveRecursionLimit(agentsEConfig, this.options.agent), @@ -3541,11 +3544,13 @@ class AgentClient extends BaseClient { last_agent_index: this.agentConfigs?.size ?? 0, user_id: this.user ?? this.options.req.user?.id, hide_sequential_outputs: this.options.agent.hide_sequential_outputs, - requestBody: { - messageId: this.responseMessageId, - conversationId: this.conversationId, - parentMessageId: this.parentMessageId, - }, + requestBody: + this.options.mcpRequestBody ?? + createMCPRuntimeRequestBody({ + messageId: this.responseMessageId, + conversationId: this.conversationId, + parentMessageId: this.parentMessageId, + }), user: createSafeUser(this.options.req.user), }, recursionLimit: resolveRecursionLimit(agentsEConfig, this.options.agent), diff --git a/api/server/controllers/agents/openai.js b/api/server/controllers/agents/openai.js index edb1a32c9fc..d3dbb19b71d 100644 --- a/api/server/controllers/agents/openai.js +++ b/api/server/controllers/agents/openai.js @@ -20,6 +20,7 @@ const { buildAgentContextAttachmentsByAgentId, AgentRunEnvelopeError, createAgentRunEnvelope, + createMCPRuntimeRequestBody, loadSkillStates, sendFinalChunk, createSafeUser, @@ -97,6 +98,7 @@ function createToolLoader(signal, definitionsOnly = true) { provider, tool_options, tool_resources, + requestBody, codeExecutionContext, accessibleMcpServerNames, }) { @@ -107,6 +109,7 @@ function createToolLoader(signal, definitionsOnly = true) { res, agent, signal, + requestBody, tool_resources, codeExecutionContext, agentResourceType: ResourceType.REMOTE_AGENT, @@ -255,6 +258,17 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { const conversationId = request.conversation_id ?? nanoid(); const parentMessageId = request.parent_message_id ?? null; + let mcpParentMessageId; + if (typeof request.parent_message_id === 'string' && request.parent_message_id.trim() !== '') { + mcpParentMessageId = request.parent_message_id; + } else if (request.conversation_id == null) { + mcpParentMessageId = null; + } + const mcpRequestBody = createMCPRuntimeRequestBody({ + messageId: responseId, + conversationId, + parentMessageId: mcpParentMessageId, + }); const agentsEConfig = appConfig?.endpoints?.[EModelEndpoint.agents]; const allowedProviders = new Set(agentsEConfig?.allowedProviders); @@ -347,6 +361,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { requestFiles: [], conversationId, parentMessageId, + requestBody: mcpRequestBody, agent, endpointOption, allowedProviders, @@ -414,6 +429,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { requestFiles: [], conversationId, parentMessageId, + requestBody: mcpRequestBody, resourceType: ResourceType.REMOTE_AGENT, computeAccessibleSkillIds: (handoffAgent) => resolveAgentScopedSkillIds({ @@ -553,6 +569,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { res, agentResourceType: ResourceType.REMOTE_AGENT, conversationId, + requestBody: mcpRequestBody, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, @@ -841,10 +858,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { appConfig, signal: abortController.signal, customHandlers: handlers, - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, user: { id: userId }, tenantId: principal.tenantId, /** Bills subagent child-run model calls (reported outside the @@ -862,10 +876,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { thread_id: conversationId, user_id: userId, user: createSafeUser(req.user), - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, ...(userMCPAuthMap != null && { userMCPAuthMap }), }, recursionLimit: resolveRecursionLimit(agentsEConfig, agent), diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index 1dedc6617d0..a28a568bd64 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -28,6 +28,7 @@ const { buildRecoveredSteerPayload, deleteAgentCheckpoint, getAttachmentTitleText, + createMCPRuntimeRequestBody, } = require('@librechat/api'); const { disposeClient } = require('~/server/cleanup'); const { @@ -96,18 +97,6 @@ async function attachConversationCreatedAt(req, conversationId, conversationAnch } } -function getPreliminaryResponseMessageId({ messageId, responseMessageId }) { - if (typeof responseMessageId === 'string' && responseMessageId.length > 0) { - return responseMessageId; - } - - if (typeof messageId !== 'string' || messageId.length === 0) { - return null; - } - - return `${messageId.replace(/_+$/, '')}_`; -} - function getPreliminaryUserMessage( { messageId, parentMessageId, text, quotes, files, manualSkills, alwaysAppliedSkills }, conversationId, @@ -382,6 +371,23 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit ); } const clientRequestId = rawClientRequestId; + const rawOverrideUserMessageId = req.body?.overrideUserMessageId; + const rawOverrideConversationId = req.body?.overrideConvoId; + if ( + (rawOverrideUserMessageId != null && typeof rawOverrideUserMessageId !== 'string') || + (rawOverrideConversationId != null && typeof rawOverrideConversationId !== 'string') + ) { + startupTelemetry?.end('rejected'); + return sendGenerationJson( + res, + 400, + { + code: 'INVALID_OVERRIDE_ID', + error: 'overrideUserMessageId and overrideConvoId must be strings.', + }, + generationProtocolVersion, + ); + } const rawExpectedPredecessorCreatedAt = req.body?.expectedPredecessorCreatedAt; if ( rawExpectedPredecessorCreatedAt != null && @@ -426,7 +432,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } const recoveredSteerId = explicitRecoveredSteerId ?? legacyRecoveredSteerId; const isRecoveredSteerRequest = recoveredSteerId != null; - const recoveryUserMessageId = req.body?.overrideUserMessageId; + const recoveryUserMessageId = rawOverrideUserMessageId; const recoveredSteerPayload = isRecoveredSteerRequest ? buildRecoveredSteerPayload(text, req.body?.files) : undefined; @@ -985,6 +991,34 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit } startupTelemetry?.mark('request_admitted'); + /** Allocate the turn identities before Agent initialization. Request-scoped + * MCP transports resolve BODY placeholders while tools are discovered, so + * discovery and graph execution must receive the same response-scoped body. + * BaseClient otherwise allocates these IDs later in `sendMessage`, after MCP + * connections already exist. */ + const overrideUserMessageId = rawOverrideUserMessageId + ? rawOverrideUserMessageId.split(Constants.COMMON_DIVIDER)[0] + : undefined; + const preallocatedUserMessageId = + overrideUserMessageId ?? overrideParentMessageId ?? crypto.randomUUID(); + const overrideConversationId = rawOverrideConversationId + ? rawOverrideConversationId.split(Constants.COMMON_DIVIDER)[0] + : undefined; + const effectiveConversationId = overrideConversationId ?? conversationId; + let preallocatedResponseMessageId = editedResponseMessageId ?? crypto.randomUUID(); + if ( + (editedContent != null && !isContinued) || + (isRegenerate && preallocatedResponseMessageId.endsWith('_')) + ) { + preallocatedResponseMessageId = crypto.randomUUID(); + } + const mcpRequestBody = createMCPRuntimeRequestBody({ + messageId: preallocatedResponseMessageId, + conversationId: effectiveConversationId, + parentMessageId: + editedContent != null ? preallocatedResponseMessageId : preallocatedUserMessageId, + }); + let client = null; let jobCreatedAt; let providerExecutionId; @@ -1022,8 +1056,10 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit const endpointIconURL = getEndpointIconURL(req, endpointOption); const responseModel = getAgentResponseModel(req, endpointOption); - const preliminaryUserMessage = getPreliminaryUserMessage(req.body, conversationId); - const preliminaryResponseMessageId = getPreliminaryResponseMessageId(req.body); + const preliminaryUserMessage = getPreliminaryUserMessage( + { ...req.body, messageId: preallocatedUserMessageId }, + conversationId, + ); const job = await GenerationJobManager.createJob(streamId, userId, conversationId, { startupTelemetry, ...(recoveredSteerId && { recoveredSteerId }), @@ -1062,7 +1098,8 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit ...(req._isManualScheduledFire === true && { scheduleManual: true }), } : {}), - responseMessageId: preliminaryResponseMessageId, + responseMessageId: preallocatedResponseMessageId, + mcpRequestBody, userMessage: preliminaryUserMessage, }, }); @@ -1246,6 +1283,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit signal: job.abortController.signal, jobCreatedAt, checkpointNamespace: job.metadata?.checkpointNamespace, + requestBody: mcpRequestBody, }); startupTelemetry?.mark('client_initialized'); client = result.client; @@ -1550,6 +1588,8 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit beforeResponsePersistence: claimBeforeResponsePersistence, userMCPAuthMap: result.userMCPAuthMap, responseMessageId: editedResponseMessageId, + preallocatedUserMessageId, + preallocatedResponseMessageId, progressOptions: { res: { write: () => true, diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index 3fd5c45d3d9..021449c7ca1 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -16,6 +16,7 @@ const { buildToolSet, AgentRunEnvelopeError, createAgentRunEnvelope, + createMCPRuntimeRequestBody, buildAgentScopedContext, buildInlineMemoryContext, buildAgentContextAttachmentsByAgentId, @@ -107,6 +108,7 @@ function createToolLoader(signal, definitionsOnly = true) { provider, tool_options, tool_resources, + requestBody, codeExecutionContext, accessibleMcpServerNames, }) { @@ -117,6 +119,7 @@ function createToolLoader(signal, definitionsOnly = true) { res, agent, signal, + requestBody, tool_resources, codeExecutionContext, agentResourceType: ResourceType.REMOTE_AGENT, @@ -389,6 +392,7 @@ const executeResponse = async (envelope, { req, res }) => { const conversationId = request.previous_response_id ?? uuidv4(); const parentMessageId = null; + const mcpRequestBody = createMCPRuntimeRequestBody({ messageId: responseId, conversationId }); const agentsEConfig = appConfig?.endpoints?.[EModelEndpoint.agents]; // Build allowed providers set @@ -482,6 +486,7 @@ const executeResponse = async (envelope, { req, res }) => { requestFiles: [], conversationId, parentMessageId, + requestBody: mcpRequestBody, agent, endpointOption, allowedProviders, @@ -549,6 +554,7 @@ const executeResponse = async (envelope, { req, res }) => { requestFiles: [], conversationId, parentMessageId, + requestBody: mcpRequestBody, resourceType: ResourceType.REMOTE_AGENT, computeAccessibleSkillIds: (handoffAgent) => resolveAgentScopedSkillIds({ @@ -800,6 +806,7 @@ const executeResponse = async (envelope, { req, res }) => { res, agentResourceType: ResourceType.REMOTE_AGENT, conversationId, + requestBody: mcpRequestBody, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, @@ -867,10 +874,7 @@ const executeResponse = async (envelope, { req, res }) => { signal: abortController.signal, customHandlers: handlers, initialSessions, - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, user: { id: userId }, tenantId: principal.tenantId, /** Bills subagent child-run model calls (reported outside the @@ -893,10 +897,7 @@ const executeResponse = async (envelope, { req, res }) => { thread_id: conversationId, user_id: userId, user: createSafeUser(req.user), - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, ...(userMCPAuthMap != null && { userMCPAuthMap }), }, signal: abortController.signal, @@ -992,6 +993,7 @@ const executeResponse = async (envelope, { req, res }) => { res, agentResourceType: ResourceType.REMOTE_AGENT, conversationId, + requestBody: mcpRequestBody, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, @@ -1057,10 +1059,7 @@ const executeResponse = async (envelope, { req, res }) => { signal: abortController.signal, customHandlers: handlers, initialSessions, - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, user: { id: userId }, tenantId: principal.tenantId, /** Bills subagent child-run model calls (reported outside the @@ -1082,10 +1081,7 @@ const executeResponse = async (envelope, { req, res }) => { thread_id: conversationId, user_id: userId, user: createSafeUser(req.user), - requestBody: { - messageId: responseId, - conversationId, - }, + requestBody: mcpRequestBody, ...(userMCPAuthMap != null && { userMCPAuthMap }), }, signal: abortController.signal, diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index b2a4f450f11..5a8233be684 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -24,6 +24,7 @@ const { isSteerPreemptSupported, isStopConfirmed, toPendingSteer, + createMCPRuntimeRequestBody, } = require('@librechat/api'); const { disposeClient } = require('~/server/cleanup'); const { @@ -1190,6 +1191,13 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) signal: job.abortController.signal, jobCreatedAt: job.createdAt, checkpointNamespace, + requestBody: + job.metadata.mcpRequestBody ?? + createMCPRuntimeRequestBody({ + messageId: job.metadata.responseMessageId, + conversationId: streamId, + parentMessageId: job.metadata.userMessage?.messageId ?? Constants.NO_PARENT, + }), }); client = result.client; diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 99ebff511a7..9fcb49a1aa4 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -2753,6 +2753,33 @@ describe('MCP Routes', () => { expect(getServerConnectionStatus).toHaveBeenCalledTimes(2); }); + it('preserves request-scoped metadata when an individual status check fails', async () => { + getMCPSetupData.mockResolvedValue({ + mcpConfig: { + server1: { + source: 'config', + headers: { 'X-Message': '{{LIBRECHAT_BODY_MESSAGEID}}' }, + customUserVars: { API_KEY: { title: 'API key' } }, + }, + }, + appConnections: new Map(), + userConnections: new Map(), + oauthServers: new Set(), + }); + getServerConnectionStatus.mockRejectedValueOnce(new Error('status unavailable')); + + const response = await request(app).get('/api/mcp/connection/status'); + + expect(response.status).toBe(200); + expect(response.body.connectionStatus.server1).toEqual( + expect.objectContaining({ + connectionState: 'error', + requestScoped: true, + configurationState: 'needs_configuration', + }), + ); + }); + it('should return 500 when connection status check fails', async () => { getMCPSetupData.mockRejectedValue(new Error('Database error')); diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index 554d1104052..32bca863d8e 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -23,6 +23,7 @@ const { OAUTH_SESSION_COOKIE, mcpConfig: mcpSettings, getServerCustomUserVars, + hasCustomUserVars, requiresEphemeralUserConnection, } = require('@librechat/api'); const { @@ -926,6 +927,9 @@ router.get('/connection/status', requireJwtAuth, async (req, res) => { { connectionState: 'error', requiresOAuth: oauthServers.has(serverName), + ...(requiresEphemeralUserConnection(config) && { requestScoped: true }), + ...(requiresEphemeralUserConnection(config) && + hasCustomUserVars(config) && { configurationState: 'needs_configuration' }), authorizationState: oauthServers.has(serverName) ? 'error' : 'not_required', error: message, }, @@ -987,6 +991,8 @@ router.get('/connection/status/:serverName', requireJwtAuth, async (req, res) => serverName, connectionStatus: serverStatus.connectionState, requiresOAuth: serverStatus.requiresOAuth, + requestScoped: serverStatus.requestScoped, + configurationState: serverStatus.configurationState, authorizationState: serverStatus.authorizationState, }); } catch (error) { diff --git a/api/server/services/Endpoints/agents/addedConvo.js b/api/server/services/Endpoints/agents/addedConvo.js index dee208f3fb8..9483b1535f0 100644 --- a/api/server/services/Endpoints/agents/addedConvo.js +++ b/api/server/services/Endpoints/agents/addedConvo.js @@ -42,6 +42,7 @@ const loadAddedAgent = (params) => * @param {Array} params.requestFiles - Request files * @param {string} params.conversationId - The conversation ID * @param {string} [params.parentMessageId] - The parent message ID for thread filtering + * @param {import('@librechat/api').MCPRuntimeRequestBody} [params.requestBody] * @param {Set} params.allowedProviders - Set of allowed providers * @param {Map} params.agentConfigs - Map of agent configs to add to * @param {string} params.primaryAgentId - The primary agent ID @@ -70,6 +71,7 @@ const processAddedConvo = async ({ requestFiles, conversationId, parentMessageId, + requestBody, allowedProviders, agentConfigs, primaryAgentId, @@ -170,6 +172,7 @@ const processAddedConvo = async ({ requestFiles, conversationId, parentMessageId, + requestBody, agent: addedAgent, endpointOption, allowedProviders, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index ded114cdead..4e2dfd031ff 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -102,6 +102,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false, jobC provider, tool_options, tool_resources, + requestBody, codeExecutionContext, accessibleMcpServerNames, }) { @@ -114,6 +115,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false, jobC signal, streamId, jobCreatedAt, + requestBody, tool_resources, codeExecutionContext, definitionsOnly, @@ -137,6 +139,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false, jobC * @param {Object} params.endpointOption * @param {number} [params.jobCreatedAt] * @param {string} [params.checkpointNamespace] Immutable saver-level generation scope + * @param {import('@librechat/api').MCPRuntimeRequestBody} [params.requestBody] */ const initializeClient = async ({ req, @@ -145,6 +148,7 @@ const initializeClient = async ({ endpointOption, jobCreatedAt, checkpointNamespace, + requestBody, }) => { if (!endpointOption) { throw new Error('Endpoint option not provided'); @@ -154,6 +158,7 @@ const initializeClient = async ({ * that trusted document for child-thread execution policy; resume and direct * callers fall back to the same owner-scoped lookup. */ const conversationId = req.body?.conversationId; + const runtimeRequestBody = requestBody ?? req.body; let requestConversationPromise = Promise.resolve(null); if (Object.prototype.hasOwnProperty.call(req, 'resolvedConversation')) { requestConversationPromise = Promise.resolve(req.resolvedConversation); @@ -334,6 +339,7 @@ const initializeClient = async ({ signal, streamId, conversationId, + requestBody: runtimeRequestBody, toolNames, agent: ctx.agent, toolRegistry: ctx.toolRegistry, @@ -496,6 +502,7 @@ const initializeClient = async ({ requestFiles, conversationId, parentMessageId, + requestBody: runtimeRequestBody, agent: primaryAgent, endpointOption, allowedProviders, @@ -561,6 +568,7 @@ const initializeClient = async ({ requestFiles, conversationId, parentMessageId, + requestBody: runtimeRequestBody, computeAccessibleSkillIds: (agent) => resolveAgentScopedSkillIds({ agent, @@ -651,6 +659,7 @@ const initializeClient = async ({ userMCPAuthMap, conversationId, parentMessageId, + requestBody: runtimeRequestBody, allowedProviders, primaryAgentId: primaryConfig.id, accessibleSkillIds, @@ -940,6 +949,7 @@ const initializeClient = async ({ requestFiles, conversationId, parentMessageId, + requestBody: runtimeRequestBody, endpointOption: { ...endpointOption, endpoint: EModelEndpoint.agents }, allowedProviders, accessibleSkillIds: scopedSkillIds, @@ -1417,6 +1427,7 @@ const initializeClient = async ({ toolInputValidationErrors, jobCreatedAt, checkpointNamespace, + mcpRequestBody: runtimeRequestBody, }); if (streamId) { diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 59910aa2a0c..927291b6b19 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -617,7 +617,7 @@ describe('initializeClient — subagent loading', () => { agentClientArgs = undefined; capturedToolExecuteOptions = undefined; mockLoadToolsForExecution.mockReset(); - mockLoadToolsForExecution.mockResolvedValue({ loadedTools: [] }); + mockLoadToolsForExecution.mockResolvedValue({ loadedTools: [], configurable: {} }); testUser = await User.create({ email: 'subagent@example.com', @@ -741,6 +741,32 @@ describe('initializeClient — subagent loading', () => { }); }); + it('uses one normalized MCP body for discovery, deferred execution, and AgentClient', async () => { + const requestBody = Object.freeze({ + messageId: 'response-message', + conversationId: 'conv_sub', + parentMessageId: 'user-message', + }); + mockInitializeAgent.mockResolvedValue(makePrimaryConfig({})); + mockLoadToolsForExecution.mockResolvedValue({ loadedTools: [], configurable: {} }); + + await initializeClient({ + req: makeSubagentReq(), + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + requestBody, + }); + + expect(mockInitializeAgent.mock.calls[0][0].requestBody).toBe(requestBody); + expect(agentClientArgs.mcpRequestBody).toBe(requestBody); + + await capturedToolExecuteOptions.loadTools([], PRIMARY_ID); + expect(mockLoadToolsForExecution).toHaveBeenCalledWith( + expect.objectContaining({ requestBody }), + ); + }); + it('keeps an existing detached task controllable after subagent config is disabled', async () => { mockInitializeAgent.mockResolvedValue( makePrimaryConfig({ diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 15bbecc025a..359aa8979b9 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -29,6 +29,7 @@ const { buildMCPAuthRunStepDeltaEvent, buildMCPAuthRunStepEndDeltaEvent, isUserSourced, + hasCustomUserVars, checkAccessWithRequestCache, getMissingCustomUserVars, getUserMCPAuthMap, @@ -1461,6 +1462,19 @@ async function hasDurableMCPAuthorization(userId, serverName, config, runtimeCon }); } +async function getMCPUserConfigurationState(serverName, config, runtimeContext = {}) { + if (!hasCustomUserVars(config)) { + return undefined; + } + + const userMCPAuthMap = + runtimeContext.userMCPAuthMap ?? (await runtimeContext.loadUserMCPAuthMap?.()); + const customUserVars = getServerCustomUserVars(userMCPAuthMap, serverName); + return getMissingCustomUserVars(config, customUserVars).length > 0 + ? 'needs_configuration' + : 'configured'; +} + function canDetectMCPRuntimeOAuth(config) { return config.requiresOAuth == null && config.apiKey == null && hasRuntimeUrlPlaceholders(config); } @@ -1474,7 +1488,7 @@ function canDetectMCPRuntimeOAuth(config) { * @param {Map} userConnections - User-level connections * @param {Set} oauthServers - Set of OAuth servers * @param {{ user?: Partial, userMCPAuthMap?: Record>, loadUserMCPAuthMap?: () => Promise> | undefined>, loadMCPAllowlists?: () => Promise<{ allowedDomains?: string[] | null, allowedAddresses?: string[] | null }> }} [runtimeContext] - * @returns {Object} Object containing requiresOAuth and connectionState + * @returns {Object} Object containing requiresOAuth, requestScoped, connectionState, and authorizationState */ async function getServerConnectionStatus( userId, @@ -1491,6 +1505,10 @@ async function getServerConnectionStatus( const liveConnectionOAuth = connection?.usesOAuth?.() === true; const runtimeOAuthCandidate = canDetectMCPRuntimeOAuth(config); const effectiveOAuth = configuredOAuth || liveConnectionOAuth; + const requestScoped = requiresEphemeralUserConnection(config); + const configurationState = requestScoped + ? await getMCPUserConfigurationState(serverName, config, runtimeContext) + : undefined; const baseConnectionState = isStaleOrDoNotExist ? 'disconnected' @@ -1535,6 +1553,8 @@ async function getServerConnectionStatus( return { requiresOAuth, + ...(requestScoped && { requestScoped: true }), + ...(configurationState && { configurationState }), connectionState: finalConnectionState, authorizationState, }; diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index c73fa2968b5..36a1ac57edb 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -546,6 +546,60 @@ describe('tests for the new helper functions used by the MCP connection status e }); }); + it('marks BODY placeholder servers as request-scoped while they are idle', async () => { + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + { + ...mockConfig, + source: 'yaml', + headers: { 'X-Parent-Message': '{{LIBRECHAT_BODY_PARENTMESSAGEID}}' }, + }, + new Map(), + new Map(), + new Set(), + ); + + expect(result).toEqual({ + requiresOAuth: false, + requestScoped: true, + connectionState: 'disconnected', + authorizationState: 'not_required', + }); + }); + + it('reports whether custom variables are configured for request-scoped servers', async () => { + const config = { + ...mockConfig, + source: 'yaml', + headers: { 'X-Conversation': '{{LIBRECHAT_BODY_CONVERSATIONID}}' }, + customUserVars: { API_KEY: { title: 'API key' } }, + }; + const connectionArgs = [new Map(), new Map(), new Set()]; + + const missing = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + ...connectionArgs, + { userMCPAuthMap: {} }, + ); + const configured = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + ...connectionArgs, + { + userMCPAuthMap: { + [`${Constants.mcp_prefix}${mockServerName}`]: { API_KEY: 'secret' }, + }, + }, + ); + + expect(missing.configurationState).toBe('needs_configuration'); + expect(configured.configurationState).toBe('configured'); + }); + it('should prioritize app connection over user connection', async () => { const appConnections = new Map([ [ @@ -871,6 +925,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, + requestScoped: true, connectionState: 'connecting', authorizationState: 'authorizing', }); diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 81aae5aef93..efe9d535e69 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -259,6 +259,7 @@ async function processRequiredActions(client, requiredActions) { options: { processFileURL, req: client.req, + res: client.res, uploadImageBuffer, openAIApiKey: client.apiKey, returnMetadata: true, @@ -565,6 +566,7 @@ const isBuiltInTool = (toolName) => * @param {ServerRequest} params.req - The request object * @param {ServerResponse} [params.res] - The response object for SSE events * @param {Object} params.agent - The agent configuration + * @param {import('@librechat/api').RequestBody} [params.requestBody] - Normalized MCP body * @param {string} [params.agentResourceType] - Permission resource type for the authorized agent route * @param {string|null} [params.streamId] - Stream ID for resumable mode * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events @@ -580,6 +582,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, + requestBody, agentResourceType, streamId = null, jobCreatedAt, @@ -599,6 +602,7 @@ async function loadToolDefinitionsWrapper({ } const appConfig = req.config; + const runtimeRequestBody = requestBody ?? req.body; const hasExpectedMCPTools = agent.tools.some(isExpectedMCPTool); const enabledCapabilities = await resolveAgentCapabilities(req, appConfig, agent.id); @@ -620,7 +624,7 @@ async function loadToolDefinitionsWrapper({ environment: agent.stateful_code_environment, userId: req.user.id, agentId: agent.id, - conversationId: req.body?.conversationId, + conversationId: runtimeRequestBody?.conversationId, }); const hasMCPTools = agent.tools?.some((tool) => tool?.includes(Constants.mcp_delimiter)); const mcpPermissionContext = createMCPPermissionContext(req); @@ -927,7 +931,7 @@ async function loadToolDefinitionsWrapper({ serverName, configServers, userMCPAuthMap, - requestBody: req.body, + requestBody: runtimeRequestBody, requestScopedConnections, }); @@ -954,7 +958,7 @@ async function loadToolDefinitionsWrapper({ serverName, configServers, userMCPAuthMap, - requestBody: req.body, + requestBody: runtimeRequestBody, requestScopedConnections, }); @@ -1082,7 +1086,7 @@ async function loadToolDefinitionsWrapper({ configServers, userMCPAuthMap, flowManager, - requestBody: req.body, + requestBody: runtimeRequestBody, returnOnOAuth: false, oauthStart, oauthEnd: createOAuthEndEmitter(serverName), @@ -1255,6 +1259,7 @@ async function loadToolDefinitionsWrapper({ * @param {ServerRequest} params.req - The request object * @param {ServerResponse} params.res - The response object * @param {Object} params.agent - The agent configuration + * @param {import('@librechat/api').RequestBody} [params.requestBody] - Normalized MCP body * @param {string} [params.agentResourceType] - Permission resource type for the authorized agent route * @param {AbortSignal} [params.signal] - Abort signal * @param {Object} [params.tool_resources] - Tool resources @@ -1269,6 +1274,7 @@ async function loadAgentTools({ req, res, agent, + requestBody, agentResourceType, signal, tool_resources, @@ -1285,6 +1291,7 @@ async function loadAgentTools({ req, res, agent, + requestBody, agentResourceType, streamId, jobCreatedAt, @@ -1402,7 +1409,7 @@ async function loadAgentTools({ environment: agent.stateful_code_environment, userId: req.user.id, agentId: agent.id, - conversationId: req.body?.conversationId, + conversationId: requestBody?.conversationId ?? req.body?.conversationId, }); const { loadedTools, toolContextMap, dynamicToolContextMap, primedCodeFiles } = await loadTools({ @@ -1415,6 +1422,7 @@ async function loadAgentTools({ options: { req, res, + requestBody, agentResourceType, mcpServerContext, jobCreatedAt, @@ -1676,6 +1684,7 @@ async function loadAgentTools({ * @param {ServerResponse} params.res - The response object * @param {AbortSignal} [params.signal] - Abort signal * @param {Object} params.agent - The agent object + * @param {import('@librechat/api').RequestBody} [params.requestBody] - Normalized MCP body * @param {string} [params.agentResourceType] - Permission resource type for the authorized agent route * @param {string[]} params.toolNames - Names of tools to load * @param {Map} [params.toolRegistry] - Tool registry @@ -1695,6 +1704,7 @@ async function loadToolsForExecution({ res, signal, agent, + requestBody, agentResourceType, toolNames, toolRegistry, @@ -1712,8 +1722,13 @@ async function loadToolsForExecution({ }) { const appConfig = req.config; const allLoadedTools = []; + const runtimeRequestBody = requestBody ?? req.body; const mcpRequestScopedConnections = requestScopedConnections ?? getMCPRequestContext(req, res); - const configurable = { userMCPAuthMap, requestScopedConnections: mcpRequestScopedConnections }; + const configurable = { + userMCPAuthMap, + requestBody: runtimeRequestBody, + requestScopedConnections: mcpRequestScopedConnections, + }; /** Per-agent set of tools that received the injected `run_in_background` * param; the event-driven executor gates background dispatch and the * `check_background_task` poll tool on this reliable per-agent channel. */ @@ -1770,7 +1785,7 @@ async function loadToolsForExecution({ environment: agent?.stateful_code_environment, userId: req.user.id, agentId: agent?.id, - conversationId: conversationId ?? req.body?.conversationId, + conversationId: conversationId ?? runtimeRequestBody?.conversationId, }); configurable.codeExecutionContext = codeExecutionContext; @@ -1900,6 +1915,7 @@ async function loadToolsForExecution({ options: { req, res, + requestBody: runtimeRequestBody, agentResourceType, jobCreatedAt, tool_resources, diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index 7b2cb745422..04e35ca28c7 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -1429,6 +1429,33 @@ describe('ToolService - Action Capability Gating', () => { ); }); + it('threads the normalized MCP body through deferred tool loading', async () => { + const capabilities = [AgentCapabilities.tools, AgentCapabilities.web_search]; + const req = createMockReq(capabilities); + const requestBody = { + messageId: 'response-1', + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + }; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + + const result = await loadToolsForExecution({ + req, + res: {}, + requestBody, + agent: { id: 'agent_123', tools: [Tools.web_search] }, + toolNames: [Tools.web_search], + actionsEnabled: false, + }); + + expect(mockLoadToolsUtil).toHaveBeenCalledWith( + expect.objectContaining({ + options: expect.objectContaining({ requestBody }), + }), + ); + expect(result.configurable.requestBody).toBe(requestBody); + }); + const actionToolName = `get_weather${actionDelimiter}api_example_com`; const regularTool = Tools.web_search; @@ -2147,6 +2174,11 @@ describe('ToolService - Action Capability Gating', () => { // zodSchema, name, and description for assistants API"), so key // resolution assertions off the request builder path instead. expect(mockCreateActionTool).toHaveBeenCalledTimes(2); + expect(mockLoadToolsUtil).toHaveBeenCalledWith( + expect.objectContaining({ + options: expect.objectContaining({ res: client.res }), + }), + ); const builderPaths = mockCreateActionTool.mock.calls.map((c) => c[0].requestBuilder?.path); expect(builderPaths).toEqual(expect.arrayContaining(['/echo', '/items'])); // Each call must carry a distinct builder — guards against the bug diff --git a/client/src/Providers/AgentPanelContext.tsx b/client/src/Providers/AgentPanelContext.tsx index 493ef03a935..afb3a1e0df8 100644 --- a/client/src/Providers/AgentPanelContext.tsx +++ b/client/src/Providers/AgentPanelContext.tsx @@ -14,6 +14,7 @@ import { useMCPConnectionStatus, useMCPServerManager, } from '~/hooks'; +import { isMCPServerReadyForAgent } from '~/components/MCP/mcpServerUtils'; import { Panel, isEphemeralAgent } from '~/common'; const AgentPanelContext = createContext(undefined); @@ -71,6 +72,7 @@ export function AgentPanelProvider({ children }: { children: React.ReactNode }) for (const [serverName, serverData] of Object.entries(mcpData.servers)) { // Get title and description from config with fallbacks const serverConfig = availableMCPServersMap?.[serverName]; + const serverStatus = connectionStatus?.[serverName]; const displayName = serverConfig?.title || serverName; const displayDescription = serverConfig?.description || `${localize('com_ui_tool_collection_prefix')} ${serverName}`; @@ -98,7 +100,12 @@ export function AgentPanelProvider({ children }: { children: React.ReactNode }) serverName, tools, isConfigured: configuredServers.has(serverName), - isConnected: connectionStatus?.[serverName]?.connectionState === 'connected', + isConnected: serverStatus?.connectionState === 'connected', + isReadyForAgent: isMCPServerReadyForAgent( + serverStatus, + serverConfig?.requestScoped === true, + Object.keys(serverConfig?.customUserVars ?? {}).length > 0, + ), requestScoped: serverConfig?.requestScoped, metadata, consumeOnly: serverConfig?.consumeOnly, @@ -113,6 +120,7 @@ export function AgentPanelProvider({ children }: { children: React.ReactNode }) } // Get title and description from config with fallbacks const serverConfig = availableMCPServersMap?.[mcpServerName]; + const serverStatus = connectionStatus?.[mcpServerName]; const displayName = serverConfig?.title || mcpServerName; const displayDescription = serverConfig?.description || @@ -130,7 +138,12 @@ export function AgentPanelProvider({ children }: { children: React.ReactNode }) metadata, isConfigured: true, serverName: mcpServerName, - isConnected: connectionStatus?.[mcpServerName]?.connectionState === 'connected', + isConnected: serverStatus?.connectionState === 'connected', + isReadyForAgent: isMCPServerReadyForAgent( + serverStatus, + serverConfig?.requestScoped === true, + Object.keys(serverConfig?.customUserVars ?? {}).length > 0, + ), requestScoped: serverConfig?.requestScoped, consumeOnly: serverConfig?.consumeOnly, }); diff --git a/client/src/common/types.ts b/client/src/common/types.ts index a86ca07d70f..28d2872217c 100644 --- a/client/src/common/types.ts +++ b/client/src/common/types.ts @@ -211,6 +211,8 @@ export interface MCPServerInfo { tools: t.AgentToolType[]; isConfigured: boolean; isConnected: boolean; + /** True when the server can be attached to an agent, even if its transport is request-scoped. */ + isReadyForAgent?: boolean; /** True when tools can only be discovered with live chat request fields. */ requestScoped?: boolean; consumeOnly?: boolean; diff --git a/client/src/components/MCP/MCPServerStatusIcon.spec.tsx b/client/src/components/MCP/MCPServerStatusIcon.spec.tsx new file mode 100644 index 00000000000..442ed57a6ed --- /dev/null +++ b/client/src/components/MCP/MCPServerStatusIcon.spec.tsx @@ -0,0 +1,71 @@ +import React from 'react'; +import '@testing-library/jest-dom/extend-expect'; +import { render, screen } from '@testing-library/react'; +import type { MCPServerStatus } from 'librechat-data-provider'; +import MCPServerStatusIcon from './MCPServerStatusIcon'; + +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string) => key, +})); + +jest.mock('@librechat/client', () => ({ + Spinner: (props: React.ComponentProps<'span'>) => , + TooltipAnchor: ({ render }: { render: React.ReactNode }) => render, + Button: ({ + children, + variant: _variant, + size: _size, + ...props + }: React.ComponentProps<'button'> & { variant?: string; size?: string }) => ( + + ), +})); + +const requestScopedStatus: MCPServerStatus = { + connectionState: 'disconnected', + authorizationState: 'not_required', + requiresOAuth: false, + requestScoped: true, + configurationState: 'needs_configuration', +}; + +describe('MCPServerStatusIcon', () => { + it('shows Configure instead of Connect for idle request-scoped custom variables', () => { + render( + , + ); + + expect( + screen.getByRole('button', { name: 'com_nav_mcp_configure_server' }), + ).toBeInTheDocument(); + expect( + screen.queryByRole('button', { name: 'com_nav_mcp_connect_server' }), + ).not.toBeInTheDocument(); + }); + + it('keeps idle request-scoped servers without custom variables actionless', () => { + const { container } = render( + , + ); + + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/client/src/components/MCP/MCPServerStatusIcon.tsx b/client/src/components/MCP/MCPServerStatusIcon.tsx index 01ee2f6ebbb..8d6b0b01bd3 100644 --- a/client/src/components/MCP/MCPServerStatusIcon.tsx +++ b/client/src/components/MCP/MCPServerStatusIcon.tsx @@ -71,7 +71,7 @@ export default function MCPServerStatusIcon({ return null; } - const { connectionState } = serverStatus; + const { connectionState, requestScoped } = serverStatus; // Connecting: show spinner, with cancel when an OAuth flow is pending. if (connectionState === 'connecting') { @@ -89,6 +89,13 @@ export default function MCPServerStatusIcon({ return ; } + // Request-scoped servers can only be connected while serving an MCP request. + if ((connectionState === 'disconnected' || connectionState === 'error') && requestScoped) { + return hasCustomUserVars ? ( + + ) : null; + } + // Disconnected or Error: show connect button (PlugZap icon) if (connectionState === 'disconnected' || connectionState === 'error') { return ; @@ -126,10 +133,12 @@ function CompactStatusDot({ serverStatus, isInitializing }: CompactStatusDotProp const { connectionState, requiresOAuth } = serverStatus; let colorClass = 'bg-status-neutral'; - if (connectionState === 'connected') { - colorClass = 'bg-status-success'; - } else if (connectionState === 'connecting') { + if (connectionState === 'connecting') { + colorClass = 'bg-status-info'; + } else if (serverStatus.requestScoped) { colorClass = 'bg-status-info'; + } else if (connectionState === 'connected') { + colorClass = 'bg-status-success'; } else if (connectionState === 'error') { colorClass = 'bg-status-error'; } else if (connectionState === 'disconnected' && requiresOAuth) { diff --git a/client/src/components/MCP/ServerInitializationSection.spec.tsx b/client/src/components/MCP/ServerInitializationSection.spec.tsx new file mode 100644 index 00000000000..d2281f442bd --- /dev/null +++ b/client/src/components/MCP/ServerInitializationSection.spec.tsx @@ -0,0 +1,78 @@ +import React from 'react'; +import '@testing-library/jest-dom/extend-expect'; +import { fireEvent, render, screen } from '@testing-library/react'; +import type { MCPServerStatus } from 'librechat-data-provider'; +import ServerInitializationSection from './ServerInitializationSection'; + +const mockInitializeServer = jest.fn(); +const mockConnectionStatus = jest.fn((): MCPServerStatus | undefined => undefined); + +jest.mock('~/hooks', () => ({ + useLocalize: () => (key: string) => key, + useMCPConnectionStatus: () => ({ + connectionStatus: { server: mockConnectionStatus() }, + }), + useMCPServerManager: () => ({ + getOAuthUrl: () => undefined, + isCancellable: () => false, + isInitializing: () => false, + cancelOAuthFlow: jest.fn(), + initializeServer: mockInitializeServer, + availableMCPServers: [{ serverName: 'server' }], + availableMCPServersMap: { server: { requestScoped: true } }, + revokeOAuthForServer: jest.fn(), + }), +})); + +jest.mock('@librechat/client', () => ({ + Spinner: (props: React.ComponentProps<'span'>) => , + Button: ({ + children, + variant: _variant, + size: _size, + ...props + }: React.ComponentProps<'button'> & { variant?: string; size?: string }) => ( + + ), +})); + +describe('ServerInitializationSection', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('offers deferred initialization after request-scoped custom variables are configured', () => { + mockConnectionStatus.mockReturnValue({ + connectionState: 'disconnected', + requiresOAuth: false, + requestScoped: true, + configurationState: 'needs_configuration', + }); + const { rerender } = render( + , + ); + + expect(screen.queryByRole('button', { name: 'com_ui_mcp_initialize' })).not.toBeInTheDocument(); + + mockConnectionStatus.mockReturnValue({ + connectionState: 'disconnected', + requiresOAuth: false, + requestScoped: true, + configurationState: 'configured', + }); + rerender( + , + ); + + fireEvent.click(screen.getByRole('button', { name: 'com_ui_mcp_initialize' })); + expect(mockInitializeServer).toHaveBeenCalledWith('server', false); + }); +}); diff --git a/client/src/components/MCP/ServerInitializationSection.tsx b/client/src/components/MCP/ServerInitializationSection.tsx index bbed167da6c..f42aa15a244 100644 --- a/client/src/components/MCP/ServerInitializationSection.tsx +++ b/client/src/components/MCP/ServerInitializationSection.tsx @@ -29,6 +29,7 @@ export default function ServerInitializationSection({ cancelOAuthFlow, initializeServer, availableMCPServers, + availableMCPServersMap, revokeOAuthForServer, } = useMCPServerManager({ conversationId, storageContextKey }); @@ -46,10 +47,23 @@ export default function ServerInitializationSection({ const isServerInitializing = isInitializing(serverName); const serverOAuthUrl = getOAuthUrl(serverName); - const shouldShowReinit = isConnected && (requiresOAuth || hasCustomUserVars); - const shouldShowInit = !isConnected && !serverOAuthUrl && !hasPendingOAuth; + const requestScoped = + serverStatus?.requestScoped === true || + availableMCPServersMap?.[serverName]?.requestScoped === true; + const shouldShowReinit = isConnected && !requestScoped && (requiresOAuth || hasCustomUserVars); + /** Saving custom variables makes an on-demand server ready, but it still + * needs one explicit initialization attempt so callers waiting to attach the + * runtime wildcard observe `connectionDeferred`. */ + const canDeferRequestScopedConnection = + requestScoped && hasCustomUserVars && serverStatus?.configurationState === 'configured'; + const shouldShowInit = + !isConnected && + (!requestScoped || canDeferRequestScopedConnection) && + !serverOAuthUrl && + !hasPendingOAuth; + const shouldShowRevoke = requiresOAuth && revokeOAuthForServer != null; - if (!shouldShowReinit && !shouldShowInit && !serverOAuthUrl) { + if (!shouldShowReinit && !shouldShowInit && !shouldShowRevoke && !serverOAuthUrl) { if (!hasPendingOAuth) { return null; } @@ -114,27 +128,29 @@ export default function ServerInitializationSection({ return (
- {requiresOAuth && revokeOAuthForServer && ( + {shouldShowRevoke && ( )} - + {(shouldShowReinit || shouldShowInit) && ( + + )}
); } diff --git a/client/src/components/MCP/mcpServerUtils.spec.ts b/client/src/components/MCP/mcpServerUtils.spec.ts new file mode 100644 index 00000000000..8425de838aa --- /dev/null +++ b/client/src/components/MCP/mcpServerUtils.spec.ts @@ -0,0 +1,66 @@ +import type { MCPServerStatus } from 'librechat-data-provider'; +import { isMCPServerReadyForAgent, shouldShowActionButton } from './mcpServerUtils'; + +const status = ( + connectionState: MCPServerStatus['connectionState'], + authorizationState: MCPServerStatus['authorizationState'], +): MCPServerStatus => ({ connectionState, authorizationState, requiresOAuth: false }); + +describe('isMCPServerReadyForAgent', () => { + it('treats an authorized idle request-scoped server as ready', () => { + expect(isMCPServerReadyForAgent(status('disconnected', 'authorized'), true)).toBe(true); + }); + + it('treats an idle request-scoped server without an auth requirement as ready', () => { + expect(isMCPServerReadyForAgent(status('disconnected', 'not_required'), true)).toBe(true); + }); + + it('keeps request-scoped servers gated while authorization is incomplete or failed', () => { + expect(isMCPServerReadyForAgent(status('disconnected', 'needs_authorization'), true)).toBe( + false, + ); + expect(isMCPServerReadyForAgent(status('error', 'error'), true)).toBe(false); + }); + + it('requires declared custom variables before an on-demand server is ready', () => { + const missingConfiguration = { + ...status('disconnected', 'not_required'), + configurationState: 'needs_configuration' as const, + }; + const configured = { + ...missingConfiguration, + configurationState: 'configured' as const, + }; + + expect(isMCPServerReadyForAgent(missingConfiguration, true, true)).toBe(false); + expect(isMCPServerReadyForAgent(configured, true, true)).toBe(true); + }); + + it('requires a live connection for servers that are not request-scoped', () => { + expect(isMCPServerReadyForAgent(status('disconnected', 'not_required'), false)).toBe(false); + expect(isMCPServerReadyForAgent(status('connected', 'not_required'), false)).toBe(true); + }); +}); + +describe('shouldShowActionButton', () => { + it('keeps configuration actionable for an idle request-scoped server', () => { + const serverStatus: MCPServerStatus = { + connectionState: 'disconnected', + authorizationState: 'not_required', + requiresOAuth: false, + requestScoped: true, + configurationState: 'needs_configuration', + }; + const baseProps = { + serverName: 'server', + serverStatus, + isInitializing: false, + canCancel: false, + onCancel: jest.fn(), + onConfigClick: jest.fn(), + }; + + expect(shouldShowActionButton({ ...baseProps, hasCustomUserVars: true })).toBe(true); + expect(shouldShowActionButton({ ...baseProps, hasCustomUserVars: false })).toBe(false); + }); +}); diff --git a/client/src/components/MCP/mcpServerUtils.ts b/client/src/components/MCP/mcpServerUtils.ts index 7cd5f2d0099..47320a3ac7f 100644 --- a/client/src/components/MCP/mcpServerUtils.ts +++ b/client/src/components/MCP/mcpServerUtils.ts @@ -64,7 +64,7 @@ export function getSelectedServerIcons( /** * Unified status color system following UX best practices: * - Green: Connected/Active (success) - * - Blue: Connecting/In-progress (processing) + * - Blue: Connecting/In-progress or request-scoped on-demand * - Amber: Needs user action (OAuth required, config missing) * - Gray: Disconnected/Inactive (neutral - server is simply off) * - Red: Error (failed, needs retry) @@ -87,13 +87,17 @@ export function getStatusColor( return 'bg-status-neutral'; } - const { connectionState, requiresOAuth } = status; + const { connectionState, requiresOAuth, requestScoped } = status; // Connecting: blue (in progress) if (connectionState === 'connecting') { return 'bg-status-info'; } + if (requestScoped) { + return 'bg-status-info'; + } + // Connected: green (success) if (connectionState === 'connected') { return 'bg-status-success'; @@ -131,7 +135,15 @@ export function getStatusTextKey( return 'com_nav_mcp_status_unknown'; } - const { connectionState, requiresOAuth } = status; + const { connectionState, requiresOAuth, requestScoped } = status; + + if (connectionState === 'connecting') { + return 'com_nav_mcp_status_connecting'; + } + + if (requestScoped) { + return 'com_nav_mcp_status_on_demand'; + } // Special case: disconnected but needs OAuth shows different text if (connectionState === 'disconnected' && requiresOAuth) { @@ -157,7 +169,9 @@ export function serverNeedsAction( _hasCustomUserVars?: boolean, ): boolean { if (!serverStatus) return false; - const { connectionState, requiresOAuth } = serverStatus; + const { connectionState, requiresOAuth, requestScoped } = serverStatus; + + if (requestScoped && connectionState !== 'connecting') return false; // Needs OAuth authentication if (connectionState === 'disconnected' && requiresOAuth) return true; @@ -168,6 +182,31 @@ export function serverNeedsAction( return false; } +/** + * Request-scoped servers are usable without an idle transport connection once + * their authorization requirement is satisfied. Agent tooling uses this + * readiness signal to attach the runtime wildcard instead of waiting for a + * tool catalog that can only be discovered during a chat request. + */ +export function isMCPServerReadyForAgent( + status: MCPServerStatus | undefined, + requestScoped: boolean, + hasCustomUserVars = false, +): boolean { + if (requestScoped && hasCustomUserVars && status?.configurationState !== 'configured') { + return false; + } + if (status?.connectionState === 'connected') { + return true; + } + if (!requestScoped) { + return false; + } + return ( + status?.authorizationState === 'not_required' || status?.authorizationState === 'authorized' + ); +} + /** * Determines if an action button should be shown for a server status. * Returns true only when the button would be actionable (not just informational). @@ -183,8 +222,14 @@ export function shouldShowActionButton(statusIconProps?: MCPServerStatusIconProp if (isInitializing) return false; if (!serverStatus) return false; - const { connectionState, requiresOAuth } = serverStatus; + const { connectionState, requiresOAuth, requestScoped } = serverStatus; + // Request-scoped servers can only be initialized with an active MCP request context, + // but their per-user variables must remain configurable while idle. + if ((connectionState === 'disconnected' || connectionState === 'error') && requestScoped) { + return hasCustomUserVars === true; + } + if (connectionState === 'connected' && requestScoped) return hasCustomUserVars === true; // Show for disconnected/error (can reconnect/configure) if (connectionState === 'disconnected' || connectionState === 'error') return true; // Show a cancel action for pending OAuth connections. diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx index eac90450d2b..8b03aa24436 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/__tests__/McpSection.spec.tsx @@ -8,6 +8,7 @@ const mockSetValue = jest.fn(); const mockGetValues = jest.fn((): string[] => []); const mockGetToolOptions = jest.fn((): Record | undefined => undefined); const mockMcpServersMap = jest.fn((): Map => new Map()); +const mockGetServerStatusIconProps = jest.fn((): object | null => null); const mockInitializeServer = jest.fn(); const mockIsConnectionDeferred = jest.fn((): boolean => false); const mockToggleIntentAll = jest.fn(); @@ -20,6 +21,9 @@ const mockCapabilities = { backgroundToolsEnabled: false, toolIntentsEnabled: false, }; +const mockLocalize = jest.fn((key: string, values?: Record) => + key === 'com_nav_mcp_status_connecting' ? `${values?.[0]} - Connecting` : key, +); jest.mock('react-hook-form', () => ({ useFormContext: () => ({ control: {}, setValue: mockSetValue, getValues: mockGetValues }), @@ -46,12 +50,12 @@ jest.mock('~/components/ui', () => ({ })); jest.mock('~/hooks', () => ({ - useLocalize: () => (key: string) => key, + useLocalize: () => mockLocalize, useCopyToClipboard: () => jest.fn(), useAgentCapabilities: () => mockCapabilities, useGetAgentsConfig: () => ({ agentsConfig: { capabilities: [] } }), useMCPServerManager: () => ({ - getServerStatusIconProps: () => null, + getServerStatusIconProps: mockGetServerStatusIconProps, getConfigDialogProps: () => null, initializeServer: mockInitializeServer, isConnectionDeferred: mockIsConnectionDeferred, @@ -121,6 +125,7 @@ jest.mock('@librechat/client', () => { const React = jest.requireActual('react'); return { TooltipAnchor: ({ render }: { render: React.ReactElement }) => render, + Spinner: ({ className }: { className?: string }) => React.createElement('span', { className }), Button: ({ children, variant: _variant, @@ -181,6 +186,9 @@ describe('McpSection', () => { mockGetToolOptions.mockReturnValue(undefined); mockMcpServersMap.mockReset(); mockMcpServersMap.mockReturnValue(new Map()); + mockGetServerStatusIconProps.mockReset(); + mockGetServerStatusIconProps.mockReturnValue(null); + mockLocalize.mockClear(); mockCodeInterpreterSelected.mockReset(); mockCodeInterpreterSelected.mockReturnValue(false); mockCapabilities.codeEnabled = false; @@ -196,6 +204,21 @@ describe('McpSection', () => { expect(screen.getByTestId('tool-mcp:srv:b')).toBeInTheDocument(); }); + test('interpolates the server name when another manager reports a connecting state', () => { + mockGetServerStatusIconProps.mockReturnValue({ + serverStatus: { + connectionState: 'connecting', + requiresOAuth: true, + }, + isInitializing: false, + }); + + render(); + + expect(screen.getByText('srv - Connecting')).toBeInTheDocument(); + expect(mockLocalize).toHaveBeenCalledWith('com_nav_mcp_status_connecting', { 0: 'srv' }); + }); + test('toggling a tool writes its id plus the server token into agent.tools', () => { render(); fireEvent.click(screen.getByTestId('tool-mcp:srv:a')); @@ -282,13 +305,14 @@ describe('McpSection', () => { expect(screen.getByText('com_ui_tools_mcp_no_tools')).toBeInTheDocument(); }); - test('lets an already-connected request-scoped server attach its runtime tools', () => { + test('lets a ready request-scoped server attach its runtime tools', () => { const runtimeItem: McpItem = { ...item, server: { ...item.server, tools: [], - isConnected: true, + isConnected: false, + isReadyForAgent: true, requestScoped: true, } as never, toolCount: 0, @@ -320,6 +344,7 @@ describe('McpSection', () => { ...item.server, tools: [], isConnected: true, + isReadyForAgent: true, requestScoped: true, } as never, toolCount: 0, diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx index 95248599d6e..accc7941b30 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx @@ -10,6 +10,7 @@ import { buildServerNameAliases, stripServerNamePrefix, } from 'librechat-data-provider'; +import type { MCPServerStatus } from 'librechat-data-provider'; import type { MouseEvent } from 'react'; import type { TranslationKeys } from '~/hooks/useLocalize'; import type { McpItem } from '../../items/types'; @@ -21,6 +22,7 @@ import { useMCPToolOptions, } from '~/hooks'; import { matchesMcpServer, mcpAllToken, mcpServerToken } from '../../items/selectors'; +import { getStatusColor, getStatusTextKey } from '~/components/MCP/mcpServerUtils'; import MCPServerStatusIcon from '~/components/MCP/MCPServerStatusIcon'; import MCPConfigDialog from '~/components/MCP/MCPConfigDialog'; import McpOAuthDialog from '~/components/MCP/McpOAuthDialog'; @@ -38,29 +40,23 @@ interface StatusDisplay { } function getStatusDisplay( - connectionState: string | undefined, + serverName: string, + serverStatus: MCPServerStatus | undefined, isInitializing: boolean, isConfigured: boolean, ): StatusDisplay { - if (isInitializing || connectionState === 'connecting') { - return { - labelKey: 'com_nav_mcp_status_initializing', - dotClass: 'bg-blue-500 animate-pulse', - }; + if (!serverStatus && !isInitializing && !isConfigured) { + return { labelKey: 'com_ui_tools_mcp_status_unconfigured', dotClass: 'bg-status-neutral' }; } - if (connectionState === 'connected') { - return { labelKey: 'com_nav_mcp_status_connected', dotClass: 'bg-emerald-500' }; - } - if (connectionState === 'error') { - return { labelKey: 'com_nav_mcp_status_error', dotClass: 'bg-red-500' }; - } - if (connectionState === 'disconnected') { - return { labelKey: 'com_nav_mcp_status_disconnected', dotClass: 'bg-amber-500' }; - } - if (!isConfigured) { - return { labelKey: 'com_ui_tools_mcp_status_unconfigured', dotClass: 'bg-gray-400' }; - } - return { labelKey: 'com_nav_mcp_status_unknown', dotClass: 'bg-gray-400' }; + const connectionStatus = serverStatus ? { [serverName]: serverStatus } : undefined; + const initializing = () => isInitializing; + return { + labelKey: getStatusTextKey(serverName, connectionStatus, initializing) as TranslationKeys, + dotClass: cn( + getStatusColor(serverName, connectionStatus, initializing), + (isInitializing || serverStatus?.connectionState === 'connecting') && 'animate-pulse', + ), + }; } interface Props { @@ -80,7 +76,7 @@ export default function McpSection({ item }: Props) { } = useMCPServerManager(); const [oauthOpen, setOauthOpen] = useState(false); const [oauthUrl, setOauthUrl] = useState(null); - const [prevConnected, setPrevConnected] = useState(false); + const [prevReadyForAgent, setPrevReadyForAgent] = useState(false); const [autoSelectPending, setAutoSelectPending] = useState(false); const { mcpServersMap, mcpToolsLoading } = useAgentPanelContext(); const { agentsConfig } = useGetAgentsConfig(); @@ -335,21 +331,27 @@ export default function McpSection({ item }: Props) { const configDialogProps = getConfigDialogProps(); const connectionState = statusIconProps?.serverStatus?.connectionState; const isInitializing = statusIconProps?.isInitializing ?? false; - const statusDisplay = getStatusDisplay(connectionState, isInitializing, liveServer.isConfigured); + const statusDisplay = getStatusDisplay( + serverName, + statusIconProps?.serverStatus, + isInitializing, + liveServer.isConfigured, + ); /** A connected server's tools arrive with the (cold-cache) MCP tools fetch, and * the server is also briefly toolless while initializing — show a skeleton in * both cases instead of a misleading "no tools" message. */ const toolsLoading = !hasTools && (mcpToolsLoading || isInitializing || connectionState === 'connecting'); const isConnected = connectionState === 'connected' || liveServer.isConnected === true; + const isReadyForAgent = liveServer.isReadyForAgent ?? isConnected; const isBusy = isInitializing || connectionState === 'connecting'; - /** Close + clear the OAuth dialog once the server connects, and don't let it + /** Close + clear the OAuth dialog once the server is ready, and don't let it * reopen on its own if the connection later drops. No useEffect — adjust state * during render by comparing against the previous connection result. */ - if (prevConnected !== isConnected) { - setPrevConnected(isConnected); - if (isConnected) { + if (prevReadyForAgent !== isReadyForAgent) { + setPrevReadyForAgent(isReadyForAgent); + if (isReadyForAgent) { setOauthOpen(false); setOauthUrl(null); } @@ -370,7 +372,7 @@ export default function McpSection({ item }: Props) { const initConnectionDeferred = isConnectionDeferred(serverName); const requestScoped = liveServer.requestScoped === true; const runtimeToolsAvailable = - !hasTools && !toolsLoading && (isWildcardAttached || (requestScoped && isConnected)); + !hasTools && !toolsLoading && (isWildcardAttached || (requestScoped && isReadyForAgent)); const runtimeToolsMessage = isWildcardAttached ? 'com_ui_tools_mcp_runtime_tools' : 'com_ui_tools_mcp_runtime_tools_available'; @@ -447,19 +449,19 @@ export default function McpSection({ item }: Props) { aria-hidden="true" /> - {localize(statusDisplay.labelKey)} + {localize(statusDisplay.labelKey, { 0: serverName })}
- {isConnected && statusIconProps && } + {isReadyForAgent && statusIconProps && }
- {/* Connect collapses smoothly once connected. Its top spacing lives inside + {/* Connect collapses smoothly once ready. Its top spacing lives inside * the reveal so the parent's flex gap never leaves a hole when it's gone, * and the auto-height dialog follows the grid-rows tween in one motion. */}
@@ -468,8 +470,8 @@ export default function McpSection({ item }: Props) { variant="submit" className="mt-5 w-full gap-2" disabled={isBusy} - tabIndex={isConnected ? -1 : undefined} - aria-hidden={isConnected || undefined} + tabIndex={isReadyForAgent ? -1 : undefined} + aria-hidden={isReadyForAgent || undefined} onClick={handleConnect} > {isBusy && } @@ -614,7 +616,7 @@ export default function McpSection({ item }: Props) { {configDialogProps && } { ); }); - test('clicking a connected request-scoped zero-tool server attaches its runtime wildcard', () => { + test('clicking a ready request-scoped zero-tool server attaches its runtime wildcard', () => { mockMcpServersMap = new Map([ [ 'runtime', @@ -237,7 +237,8 @@ describe('ToolsMarketplaceDialog', () => { serverName: 'runtime', tools: [], isConfigured: true, - isConnected: true, + isConnected: false, + isReadyForAgent: true, requestScoped: true, metadata: { name: 'runtime', pluginKey: 'runtime', description: '' }, }, @@ -272,6 +273,7 @@ describe('ToolsMarketplaceDialog', () => { tools: [], isConfigured: true, isConnected: true, + isReadyForAgent: true, requestScoped: true, metadata: { name: 'runtime', pluginKey: 'runtime', description: '' }, }, diff --git a/client/src/components/SidePanel/MCPBuilder/MCPCardActions.spec.tsx b/client/src/components/SidePanel/MCPBuilder/MCPCardActions.spec.tsx index 6067c16ff8d..e21a057f66a 100644 --- a/client/src/components/SidePanel/MCPBuilder/MCPCardActions.spec.tsx +++ b/client/src/components/SidePanel/MCPBuilder/MCPCardActions.spec.tsx @@ -47,4 +47,59 @@ describe('MCPCardActions', () => { expect(revokeButton).toHaveClass('hover:text-text-secondary'); expect(revokeButton.querySelector('svg')).toHaveClass('text-text-destructive'); }); + + test.each([ + ['disconnected', 'com_nav_mcp_connect'], + ['error', 'com_nav_mcp_connect'], + ['connected', 'com_nav_mcp_reconnect'], + ] as const)( + 'does not render a manual connection action when %s and on-demand', + (state, label) => { + render( + , + ); + + expect(screen.queryByRole('button', { name: label })).not.toBeInTheDocument(); + }, + ); + + test('keeps custom-variable configuration available while an on-demand server is idle', () => { + render( + , + ); + + expect(screen.getByRole('button', { name: 'com_ui_configure' })).toBeInTheDocument(); + expect(screen.queryByRole('button', { name: 'com_nav_mcp_connect' })).not.toBeInTheDocument(); + }); }); diff --git a/client/src/components/SidePanel/MCPBuilder/MCPCardActions.tsx b/client/src/components/SidePanel/MCPBuilder/MCPCardActions.tsx index 6becb23931f..14a5770ab00 100644 --- a/client/src/components/SidePanel/MCPBuilder/MCPCardActions.tsx +++ b/client/src/components/SidePanel/MCPBuilder/MCPCardActions.tsx @@ -125,7 +125,7 @@ export default function MCPCardActions({ )} {/* Connect button - for disconnected or error states */} - {(isDisconnected || isError) && ( + {(isDisconnected || isError) && !serverStatus?.requestScoped && ( )} - {/* Configure button - for connected servers with custom vars */} - {isConnected && hasCustomUserVars && ( + {/* On-demand servers stay idle between requests, so their user variables + must remain configurable without a live transport connection. */} + {(isConnected || serverStatus?.requestScoped) && hasCustomUserVars && ( ({ + useLocalize: () => (key: string) => key, +})); + +jest.mock('@librechat/client', () => ({ + Spinner: (props: React.ComponentProps<'span'>) => , +})); + +describe('MCPStatusBadge', () => { + test.each(['disconnected', 'connected', 'error'] as const)( + 'renders the %s request-scoped state as on-demand', + (connectionState) => { + const serverStatus: MCPServerStatus = { + connectionState, + requiresOAuth: true, + requestScoped: true, + }; + + render(); + + expect(screen.getByRole('status')).toHaveTextContent('com_nav_mcp_status_on_demand'); + expect(getStatusDotColor(serverStatus)).toBe('bg-status-info'); + }, + ); + + it('preserves the active connecting state for a request-scoped OAuth flow', () => { + const serverStatus: MCPServerStatus = { + connectionState: 'connecting', + requiresOAuth: true, + requestScoped: true, + }; + + render(); + + expect(screen.getByRole('status')).toHaveTextContent('com_nav_mcp_status_connecting'); + }); +}); diff --git a/client/src/components/SidePanel/MCPBuilder/MCPStatusBadge.tsx b/client/src/components/SidePanel/MCPBuilder/MCPStatusBadge.tsx index fd3eefa769f..c3306d8b86d 100644 --- a/client/src/components/SidePanel/MCPBuilder/MCPStatusBadge.tsx +++ b/client/src/components/SidePanel/MCPBuilder/MCPStatusBadge.tsx @@ -1,5 +1,5 @@ import { Spinner } from '@librechat/client'; -import { Check, PlugZap } from 'lucide-react'; +import { Check, PlugZap, Zap } from 'lucide-react'; import type { MCPServerStatus } from 'librechat-data-provider'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; @@ -15,7 +15,7 @@ interface MCPStatusBadgeProps { * * Unified color system: * - Green: Connected/Active (success) - * - Blue: Connecting/In-progress + * - Blue: Connecting/In-progress or request-scoped on-demand * - Amber: Needs user action (OAuth required) * - Gray: Disconnected/Inactive (neutral) * - Red: Error @@ -66,6 +66,15 @@ export default function MCPStatusBadge({ ); } + if (serverStatus.requestScoped) { + return ( +
+
+ ); + } + // Disconnected state - check if needs action if (connectionState === 'disconnected') { if (requiresOAuth) { @@ -121,7 +130,7 @@ export default function MCPStatusBadge({ * * Colors: * - Green: Connected - * - Blue: Connecting/Initializing + * - Blue: Connecting/Initializing or request-scoped on-demand * - Amber: Needs action (OAuth required while disconnected) * - Gray: Disconnected (neutral) * - Red: Error @@ -144,6 +153,10 @@ export function getStatusDotColor( return 'bg-status-info'; } + if (serverStatus.requestScoped) { + return 'bg-status-info'; + } + if (connectionState === 'connected') { return 'bg-status-success'; } @@ -153,8 +166,10 @@ export function getStatusDotColor( } if (connectionState === 'disconnected') { - // Needs OAuth = amber, otherwise gray - return requiresOAuth ? 'bg-status-warning' : 'bg-status-neutral'; + if (requiresOAuth) { + return 'bg-status-warning'; + } + return 'bg-status-neutral'; } return 'bg-status-neutral'; diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index 910bba288be..46718b59bfd 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -169,9 +169,26 @@ export function useMCPServerManager({ // Poll intervals are kept local (not serializable) const pollIntervalsRef = useRef({}); - const { connectionStatus } = useMCPConnectionStatus({ + const { connectionStatus: polledConnectionStatus } = useMCPConnectionStatus({ enabled: !isLoading && availableMCPServers.length > 0, }); + const connectionStatus = useMemo(() => { + if (!polledConnectionStatus) { + return polledConnectionStatus; + } + + let changed = false; + const nextStatus: MCPConnectionStatusResponse['connectionStatus'] = {}; + for (const [serverName, status] of Object.entries(polledConnectionStatus)) { + if (status.requestScoped === true || loadedServers?.[serverName]?.requestScoped !== true) { + nextStatus[serverName] = status; + continue; + } + changed = true; + nextStatus[serverName] = { ...status, requestScoped: true }; + } + return changed ? nextStatus : polledConnectionStatus; + }, [polledConnectionStatus, loadedServers]); const updateServerInitState = useCallback( (serverName: string, updates: Partial) => { diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 5ad9bdf4d88..8b1b216a02b 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -610,6 +610,7 @@ "com_nav_mcp_status_error": "Error", "com_nav_mcp_status_initializing": "Initializing", "com_nav_mcp_status_needs_auth": "Needs Auth", + "com_nav_mcp_status_on_demand": "On-demand", "com_nav_mcp_status_unknown": "Unknown", "com_nav_mcp_vars_update_error": "Error updating MCP custom user variables", "com_nav_mcp_vars_updated": "MCP custom user variables updated successfully.", diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index 8138eb2e047..b517ae49789 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -2540,6 +2540,32 @@ describe('initializeAgent — run-scoped MCP tool definitions', () => { ); }); + it('threads the normalized MCP request body into tool discovery', async () => { + const { agent, req, res, loadTools, db } = createMocks(); + agent.tools = ['custom_tool']; + const requestBody = { + messageId: 'message-1', + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + }; + + await initializeAgent( + { + req, + res, + agent, + loadTools, + requestBody, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + }, + db, + ); + + expect(loadTools).toHaveBeenCalledWith(expect.objectContaining({ requestBody })); + }); + it('unions snapshot config names into the audit when the merged read omits them', async () => { /** The registry's merged read tolerates config-server init failures and * can silently drop config-only servers — the heal audit must restore diff --git a/packages/api/src/agents/discovery.spec.ts b/packages/api/src/agents/discovery.spec.ts index 10ee3a5190c..dce09b5425a 100644 --- a/packages/api/src/agents/discovery.spec.ts +++ b/packages/api/src/agents/discovery.spec.ts @@ -273,6 +273,40 @@ describe('discoverConnectedAgents', () => { ); }); + it('forwards normalized request metadata to every handoff initializeAgent call', async () => { + const primaryConfig = makeConfig('A', [{ from: 'A', to: 'B', edgeType: 'handoff' }]); + const getAgent = jest.fn(async () => makeAgent('B', [])); + const checkPermission = jest.fn().mockResolvedValue(true); + const requestBody = { + messageId: 'message-1', + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + }; + + await discoverConnectedAgents( + { + req: makeReq(), + res: makeRes(), + primaryConfig, + allowedProviders: new Set(), + modelsConfig: { openai: ['gpt-4o'] }, + loadTools: jest.fn(), + requestBody, + }, + { + getAgent, + checkPermission, + logViolation: jest.fn(), + db: {} as never, + }, + ); + + expect(mockInitializeAgent).toHaveBeenCalledWith( + expect.objectContaining({ requestBody }), + expect.anything(), + ); + }); + it('forwards codeEnvAvailable=false verbatim so handoff agents respect disabled capability', async () => { /* Symmetric to the "true" case: when the primary resolved `codeEnvAvailable = false`, handoffs must NOT accidentally diff --git a/packages/api/src/agents/discovery.ts b/packages/api/src/agents/discovery.ts index 2fb107940e4..f525481f237 100644 --- a/packages/api/src/agents/discovery.ts +++ b/packages/api/src/agents/discovery.ts @@ -69,6 +69,8 @@ export interface DiscoverConnectedAgentsParams { requestFiles?: InitializeAgentParams['requestFiles']; conversationId?: string | null; parentMessageId?: string | null; + /** Normalized runtime request metadata forwarded to MCP tool loading. */ + requestBody?: InitializeAgentParams['requestBody']; /** * ResourceType to check each sub-agent's access against. Defaults to * `AGENT` for the in-app chat flow. Callers whose entry-point gates on @@ -230,6 +232,7 @@ async function initializeReferencedAgent( requestFiles: params.requestFiles, conversationId: params.conversationId, parentMessageId: params.parentMessageId, + requestBody: params.requestBody, endpointOption: { ...(params.endpointOption ?? {}), endpoint: EModelEndpoint.agents, diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 4e0442b4ef6..48d02050602 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -78,6 +78,8 @@ export interface ToolExecuteOptions { loadTools: ( toolNames: string[], agentId?: string, + /** Immutable run configuration available before deferred tools connect. */ + configurable?: Record, ) => Promise<{ loadedTools: StructuredToolInterface[]; /** Additional configurable properties to merge (e.g., userMCPAuthMap) */ @@ -3848,12 +3850,13 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand await runOutsideTracing(async () => { try { const toolNames = [...new Set(toolCalls.map((tc: ToolCallRequest) => tc.name))]; + const sourceConfigurable = configurable as Record | undefined; const { loadedTools, configurable: toolConfigurable } = await loadTools( toolNames, agentId, + sourceConfigurable, ); const toolMap = new Map(loadedTools.map((t) => [t.name, t])); - const sourceConfigurable = configurable as Record | undefined; const loadedConfigurable = toolConfigurable as Record | undefined; const mergedConfigurable = mergeToolConfigurables( sourceConfigurable, diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 9c429832d25..bcc08733197 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -26,18 +26,19 @@ import type { import type { GenericTool, LCToolRegistry, ToolMap, LCTool } from '@librechat/agents'; import type { IMongoFile, FileOwnerScope } from '@librechat/data-schemas'; import type { Response as ServerResponse } from 'express'; -import type { - ResolvedManualSkill, - ResolvedAlwaysApplySkill, - TListSkillsByAccess, - TGetSkillByName, -} from './skills'; import type { ServerRequest, + RequestBody, EndpointDbMethods, EndpointTokenConfig, InitializeResultBase, } from '~/types'; +import type { + ResolvedManualSkill, + ResolvedAlwaysApplySkill, + TListSkillsByAccess, + TGetSkillByName, +} from './skills'; import type { LCAvailableTools, RequestScopedMCPConnectionStore } from '../mcp/types'; import type { TFilterFilesByAgentAccess } from './resources'; import type { MCPToolAlias } from '~/tools/classification'; @@ -417,6 +418,8 @@ export interface InitializeAgentParams { conversationId?: string | null; /** Parent message ID for determining the current thread (optional) */ parentMessageId?: string | null; + /** Normalized body used by MCP runtime placeholders during tool discovery. */ + requestBody?: RequestBody; /** Request files */ requestFiles?: IMongoFile[]; /** Function to load agent tools */ @@ -429,6 +432,7 @@ export interface InitializeAgentParams { model: string | null; tool_options: AgentToolOptions | undefined; tool_resources: AgentToolResources | undefined; + requestBody?: RequestBody; /** Trusted endpoint/profile resolved for this agent before any code-file priming. */ codeExecutionContext: CodeExecutionContext; /** Full accessible MCP server names (operator + user DB) when the heal @@ -613,6 +617,7 @@ export async function initializeAgent( conversationId, endpointOption, parentMessageId, + requestBody, allowedProviders, isInitialAgent = false, } = params; @@ -1068,6 +1073,7 @@ export async function initializeAgent( model: agent.model, tool_options: agent.tool_options, tool_resources, + requestBody, codeExecutionContext, accessibleMcpServerNames: resolvedAuditNames, }); diff --git a/packages/api/src/agents/openai/service.spec.ts b/packages/api/src/agents/openai/service.spec.ts index 75bcf091724..bc00e0e3bc1 100644 --- a/packages/api/src/agents/openai/service.spec.ts +++ b/packages/api/src/agents/openai/service.spec.ts @@ -1,3 +1,4 @@ +import { GraphEvents } from '@librechat/agents'; import { ErrorTypes } from 'librechat-data-provider'; import type { ChatCompletionDependencies } from './service'; import { createAgentChatCompletion } from './service'; @@ -15,6 +16,7 @@ type CreateRunArgs = { user?: Record; tenantId?: string; appConfig?: Record; + requestBody?: Record; }; type ProcessStreamConfig = { configurable?: Record }; @@ -110,6 +112,88 @@ describe('createAgentChatCompletion - MCP permission user propagation', () => { expect(streamConfig.configurable?.user).not.toHaveProperty('role'); }); + it('threads the parent message id into the run and execution context', async () => { + const req = createMockReq({ id: 'user-123', role: 'USER' }) as unknown as { + body: Record; + }; + req.body.parent_message_id = 'parent-123'; + + await createAgentChatCompletion(req as never, createMockRes(), deps); + + expect(deps.initializeAgent).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: expect.objectContaining({ parentMessageId: 'parent-123' }), + }), + ); + const runArgs = createRun.mock.calls[0][0] as CreateRunArgs; + expect(runArgs.requestBody).toEqual(expect.objectContaining({ parentMessageId: 'parent-123' })); + const streamConfig = processStream.mock.calls[0][1] as ProcessStreamConfig; + expect(streamConfig.configurable?.requestBody).toEqual(runArgs.requestBody); + }); + + it('forwards the normalized MCP body to deferred execution loaders', async () => { + const req = createMockReq({ id: 'user-123', role: 'USER' }) as unknown as { + body: Record; + }; + req.body.stream = true; + req.body.parent_message_id = 'parent-123'; + const loadTools = jest.fn().mockResolvedValue({ loadedTools: [] }); + deps.toolExecuteOptions = { loadTools }; + + await createAgentChatCompletion(req as never, createMockRes(), deps); + + const runArgs = createRun.mock.calls[0][0] as CreateRunArgs & { + customHandlers: Record Promise }>; + }; + const streamConfig = processStream.mock.calls[0][1] as ProcessStreamConfig; + const resolve = jest.fn(); + const reject = jest.fn(); + await runArgs.customHandlers[GraphEvents.ON_TOOL_EXECUTE].handle(GraphEvents.ON_TOOL_EXECUTE, { + toolCalls: [{ id: 'tool-call-1', name: 'deferred_mcp_tool', args: {} }], + agentId: 'agent_test', + configurable: streamConfig.configurable, + metadata: {}, + resolve, + reject, + }); + + expect(loadTools).toHaveBeenCalledWith( + ['deferred_mcp_tool'], + 'agent_test', + expect.objectContaining({ requestBody: runArgs.requestBody }), + ); + }); + + it('uses the root parent sentinel when chat completions omit a parent id', async () => { + const req = createMockReq({ id: 'user-123', role: 'USER' }); + + await createAgentChatCompletion(req, createMockRes(), deps); + + expect(deps.initializeAgent).toHaveBeenCalledWith( + expect.objectContaining({ + requestBody: expect.objectContaining({ + parentMessageId: '00000000-0000-0000-0000-000000000000', + }), + }), + ); + }); + + it('omits an unavailable parent for an existing chat-completions conversation', async () => { + const req = createMockReq({ id: 'user-123', role: 'USER' }) as unknown as { + body: Record; + }; + req.body.conversation_id = 'conversation-123'; + + await createAgentChatCompletion(req as never, createMockRes(), deps); + + const requestBody = (deps.initializeAgent as jest.Mock).mock.calls[0][0].requestBody; + expect(requestBody).toEqual({ + messageId: expect.any(String), + conversationId: 'conversation-123', + }); + expect(requestBody).not.toHaveProperty('parentMessageId'); + }); + it('forwards appConfig and tenantId to createRun', async () => { const appConfig = { endpoints: { diff --git a/packages/api/src/agents/openai/service.ts b/packages/api/src/agents/openai/service.ts index 193e8fa92fc..62b7d8df8e8 100644 --- a/packages/api/src/agents/openai/service.ts +++ b/packages/api/src/agents/openai/service.ts @@ -32,6 +32,7 @@ import type { ToolCall, } from './types'; import type { OpenAIStreamHandlerConfig, EventHandler } from './handlers'; +import type { MCPRuntimeRequestBody } from '~/mcp/request'; import type { ToolExecuteOptions } from '../handlers'; import { createOpenAIContentAggregator, @@ -41,6 +42,7 @@ import { createChunk, writeSSE, } from './handlers'; +import { createMCPRuntimeRequestBody } from '~/mcp/request'; import { createSafeUser } from '~/utils'; /** @@ -135,6 +137,7 @@ interface InitializeAgentParams { agent: Agent; conversationId?: string | null; parentMessageId?: string | null; + requestBody?: MCPRuntimeRequestBody; requestFiles?: unknown[]; loadTools?: LoadToolsFn; endpointOption?: Record; @@ -191,6 +194,7 @@ type LoadToolsFn = (params: { model: string | null; tool_options: unknown; tool_resources: unknown; + requestBody?: MCPRuntimeRequestBody; }) => Promise<{ tools: unknown[]; toolContextMap: Record; @@ -435,6 +439,17 @@ export async function createAgentChatCompletion( // Generate IDs const requestId = `chatcmpl-${nanoid()}`; const conversationId = request.conversation_id ?? nanoid(); + let mcpParentMessageId: string | null | undefined; + if (typeof request.parent_message_id === 'string' && request.parent_message_id.trim() !== '') { + mcpParentMessageId = request.parent_message_id; + } else if (request.conversation_id == null) { + mcpParentMessageId = null; + } + const mcpRequestBody = createMCPRuntimeRequestBody({ + messageId: requestId, + conversationId, + parentMessageId: mcpParentMessageId, + }); const created = Math.floor(Date.now() / 1000); // Build response context @@ -502,6 +517,7 @@ export async function createAgentChatCompletion( agent, conversationId, parentMessageId: request.parent_message_id, + requestBody: mcpRequestBody, loadTools: deps.loadAgentTools, endpointOption: { endpoint: agent.provider, @@ -570,17 +586,13 @@ export async function createAgentChatCompletion( * correctly leaves MCP gated. */ const safeUser: Record = { ...createSafeUser(reqUser), id: userId }; - const run = await deps.createRun({ agents: [initializedAgent], messages, runId: requestId, signal: abortController.signal, customHandlers: eventHandlers, - requestBody: { - messageId: requestId, - conversationId, - }, + requestBody: mcpRequestBody, user: safeUser, tenantId: typeof reqUser?.tenantId === 'string' ? reqUser.tenantId : undefined, appConfig: deps.appConfig @@ -600,6 +612,7 @@ export async function createAgentChatCompletion( thread_id: conversationId, user_id: userId, user: safeUser, + requestBody: mcpRequestBody, /** Same per-agent channel the in-repo controllers thread via * `loadTools`: without it, the executor's PTC path cannot * strip host-injected `intent` params from the schemas the diff --git a/packages/api/src/mcp/__tests__/request.test.ts b/packages/api/src/mcp/__tests__/request.test.ts index e58c972134b..47cd41c18f6 100644 --- a/packages/api/src/mcp/__tests__/request.test.ts +++ b/packages/api/src/mcp/__tests__/request.test.ts @@ -1,6 +1,11 @@ import { EventEmitter } from 'events'; -import { getMCPRequestContext, cleanupMCPRequestContextForReq } from '~/mcp/request'; +import { + createMCPRuntimeRequestBody, + getMCPRequestContext, + cleanupMCPRequestContextForReq, +} from '~/mcp/request'; +import { getMissingRuntimeBodyPlaceholderFields } from '~/mcp/utils'; jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -105,3 +110,47 @@ describe('MCP request context', () => { expect(connection.disconnect).not.toHaveBeenCalled(); }); }); + +describe('MCP runtime request body', () => { + it('preserves a supplied parent message id', () => { + expect( + createMCPRuntimeRequestBody({ + messageId: 'response-1', + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + }), + ).toEqual({ + messageId: 'response-1', + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + }); + }); + + it('uses the root-turn parent sentinel for an explicit root parent', () => { + expect( + createMCPRuntimeRequestBody({ + messageId: 'response-1', + conversationId: 'conversation-1', + parentMessageId: null, + }), + ).toEqual(expect.objectContaining({ parentMessageId: '00000000-0000-0000-0000-000000000000' })); + }); + + it('leaves the parent absent when the protocol cannot supply that identity', () => { + const requestBody = createMCPRuntimeRequestBody({ + messageId: 'response-1', + conversationId: 'conversation-1', + }); + + expect(requestBody).toEqual({ messageId: 'response-1', conversationId: 'conversation-1' }); + expect( + getMissingRuntimeBodyPlaceholderFields( + { + source: 'yaml', + headers: { 'X-Parent': '{{LIBRECHAT_BODY_PARENTMESSAGEID}}' }, + }, + requestBody, + ), + ).toEqual(['parentMessageId']); + }); +}); diff --git a/packages/api/src/mcp/__tests__/scope.integration.test.ts b/packages/api/src/mcp/__tests__/scope.integration.test.ts index 9f49d161de2..59fe600f6f6 100644 --- a/packages/api/src/mcp/__tests__/scope.integration.test.ts +++ b/packages/api/src/mcp/__tests__/scope.integration.test.ts @@ -51,6 +51,7 @@ interface RequestScopedTestServer { sessionsCreated: () => number; toolCallCount: () => number; observedRunIds: () => string[]; + observedParentMessageIds: () => string[]; } function trackSockets(httpServer: http.Server): () => Promise { @@ -72,6 +73,7 @@ function trackSockets(httpServer: http.Server): () => Promise { async function createRequestScopedTestServer(): Promise { const sessions = new Map(); const runIds: string[] = []; + const parentMessageIds: string[] = []; let created = 0; let deletes = 0; let toolCalls = 0; @@ -87,6 +89,10 @@ async function createRequestScopedTestServer(): Promise if (typeof runId === 'string') { runIds.push(runId); } + const parentMessageId = req.headers['x-parent-message']; + if (typeof parentMessageId === 'string') { + parentMessageIds.push(parentMessageId); + } } else if (req.method === 'DELETE') { deletes += 1; } @@ -127,6 +133,7 @@ async function createRequestScopedTestServer(): Promise sessionsCreated: () => created, toolCallCount: () => toolCalls, observedRunIds: () => [...runIds], + observedParentMessageIds: () => [...parentMessageIds], close: async () => { const closing = [...sessions.values()].map((transport) => transport.close().catch(() => undefined), @@ -157,7 +164,10 @@ function createServerConfig(url: string): ParsedServerConfig { source: 'yaml', requiresOAuth: false, initTimeout: 500, - headers: { 'X-Run-Id': '{{LIBRECHAT_BODY_MESSAGEID}}' }, + headers: { + 'X-Run-Id': '{{LIBRECHAT_BODY_MESSAGEID}}', + 'X-Parent-Message': '{{LIBRECHAT_BODY_PARENTMESSAGEID}}', + }, }; } @@ -246,6 +256,7 @@ describe('request-scoped MCP lifecycle integration', () => { expect(server.liveSessionCount()).toBe(1); expect(server.toolCallCount()).toBe(burstSize); expect(new Set(server.observedRunIds())).toEqual(new Set(['run-1'])); + expect(new Set(server.observedParentMessageIds())).toEqual(new Set(['parent-1'])); expect(manager.getConnectionStats().activityEntries).toBe(0); await cleanupMCPRequestContext(firstRun); @@ -266,6 +277,7 @@ describe('request-scoped MCP lifecycle integration', () => { expect(server.sessionsCreated()).toBe(2); expect(server.liveSessionCount()).toBe(1); expect(new Set(server.observedRunIds())).toEqual(new Set(['run-1', 'run-2'])); + expect(new Set(server.observedParentMessageIds())).toEqual(new Set(['parent-1'])); }); it('clears a failed run so the same server can recover in a fresh run', async () => { diff --git a/packages/api/src/mcp/__tests__/utils.test.ts b/packages/api/src/mcp/__tests__/utils.test.ts index 0f85636cc46..53cb788f29e 100644 --- a/packages/api/src/mcp/__tests__/utils.test.ts +++ b/packages/api/src/mcp/__tests__/utils.test.ts @@ -13,7 +13,7 @@ import { getMissingCustomUserVars, hasCustomUserVars, hasRuntimeUrlPlaceholders, - hasRuntimeBodyPlaceholders, + getMCPRequestScope, hasRuntimeContextPlaceholders, getRuntimeBodyPlaceholderFields, getMissingRuntimeBodyPlaceholderFields, @@ -814,32 +814,32 @@ describe('hasRuntimeUrlPlaceholders', () => { }); }); -describe('hasRuntimeBodyPlaceholders', () => { +describe('getMCPRequestScope', () => { it('detects trusted runtime BODY placeholders across connection fields', () => { expect( - hasRuntimeBodyPlaceholders({ + getMCPRequestScope({ source: 'yaml', url: 'https://example.com/conversations/{{LIBRECHAT_BODY_CONVERSATIONID}}/mcp', - }), + }).requestScoped, ).toBe(true); expect( - hasRuntimeBodyPlaceholders({ + getMCPRequestScope({ source: 'config', headers: { 'X-Message': '{{LIBRECHAT_BODY_MESSAGEID}}', }, - }), + }).requestScoped, ).toBe(true); }); it('ignores BODY placeholders in user-sourced configs', () => { expect( - hasRuntimeBodyPlaceholders({ + getMCPRequestScope({ source: 'user', dbId: 'server-123', url: 'https://example.com/{{LIBRECHAT_BODY_MESSAGEID}}/mcp', - }), + }).requestScoped, ).toBe(false); }); @@ -851,7 +851,7 @@ describe('hasRuntimeBodyPlaceholders', () => { expect(hasRuntimeContextPlaceholders(config)).toBe(false); expect(hasRuntimeUrlPlaceholders(config)).toBe(false); - expect(hasRuntimeBodyPlaceholders(config)).toBe(false); + expect(getMCPRequestScope(config).requestScoped).toBe(false); expect(getRuntimeBodyPlaceholderFields(config)).toEqual([]); expect(getMissingRuntimeBodyPlaceholderFields(config)).toEqual([]); expect(requiresEphemeralUserConnection(config)).toBe(false); @@ -869,7 +869,7 @@ describe('hasRuntimeBodyPlaceholders', () => { expect(hasRuntimeContextPlaceholders(config)).toBe(false); expect(hasRuntimeUrlPlaceholders(config)).toBe(false); - expect(hasRuntimeBodyPlaceholders(config)).toBe(false); + expect(getMCPRequestScope(config).requestScoped).toBe(false); expect(getRuntimeBodyPlaceholderFields(config)).toEqual([]); expect(getMissingRuntimeBodyPlaceholderFields(config)).toEqual([]); expect(requiresEphemeralUserConnection(config)).toBe(false); diff --git a/packages/api/src/mcp/request.ts b/packages/api/src/mcp/request.ts index 68b76027f01..4c2277181fc 100644 --- a/packages/api/src/mcp/request.ts +++ b/packages/api/src/mcp/request.ts @@ -1,6 +1,33 @@ import { logger } from '@librechat/data-schemas'; - -import type { RequestScopedMCPConnectionStore } from './types'; +import { Constants } from 'librechat-data-provider'; + +import type { MCPRuntimeRequestBody, RequestScopedMCPConnectionStore } from './types'; + +export type { MCPRuntimeRequestBody } from './types'; + +/** + * Builds the complete request context that runtime MCP placeholders may resolve. + * An explicit null parent means a known root turn and becomes the root sentinel. + * An omitted parent stays omitted so protocols without parent-message identity + * fail closed for configurations that require that BODY placeholder. + */ +export function createMCPRuntimeRequestBody({ + messageId, + conversationId, + parentMessageId, +}: { + messageId: string; + conversationId: string; + parentMessageId?: string | null; +}): MCPRuntimeRequestBody { + return { + messageId, + conversationId, + ...(parentMessageId !== undefined && { + parentMessageId: parentMessageId ?? Constants.NO_PARENT, + }), + }; +} export interface MCPRequestContext extends RequestScopedMCPConnectionStore { cleanupStarted: boolean; diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index ae59ed4af4f..e823f421b10 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -25,6 +25,9 @@ import type { FlowStateManager } from '~/flow/manager'; import type { RequestBody } from '~/types/http'; import type * as o from '~/mcp/oauth/types'; +export type MCPRuntimeRequestBody = Required> & + Pick; + export type StdioOptions = z.infer; export type WebSocketOptions = z.infer; export type SSEOptions = z.infer; diff --git a/packages/api/src/mcp/utils.ts b/packages/api/src/mcp/utils.ts index 0abbe460347..ea0d4c04462 100644 --- a/packages/api/src/mcp/utils.ts +++ b/packages/api/src/mcp/utils.ts @@ -202,6 +202,11 @@ type PlaceholderValue = | readonly PlaceholderValue[] | { readonly [key: string]: PlaceholderValue }; +export interface MCPRequestScope { + requestScoped: boolean; + requiredBodyFields: Array; +} + type UserScopedConnectionConfig = Pick< ParsedServerConfig, 'requiresOAuth' | 'source' | 'dbId' | 'startup' @@ -281,7 +286,10 @@ function hasPlaceholder(value: PlaceholderValue, pattern: RegExp): boolean { return Object.values(value).some((item) => hasPlaceholder(item, pattern)); } -function addRuntimeBodyPlaceholderFields(value: PlaceholderValue, fields: Set): void { +function addRuntimeBodyPlaceholderFields( + value: PlaceholderValue, + fields: Set, +): void { if (typeof value === 'string') { for (const match of value.matchAll(RUNTIME_BODY_PLACEHOLDER_CAPTURE_PATTERN)) { const placeholderKey = match[1]; @@ -335,34 +343,32 @@ export function hasRuntimeUrlPlaceholders(config: UserScopedConnectionConfig): b return hasRuntimeContextPlaceholder(config.url); } -export function hasRuntimeBodyPlaceholders(config: UserScopedConnectionConfig): boolean { - if (!canResolveRuntimePlaceholders(config)) { - return false; - } - - return placeholderBearingFields(config).some((value) => - hasPlaceholder(value, RUNTIME_BODY_PLACEHOLDER_PATTERN), - ); -} - -export function getRuntimeBodyPlaceholderFields(config: UserScopedConnectionConfig): string[] { +export function getMCPRequestScope(config: UserScopedConnectionConfig): MCPRequestScope { if (!canResolveRuntimePlaceholders(config)) { - return []; + return { requestScoped: false, requiredBodyFields: [] }; } - const fields = new Set(); + const requiredBodyFields = new Set(); for (const value of placeholderBearingFields(config)) { - addRuntimeBodyPlaceholderFields(value, fields); + addRuntimeBodyPlaceholderFields(value, requiredBodyFields); } - return Array.from(fields); + + const fields = Array.from(requiredBodyFields); + return { requestScoped: fields.length > 0, requiredBodyFields: fields }; +} + +export function getRuntimeBodyPlaceholderFields( + config: UserScopedConnectionConfig, +): Array { + return getMCPRequestScope(config).requiredBodyFields; } export function getMissingRuntimeBodyPlaceholderFields( config: UserScopedConnectionConfig, requestBody?: RequestBody, ): string[] { - return getRuntimeBodyPlaceholderFields(config).filter((field) => { - const value = requestBody?.[field as keyof RequestBody]; + return getMCPRequestScope(config).requiredBodyFields.filter((field) => { + const value = requestBody?.[field]; return value == null || (typeof value === 'string' && value.trim() === ''); }); } @@ -380,13 +386,7 @@ export function getMissingRuntimeBodyPlaceholderFields( * connection without forcing a reconnect for every invocation. */ export function requiresEphemeralUserConnection(config: UserScopedConnectionConfig): boolean { - if (!canResolveRuntimePlaceholders(config)) { - return false; - } - - return placeholderBearingFields(config).some((value) => - hasPlaceholder(value, RUNTIME_BODY_PLACEHOLDER_PATTERN), - ); + return getMCPRequestScope(config).requestScoped; } /** diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 25febf31ab5..710a9fe0433 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -2529,6 +2529,7 @@ class GenerationJobManagerClass { userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, isRegenerate: jobData.isRegenerate, + mcpRequestBody: jobData.mcpRequestBody, sender: jobData.sender, endpoint: jobData.endpoint, iconURL: jobData.iconURL, diff --git a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts index ff3be25270e..5610d3e59cc 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.spec.ts @@ -327,6 +327,11 @@ describe('RedisJobStore', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + mcpRequestBody: { + messageId: 'response-1', + conversationId: 'overridden-conversation', + parentMessageId: 'response-1', + }, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', @@ -389,6 +394,11 @@ describe('RedisJobStore', () => { parentMessageId: 'parent-1', }, responseMessageId: 'response-1', + mcpRequestBody: { + messageId: 'response-1', + conversationId: 'overridden-conversation', + parentMessageId: 'response-1', + }, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', @@ -423,6 +433,11 @@ describe('RedisJobStore', () => { expect(storedFields).toMatchObject({ conversationId: 'conversation-1', responseMessageId: 'response-1', + mcpRequestBody: JSON.stringify({ + messageId: 'response-1', + conversationId: 'overridden-conversation', + parentMessageId: 'response-1', + }), agent_id: 'agent-1', isTemporary: '0', scheduleId: 'schedule-1', diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index e114c8d3c5a..227635cd496 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -100,6 +100,11 @@ describe('GenerationJobManager startup telemetry', () => { }, responseMessageId: 'response-1', isRegenerate: true, + mcpRequestBody: { + messageId: 'response-1', + conversationId: 'overridden-conversation', + parentMessageId: 'response-1', + }, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', @@ -131,6 +136,11 @@ describe('GenerationJobManager startup telemetry', () => { }, responseMessageId: 'response-1', isRegenerate: true, + mcpRequestBody: { + messageId: 'response-1', + conversationId: 'overridden-conversation', + parentMessageId: 'response-1', + }, sender: 'Agent', endpoint: 'agents', iconURL: 'https://example.com/icon.png', diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index b73a4d7af98..7aa87ebfa52 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -4548,6 +4548,7 @@ export class RedisJobStore implements IJobStoreV2 { userMessage: data.userMessage ? JSON.parse(data.userMessage) : undefined, responseMessageId: data.responseMessageId || undefined, isRegenerate: data.isRegenerate != null ? data.isRegenerate === '1' : undefined, + mcpRequestBody: data.mcpRequestBody ? JSON.parse(data.mcpRequestBody) : undefined, createdEventEmitted: data.createdEventEmitted === '1', sender: data.sender || undefined, syncSent: data.syncSent === '1', diff --git a/packages/api/src/stream/interfaces/IJobStore.ts b/packages/api/src/stream/interfaces/IJobStore.ts index e086b517600..a8438b6b606 100644 --- a/packages/api/src/stream/interfaces/IJobStore.ts +++ b/packages/api/src/stream/interfaces/IJobStore.ts @@ -3,6 +3,7 @@ import type { StandardGraph } from '@librechat/agents'; import type { ActivityPhaseSnapshot } from '~/agents/activityPhases/runtime'; import type { ResolvedAskUserQuestion } from '~/agents/hitl/resume'; import type { RecoveredSteerPayload } from '../SteerRecovery'; +import type { MCPRuntimeRequestBody } from '~/mcp/types'; /** * A pause owner has this long to durably persist the interrupted turn before @@ -77,6 +78,8 @@ export interface SerializableJobData { /** Whether this generation replaces an existing assistant branch. */ isRegenerate?: boolean; + /** Exact normalized MCP placeholder identity for this turn. */ + mcpRequestBody?: MCPRuntimeRequestBody; /** * Whether this run has activity labels enabled (per-endpoint @@ -298,6 +301,7 @@ export type JobMetadataPatch = Partial< SerializableJobData, | 'responseMessageId' | 'isRegenerate' + | 'mcpRequestBody' | 'sender' | 'conversationId' | 'userMessage' diff --git a/packages/api/src/stream/metadata.ts b/packages/api/src/stream/metadata.ts index da8c108b34a..d7aade47832 100644 --- a/packages/api/src/stream/metadata.ts +++ b/packages/api/src/stream/metadata.ts @@ -9,6 +9,9 @@ export function sanitizeJobMetadata(metadata: Partial): J if (metadata.isRegenerate !== undefined) { patch.isRegenerate = metadata.isRegenerate; } + if (metadata.mcpRequestBody) { + patch.mcpRequestBody = metadata.mcpRequestBody; + } if (metadata.sender) { patch.sender = metadata.sender; } diff --git a/packages/api/src/types/stream.ts b/packages/api/src/types/stream.ts index fca5d8e9438..16e8e8124c2 100644 --- a/packages/api/src/types/stream.ts +++ b/packages/api/src/types/stream.ts @@ -2,6 +2,7 @@ import type { Agents } from 'librechat-data-provider'; import type { EventEmitter } from 'events'; import type { ActivityPhaseSnapshot } from '~/agents/activityPhases/runtime'; import type { ResolvedAskUserQuestion } from '../agents/hitl/resume'; +import type { MCPRuntimeRequestBody } from '../mcp/types'; import type { ServerSentEvent } from './events'; export interface GenerationJobMetadata { @@ -19,6 +20,9 @@ export interface GenerationJobMetadata { responseMessageId?: string; /** Whether this generation replaces an existing assistant branch. */ isRegenerate?: boolean; + /** Exact normalized MCP placeholder identity for this turn. Persisted so HITL + * resume does not reconstruct a different parent or overridden conversation. */ + mcpRequestBody?: MCPRuntimeRequestBody; /** Sender label for the response (e.g., "GPT-4.1", "Claude") */ sender?: string; /** Endpoint identifier for abort handling */ diff --git a/packages/data-provider/src/types/queries.ts b/packages/data-provider/src/types/queries.ts index 2a481d3df55..1078135b422 100644 --- a/packages/data-provider/src/types/queries.ts +++ b/packages/data-provider/src/types/queries.ts @@ -212,6 +212,10 @@ export type ListRolesResponse = { export interface MCPServerStatus { requiresOAuth: boolean; + /** The server connects only inside a chat request because its config reads BODY placeholders. */ + requestScoped?: boolean; + /** Whether all declared per-user variables are present for an on-demand connection. */ + configurationState?: 'configured' | 'needs_configuration'; connectionState: 'disconnected' | 'connecting' | 'connected' | 'error'; authorizationState?: | 'not_required' @@ -232,6 +236,8 @@ export interface MCPServerConnectionStatusResponse { success: boolean; serverName: string; requiresOAuth: boolean; + requestScoped?: boolean; + configurationState?: MCPServerStatus['configurationState']; connectionStatus: 'disconnected' | 'connecting' | 'connected' | 'error'; authorizationState?: MCPServerStatus['authorizationState']; } From 21b7f78d5645f0882a9d0cd61f9e8a4fcd722f2b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 21 Aug 2026 19:50:40 -0400 Subject: [PATCH 5/8] =?UTF-8?q?=F0=9F=99=8B=20feat:=20Collapse=20Settled?= =?UTF-8?q?=20Question=20Records=20by=20Default=20(#15107)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🙋 feat: Collapse Settled Question Records by Default The durable `ask_user_question` record rendered as a permanently open card. Answers are frequently long, multi-paragraph text, so a settled Q&A buried the reply that followed it. It now reads as one collapsed tool-call line — the same `ProgressText` primitive `ToolCall`/`SkillCall` use — naming the question (or the batch count, reusing the keys `ToolCallGroup` already had) and opening on demand under the existing `autoExpandTools` preference. Only the settled record collapses; the live pause and the interim progress card are untouched. The expanded panel was also hard to read. Authored text rendered without `pre-wrap`, so a numbered or paragraphed answer collapsed into one wall; the answer ran on from its inline label; and batch items sat flush against their divider. Line breaks are now content, the answer sits under its own label behind a rule, and dividers have air on both sides. `ProgressText`'s subtitle now truncates and absorbs the flex shrink, so arbitrary authored text ellipsizes instead of pushing the line past the message column — this also fixes long MCP server names on tool cards. * 🩹 fix: Address Codex Round 1 on the Collapsed Question Record - Settle the summary tense. A live, unanswered pause returns before the header, so every state reaching it is settled — an abandoned pause read "Asking" forever, and the collapse hid the "no answer" line that used to qualify it. Past tense unconditionally, matching `ToolCallGroup`. - Move the rejection announcement out of the disclosure. `useExpandCollapse` marks the closed panel `inert`, so the failure explanation's `role="status"` could never reach the accessibility tree; it is now an sr-only status outside the panel, carrying both the label and the explanation. - Count records, not repeated text, in the Bombadil observation. With Auto-expand tool details on, one settled record shows the question in both its summary line and its panel, so the old selector double-counted it and broke the `<= 1` singularity invariant. --- .../Messages/Content/AskUserQuestionCall.tsx | 288 +++++++++++------- .../components/Chat/Messages/Content/Part.tsx | 1 + .../Chat/Messages/Content/ProgressText.tsx | 10 +- .../__tests__/AskUserQuestionCall.test.tsx | 113 ++++++- e2e/bombadil/hitl-lifecycle.specification.ts | 28 +- 5 files changed, 311 insertions(+), 129 deletions(-) diff --git a/client/src/components/Chat/Messages/Content/AskUserQuestionCall.tsx b/client/src/components/Chat/Messages/Content/AskUserQuestionCall.tsx index a3d622ab9ca..57675883eed 100644 --- a/client/src/components/Chat/Messages/Content/AskUserQuestionCall.tsx +++ b/client/src/components/Chat/Messages/Content/AskUserQuestionCall.tsx @@ -1,3 +1,5 @@ +import { useState, useEffect, useCallback } from 'react'; +import { useRecoilValue } from 'recoil'; import { MessageCircleQuestion, TriangleAlert } from 'lucide-react'; import type { Agents } from 'librechat-data-provider'; import { @@ -6,9 +8,11 @@ import { parseAskUserQuestionsArgs, } from '~/utils/approval'; import AskUserQuestionProgress from './AskUserQuestionProgress'; +import { useLocalize, useExpandCollapse } from '~/hooks'; +import ProgressText from './ProgressText'; import EmptyText from './Parts/EmptyText'; -import { useLocalize } from '~/hooks'; import Container from './Container'; +import store from '~/store'; /** * Static rendering of a COMPLETED (or abandoned) `ask_user_question` tool call — @@ -16,6 +20,12 @@ import Container from './Container'; * is wrong here: it labels a no-output call "cancelled" and shows raw JSON args. * The interactive card ({@link AskUserQuestion}) renders only while the pause is * live; this component owns the part everywhere else (history, reload, exports). + * + * Settled, it is history — so it reads as one collapsed tool-call line (status + * label plus the question itself) and opens on demand, under the same + * `autoExpandTools` preference every other tool card follows. Answers are + * frequently long, multi-paragraph text; left expanded they buried the reply + * that followed them. */ export default function AskUserQuestionCall({ args, @@ -24,6 +34,7 @@ export default function AskUserQuestionCall({ isSubmitting = false, failed = false, showCursor = false, + onExpand, }: { args: string | Record | undefined; output: string; @@ -31,8 +42,29 @@ export default function AskUserQuestionCall({ isSubmitting?: boolean; failed?: boolean; showCursor?: boolean; + onExpand?: () => void; }) { const localize = useLocalize(); + const autoExpand = useRecoilValue(store.autoExpandTools); + const [expanded, setExpanded] = useState(autoExpand); + const { style: expandStyle, ref: expandRef } = useExpandCollapse(expanded); + + useEffect(() => { + if (autoExpand) { + setExpanded(true); + } + }, [autoExpand]); + + const toggleExpanded = useCallback(() => { + setExpanded((prev) => { + const next = !prev; + if (next) { + onExpand?.(); + } + return next; + }); + }, [onExpand]); + const question = parseAskUserQuestionArgs(args); const batch = parseAskUserQuestionsArgs(args); /** @@ -88,137 +120,169 @@ export default function AskUserQuestionCall({ ) : null; - if (batch != null) { - let statusLabel = localize('com_ui_asking'); + const count = batch?.questions.length ?? 1; + /** + * Past tense unconditionally: a live, unanswered pause returns above, so + * every state that reaches this header is settled — answered, abandoned + * (the run stopped before an answer), or rejected. An abandoned pause was + * still ASKED; it explains itself with "no answer" inside the panel, and a + * present-tense summary would strand it as permanently in-flight now that + * the panel starts closed. Matches `ToolCallGroup`, which settles its own + * question header on `!isSubmitting`. + */ + const statusLabel = (() => { if (failed) { - statusLabel = localize('com_ui_question_failed'); - } else if (answered) { - statusLabel = localize('com_ui_asked'); + return localize('com_ui_question_failed'); } - return ( - <> -
-
- {failed ? ( -
    diff --git a/client/src/components/Chat/Presentation.test.tsx b/client/src/components/Chat/Presentation.test.tsx index 756fe4a8f94..6f38dc09da0 100644 --- a/client/src/components/Chat/Presentation.test.tsx +++ b/client/src/components/Chat/Presentation.test.tsx @@ -93,6 +93,7 @@ const OpenSubagentPanel = () => { const open = () => { setConversation({ conversationId: 'parent-conversation' } as TConversation); setSelection({ + host: 'conversation', parentConversationId: 'parent-conversation', parentMessageId: 'parent-message', toolCallId: 'tool-call', diff --git a/client/src/components/Chat/Presentation.tsx b/client/src/components/Chat/Presentation.tsx index 86003aa9842..cdb4fce37a9 100644 --- a/client/src/components/Chat/Presentation.tsx +++ b/client/src/components/Chat/Presentation.tsx @@ -96,7 +96,11 @@ export default function Presentation({ children }: { children: React.ReactNode } }, [artifactsElement, resetSelectedSubagent, selectedSubagent]); const subagentElement = useMemo(() => { - if (selectedSubagent == null || selectedSubagent.parentConversationId !== conversationId) { + if ( + selectedSubagent == null || + selectedSubagent.host !== 'conversation' || + selectedSubagent.parentConversationId !== conversationId + ) { return null; } return ( diff --git a/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.test.tsx b/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.test.tsx new file mode 100644 index 00000000000..d7210cadd71 --- /dev/null +++ b/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.test.tsx @@ -0,0 +1,133 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { ContentTypes } from 'librechat-data-provider'; +import { fireEvent, render, screen, waitFor } from '@testing-library/react'; +import type { TMessageContentParts } from 'librechat-data-provider'; +import SubagentCall from '~/components/Chat/Messages/Content/Parts/SubagentCall'; +import SharedSubagentActivityDialog from './SharedSubagentActivityDialog'; +import { MessageContext } from '~/Providers/MessageContext'; +import { ShareContext } from '~/Providers/ShareContext'; + +const mockUseSubagentThreadQuery = jest.fn(); + +jest.mock('~/data-provider', () => ({ + useSubagentThreadQuery: (...args: unknown[]) => mockUseSubagentThreadQuery(...args), +})); + +jest.mock('~/hooks', () => ({ + useLocalize: + () => + (key: string, values?: Record): string => { + if (key === 'com_ui_subagent_dialog_title') return `Agent ${values?.[0] ?? ''}`; + if (key === 'com_ui_subagent_complete') return 'Ran agent'; + if (key === 'com_ui_subagent_activity') return 'Agent activity'; + return key; + }, +})); + +jest.mock('~/Providers', () => ({ useAgentsMapContext: () => ({}) })); +jest.mock('~/components/Share/MessageIcon', () => ({ __esModule: true, default: () => null })); +jest.mock('~/hooks/MCP', () => ({ useMCPServerNames: () => [] })); + +jest.mock('./SubagentActivity', () => ({ + __esModule: true, + default: ({ + activity, + }: { + activity: { title: string; items: Array<{ type: string; text?: string }> }; + }) => ( +
    + {activity.title} + {activity.items.map((item, index) => ( + {item.text ?? item.type} + ))} +
    + ), +})); + +const persistedContent = (text: string): TMessageContentParts[] => [ + { type: ContentTypes.TEXT, text } as TMessageContentParts, +]; + +const detachedOutput = JSON.stringify({ + background_task_id: 'task-1', + subagent_thread_id: 'thread-1', + tool: 'subagent', + subagent_type: 'researcher', + status: 'running', + message: + 'Started subagent "researcher" background task. Poll the host background-task tool with background_task_id "task-1".', +}); + +function renderSharedCall(input: { + output?: string; + persistedContent?: TMessageContentParts[]; + detached?: boolean; +}) { + return render( + + + + + + + + , + ); +} + +describe('SharedSubagentActivityDialog', () => { + beforeEach(() => mockUseSubagentThreadQuery.mockClear()); + + it('opens readable foreground activity from the shared message payload and restores focus', async () => { + renderSharedCall({ + output: 'Legacy fallback.', + persistedContent: persistedContent('Shared review complete.'), + }); + const trigger = screen.getByRole('button', { name: 'Ran agent' }); + + fireEvent.click(trigger); + + expect(screen.getByRole('dialog')).toBeInTheDocument(); + expect(screen.getByText('Shared review complete.')).toBeInTheDocument(); + fireEvent.click(screen.getByRole('button', { name: 'Close' })); + await waitFor(() => expect(trigger).toHaveFocus()); + }); + + it('renders detached persisted activity without performing the private durable query', () => { + renderSharedCall({ + output: detachedOutput, + persistedContent: persistedContent('Detached work survived refresh.'), + detached: true, + }); + + fireEvent.click(screen.getByRole('button', { name: 'Agent activity' })); + + expect(screen.getByText('Detached work survived refresh.')).toBeInTheDocument(); + expect(mockUseSubagentThreadQuery).not.toHaveBeenCalled(); + }); + + it('makes a detached shared card without persisted activity explicitly noninteractive', () => { + renderSharedCall({ output: detachedOutput, detached: true }); + + expect(screen.getByRole('button', { name: 'Agent activity' })).toBeDisabled(); + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(); + expect(mockUseSubagentThreadQuery).not.toHaveBeenCalled(); + }); +}); diff --git a/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.tsx b/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.tsx new file mode 100644 index 00000000000..9bcde4ab49f --- /dev/null +++ b/client/src/components/Chat/Subagents/SharedSubagentActivityDialog.tsx @@ -0,0 +1,74 @@ +import { useCallback, useEffect, useMemo, useRef } from 'react'; +import { useRecoilValue, useResetRecoilState } from 'recoil'; +import { OGDialog, OGDialogContent, OGDialogHeader, OGDialogTitle } from '@librechat/client'; +import { activeSubagentPanel } from '~/store/subagents'; +import { adaptLivePersistedActivity } from './adapters'; +import SubagentActivity from './SubagentActivity'; +import { useLocalize } from '~/hooks'; + +/** Public-share fallback for subagent activity already embedded in the shared message payload. */ +export default function SharedSubagentActivityDialog({ shareId }: { shareId?: string }) { + const localize = useLocalize(); + const selected = useRecoilValue(activeSubagentPanel); + const resetSelection = useResetRecoilState(activeSubagentPanel); + const selection = selected?.host === 'share' && selected.shareId === shareId ? selected : null; + const restoreSelectionRef = useRef(selection); + if (selection != null) restoreSelectionRef.current = selection; + const title = + selection?.subagentType === 'self' + ? localize('com_ui_subagent_dialog_title_self') + : localize('com_ui_subagent_dialog_title', { 0: selection?.subagentType ?? '' }); + const activity = useMemo( + () => + adaptLivePersistedActivity({ + title, + prompt: selection?.prompt, + progress: null, + persistedContent: selection?.persistedContent, + legacyOutput: selection?.legacyOutput, + initialProgress: selection?.initialProgress ?? 1, + isSubmitting: false, + runStepStatus: selection?.runStepStatus, + approvalVisibility: 'hidden', + }), + [selection, title], + ); + + const restoreTriggerFocus = useCallback((event: Event) => { + event.preventDefault(); + const selectionToRestore = restoreSelectionRef.current; + if (selectionToRestore == null) return; + requestAnimationFrame(() => { + const trigger = Array.from( + document.querySelectorAll('[data-subagent-tool-call]'), + ).find( + (element) => + element.dataset.subagentToolCall === selectionToRestore.toolCallId && + element.dataset.subagentParentMessage === selectionToRestore.parentMessageId && + element.dataset.subagentPartIndex === String(selectionToRestore.partIndex), + ); + trigger?.focus(); + }); + }, []); + + useEffect(() => () => resetSelection(), [resetSelection]); + useEffect(() => { + if (selected?.host === 'share' && selected.shareId !== shareId) resetSelection(); + }, [resetSelection, selected, shareId]); + + return ( + !open && resetSelection()}> + + + + {activity.title} + + + + + + ); +} diff --git a/client/src/components/Chat/Subagents/SubagentThreadPanel.test.tsx b/client/src/components/Chat/Subagents/SubagentThreadPanel.test.tsx index ede8983cfaf..8cf79862afa 100644 --- a/client/src/components/Chat/Subagents/SubagentThreadPanel.test.tsx +++ b/client/src/components/Chat/Subagents/SubagentThreadPanel.test.tsx @@ -78,6 +78,7 @@ jest.mock('lucide-react', () => ({ })); const selection: ActiveSubagentPanel = { + host: 'conversation', parentConversationId: 'parent-conversation', parentMessageId: 'parent-message', toolCallId: 'tool-call', @@ -179,6 +180,7 @@ describe('SubagentThreadPanel', () => { isReadinessPending: false, }); const foreground: ActiveSubagentPanel = { + host: 'conversation', parentConversationId: 'parent-conversation', parentMessageId: 'parent-message', toolCallId: 'foreground-call', diff --git a/client/src/components/Chat/Subagents/adapters.test.ts b/client/src/components/Chat/Subagents/adapters.test.ts index 4eb639e8ffb..273efb613ab 100644 --- a/client/src/components/Chat/Subagents/adapters.test.ts +++ b/client/src/components/Chat/Subagents/adapters.test.ts @@ -145,6 +145,34 @@ describe('child activity adapters', () => { ); }); + it('keeps shared-message activity read-only by omitting approval controls', () => { + const activity = adaptLivePersistedActivity({ + title: 'researcher', + progress: null, + persistedContent: [ + { + type: ContentTypes.TOOL_CALL, + [ContentTypes.TOOL_CALL]: { + id: 'tool', + name: 'protected_tool', + args: '{}', + output: '', + progress: 0.1, + approval: { expires_at: 123 }, + }, + }, + ] as unknown as TMessageContentParts[], + initialProgress: 1, + isSubmitting: false, + approvalVisibility: 'hidden', + }); + + expect(activity.items[0]).toEqual( + expect.objectContaining({ type: 'tool', name: 'protected_tool' }), + ); + expect(activity.items[0]).not.toHaveProperty('approval'); + }); + it('uses the exact assistant row as terminal authority for an older API response', () => { const oldView = { threadId: 'thread', diff --git a/client/src/components/Chat/Subagents/adapters.ts b/client/src/components/Chat/Subagents/adapters.ts index baa0231089b..55452a82f48 100644 --- a/client/src/components/Chat/Subagents/adapters.ts +++ b/client/src/components/Chat/Subagents/adapters.ts @@ -52,6 +52,7 @@ type ContentToolCall = { const contentPartsToActivity = ( parts: TMessageContentParts[], reasoningVisibility: 'visible' | 'marker', + approvalVisibility: 'visible' | 'hidden', ): ChildActivityItem[] => parts.flatMap((part, index): ChildActivityItem[] => { if (part.type === ContentTypes.TEXT) { @@ -85,7 +86,9 @@ const contentPartsToActivity = ( ...(tool.args == null ? {} : { input: tool.args }), ...(tool.output == null ? {} : { output: tool.output }), status: runStepStatus ?? (completed ? 'completed' : 'running'), - ...(tool.approval == null ? {} : { approval: tool.approval }), + ...(tool.approval == null || approvalVisibility === 'hidden' + ? {} + : { approval: tool.approval }), }, ]; }); @@ -129,11 +132,16 @@ export function adaptLivePersistedActivity(input: { isSubmitting: boolean; runStepStatus?: PartMetadata['runStepStatus']; reasoningVisibility?: 'visible' | 'marker'; + approvalVisibility?: 'visible' | 'hidden'; }): ChildActivity { const persisted = input.persistedContent ?? []; const live = (input.progress?.contentParts ?? []) as TMessageContentParts[]; const parts = persisted.length > 0 ? persisted : live; - const items = contentPartsToActivity(parts, input.reasoningVisibility ?? 'visible'); + const items = contentPartsToActivity( + parts, + input.reasoningVisibility ?? 'visible', + input.approvalVisibility ?? 'visible', + ); if (items.length === 0 && input.legacyOutput != null && input.legacyOutput !== '') { items.push({ type: 'writing', text: input.legacyOutput }); } diff --git a/client/src/components/Share/ShareView.tsx b/client/src/components/Share/ShareView.tsx index 4d5f0c2a6a6..c76bdf4566f 100644 --- a/client/src/components/Share/ShareView.tsx +++ b/client/src/components/Share/ShareView.tsx @@ -19,6 +19,7 @@ import { TooltipAnchor, useToastContext, } from '@librechat/client'; +import SharedSubagentActivityDialog from '~/components/Chat/Subagents/SharedSubagentActivityDialog'; import { cn, DEFAULT_APP_TITLE, getResponseStatus, selectActiveBranchTail } from '~/utils'; import { ThemeSelector, LangSelector } from '~/components/Appearance'; import { ShareMessagesProvider } from './ShareMessagesProvider'; @@ -239,6 +240,7 @@ function SharedView() { {artifactsContainer}
+ ); } diff --git a/client/src/store/subagents.ts b/client/src/store/subagents.ts index 26369a77d6f..8f6d93ef8f9 100644 --- a/client/src/store/subagents.ts +++ b/client/src/store/subagents.ts @@ -46,6 +46,8 @@ export interface SubagentProgress { /** One child invocation selected for the shared read-only activity panel. */ export type ActiveSubagentPanel = { + host: 'conversation' | 'share'; + shareId?: string; parentConversationId: string; parentMessageId: string; toolCallId: string; diff --git a/packages/api/src/agents/view.spec.ts b/packages/api/src/agents/view.spec.ts index f16f5b9c644..9849af90e4a 100644 --- a/packages/api/src/agents/view.spec.ts +++ b/packages/api/src/agents/view.spec.ts @@ -5,6 +5,7 @@ import { createSubagentThreadViewHandler, SUBAGENT_THREAD_VIEW_LIMITS } from './ jest.mock('@librechat/data-schemas', () => ({ CLIENT_MESSAGE_SELECT: '-_id -user', + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT: 256 * 1024, logger: { error: jest.fn() }, })); @@ -259,6 +260,34 @@ describe('subagent thread parent-scoped view', () => { expect(JSON.stringify(json.mock.calls[0][0])).not.toContain('Wrong task.'); }); + it('falls back to the bounded final message when storage omits an oversized transcript', async () => { + const selected = { + ...message('task-1:assistant', 'completed'), + text: 'The bounded final answer.', + subagentTranscriptProjectionTruncated: true, + } as IMessage & { subagentTranscriptProjectionTruncated: boolean }; + const handler = createSubagentThreadViewHandler({ + getConvoOwnership: jest.fn().mockResolvedValue(parent), + getSubagentThreadForParent: jest + .fn() + .mockResolvedValue({ ...child, subagentThreadLease: undefined }), + getMessagesForSubagentThreadView: jest.fn().mockResolvedValue([selected]), + }); + const { response, json } = createResponse(); + + await handler(createRequest({}, { taskId: 'task-1' }), response); + + const view = json.mock.calls[0][0]; + expect(view).toEqual( + expect.objectContaining({ + activity: [], + activityTruncated: true, + messages: [expect.objectContaining({ text: 'The bounded final answer.' })], + }), + ); + expect(JSON.stringify(view)).not.toContain('subagentTranscript'); + }); + it('bounds the complete UTF-8 response while retaining the newest history', async () => { const getConvoOwnership = jest.fn().mockResolvedValue(parent); const messages = Array.from( diff --git a/packages/api/src/agents/view.ts b/packages/api/src/agents/view.ts index a850a77aa19..92503a32407 100644 --- a/packages/api/src/agents/view.ts +++ b/packages/api/src/agents/view.ts @@ -1,4 +1,4 @@ -import { logger } from '@librechat/data-schemas'; +import { logger, SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT } from '@librechat/data-schemas'; import type { ConversationMethods, MessageMethods, @@ -204,23 +204,30 @@ export function createSubagentThreadViewHandler(deps: SubagentThreadViewDependen child.subagentThreadLease != null && child.subagentThreadLease.expiresAt > now ? child.subagentThreadLease.taskId : undefined; - const selectedTranscript = + const selectedMessage = requestedTaskId == null ? undefined - : newestFirst.find((message) => message.messageId === `${requestedTaskId}:assistant`) - ?.subagentTranscript; + : newestFirst.find((message) => message.messageId === `${requestedTaskId}:assistant`); + const selectedTranscript = selectedMessage?.subagentTranscript; const selectedInput = requestedTaskId == null ? undefined : newestFirst.find((message) => message.messageId === `${requestedTaskId}:user`); - const projectedActivity = - selectedTranscript != null && selectedTranscript.taskId === requestedTaskId - ? projectSubagentActivity( - selectedTranscript.messagesJson, - selectedTranscript.mode, - selectedInput?.textProjectionTruncated === true ? undefined : selectedInput?.text, - ) - : { activity: [], truncated: selectedTranscript != null }; + let projectedActivity: ReturnType = { + activity: [], + truncated: false, + }; + if (selectedMessage?.subagentTranscriptProjectionTruncated === true) { + projectedActivity = { activity: [], truncated: true }; + } else if (selectedTranscript != null && selectedTranscript.taskId === requestedTaskId) { + projectedActivity = projectSubagentActivity( + selectedTranscript.messagesJson, + selectedTranscript.mode, + selectedInput?.textProjectionTruncated === true ? undefined : selectedInput?.text, + ); + } else if (selectedTranscript != null) { + projectedActivity = { activity: [], truncated: true }; + } const projectedNewestFirst: SubagentThreadMessage[] = []; let remainingTextBytes = MAX_RESPONSE_TEXT_BYTES; for (const message of newestFirst) { @@ -272,6 +279,7 @@ export const SUBAGENT_THREAD_VIEW_LIMITS: Readonly<{ responseBytes: number; activityItems: number; activityBytes: number; + activitySourceBytes: number; }> = { messages: MAX_THREAD_MESSAGES, messageTextBytes: MAX_MESSAGE_TEXT_BYTES, @@ -279,4 +287,5 @@ export const SUBAGENT_THREAD_VIEW_LIMITS: Readonly<{ responseBytes: MAX_RESPONSE_BYTES, activityItems: SUBAGENT_ACTIVITY_LIMITS.items, activityBytes: SUBAGENT_ACTIVITY_LIMITS.bytes, + activitySourceBytes: SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, }; diff --git a/packages/data-schemas/src/index.ts b/packages/data-schemas/src/index.ts index 9220c5f67a3..b1a56dc206a 100644 --- a/packages/data-schemas/src/index.ts +++ b/packages/data-schemas/src/index.ts @@ -8,6 +8,7 @@ export { createModels } from './models'; export { createMethods, CLIENT_MESSAGE_SELECT, + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, RoleConflictError, DEFAULT_REFRESH_TOKEN_EXPIRY, DEFAULT_SESSION_EXPIRY, diff --git a/packages/data-schemas/src/methods/index.ts b/packages/data-schemas/src/methods/index.ts index d90f48ed170..de61ec3e64f 100644 --- a/packages/data-schemas/src/methods/index.ts +++ b/packages/data-schemas/src/methods/index.ts @@ -47,6 +47,7 @@ import { createConversationTagMethods, type ConversationTagMethods } from './con import { createMessageMethods, CLIENT_MESSAGE_SELECT, + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, type MessageMethods, type SubagentThreadViewMessageRecord, type SubagentTaskResultClaim, @@ -146,7 +147,7 @@ export { }; export { tokenValues, cacheTokenValues, premiumTokenValues, defaultRate, createTxMethods }; export { permissionBitSupersets }; -export { CLIENT_MESSAGE_SELECT }; +export { CLIENT_MESSAGE_SELECT, SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT }; export { partitionIssues, validateSkillName, diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 390f5efcf11..30d3fe31056 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -3,7 +3,11 @@ import { v4 as uuidv4 } from 'uuid'; import { RetentionMode } from 'librechat-data-provider'; import { MongoMemoryServer } from 'mongodb-memory-server'; import type { IMessage } from '..'; -import { createMessageMethods, CLIENT_MESSAGE_SELECT } from './message'; +import { + createMessageMethods, + CLIENT_MESSAGE_SELECT, + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, +} from './message'; import { tenantStorage, runAsSystem } from '~/config/tenantContext'; import { createModels } from '../models'; import logger from '~/config/winston'; @@ -750,6 +754,39 @@ describe('Message Operations', () => { expect(messages[0]).toHaveProperty('messageId', 'task-a:assistant'); expect(messages[0]).toHaveProperty('subagentTranscript.taskId', 'task-a'); }); + + it('omits an oversized private transcript before returning the application result', async () => { + const conversationId = uuidv4(); + await saveMessage(mockCtx, { + messageId: 'task-large:assistant', + conversationId, + text: 'The bounded public answer remains available.', + user: 'user123', + subagentTranscript: { + taskId: 'task-large', + mode: 'append', + messagesJson: JSON.stringify([ + { + type: 'ai', + data: { content: 'x'.repeat(SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT + 1) }, + }, + ]), + }, + }); + + const messages = await getMessagesForSubagentThreadView({ + user: 'user123', + conversationId, + limit: 1, + textCodePointLimit: 8_192, + taskId: 'task-large', + }); + + expect(messages).toHaveLength(1); + expect(messages[0].text).toBe('The bounded public answer remains available.'); + expect(messages[0]).not.toHaveProperty('subagentTranscript'); + expect(messages[0].subagentTranscriptProjectionTruncated).toBe(true); + }); }); describe('deleteMessages', () => { diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index dd443669b63..4d466decdcb 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -9,6 +9,14 @@ import logger from '~/config/winston'; /** Simple UUID v4 regex to replace zod validation */ const UUID_REGEX = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; +/** + * Maximum private transcript JSON that may cross the MongoDB projection seam + * for the bounded public subagent-activity view. This gives the sanitizer + * enough source headroom while preventing multi-megabyte transcripts from + * being materialized merely to produce a 64 KiB public activity response. + */ +export const SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT: number = 256 * 1024; + /** * Exclusion projection for message reads that feed the chat client (the * conversation GET and shared-link reads). Every excluded field is either @@ -63,7 +71,10 @@ export type SubagentThreadViewMessageRecord = Pick< | 'error' | 'subagentTranscript' | 'subagentTask' -> & { textProjectionTruncated?: boolean }; +> & { + textProjectionTruncated?: boolean; + subagentTranscriptProjectionTruncated?: boolean; +}; export interface MessageMethods { saveMessage( @@ -751,6 +762,21 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa }): Promise { try { const Message = mongoose.models.Message as Model; + const selectedAssistantMessageId = + input.taskId == null ? undefined : `${input.taskId}:assistant`; + const transcriptJsonBytes = { + $strLenBytes: { + $convert: { + input: '$subagentTranscript.messagesJson', + to: 'string', + onError: '', + onNull: '', + }, + }, + }; + const transcriptIsString = { + $eq: [{ $type: '$subagentTranscript.messagesJson' }, 'string'], + }; return await Message.aggregate([ { $match: { @@ -770,6 +796,16 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa }, { $sort: { createdAt: -1, _id: -1 } }, { $limit: input.limit }, + ...(input.taskId == null + ? [] + : [ + { + $set: { + _subagentTranscriptSourceBytes: transcriptJsonBytes, + _subagentTranscriptSourceIsString: transcriptIsString, + }, + }, + ]), { $project: { _id: 0, @@ -789,8 +825,48 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa : { subagentTranscript: { $cond: [ - { $eq: ['$messageId', `${input.taskId}:assistant`] }, - '$subagentTranscript', + { + $and: [ + { $eq: ['$messageId', selectedAssistantMessageId] }, + '$_subagentTranscriptSourceIsString', + { + $lte: [ + '$_subagentTranscriptSourceBytes', + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, + ], + }, + ], + }, + { + taskId: '$subagentTranscript.taskId', + mode: '$subagentTranscript.mode', + messagesJson: '$subagentTranscript.messagesJson', + }, + '$$REMOVE', + ], + }, + subagentTranscriptProjectionTruncated: { + $cond: [ + { + $and: [ + { $eq: ['$messageId', selectedAssistantMessageId] }, + { + $ne: [{ $type: '$subagentTranscript.messagesJson' }, 'missing'], + }, + { + $or: [ + { $eq: ['$_subagentTranscriptSourceIsString', false] }, + { + $gt: [ + '$_subagentTranscriptSourceBytes', + SUBAGENT_TRANSCRIPT_SOURCE_BYTE_LIMIT, + ], + }, + ], + }, + ], + }, + true, '$$REMOVE', ], },