diff --git a/.env.example b/.env.example index 3622c660c7..0cd016b0d3 100644 --- a/.env.example +++ b/.env.example @@ -40,10 +40,10 @@ DOMAIN_SERVER=http://localhost:3080 ADMIN_PANEL_URL= # Session encryption key for the bundled admin panel (min 32 characters). -# The admin panel ships enabled in docker-compose/deploy-compose; the value -# below is a development default. Generate a unique one for production: +# Required when using the bundled admin panel in docker-compose/deploy-compose. +# Generate a unique value before starting the stack: # openssl rand -hex 32 -ADMIN_PANEL_SESSION_SECRET=0f7114d0b0b192b59fde7e3b730949b27d49a8717fa32ad3167d94c7684fede3 +ADMIN_PANEL_SESSION_SECRET= # Host port for the bundled admin panel (default docker-compose only). # In deploy-compose the panel is served at http://admin.localhost via nginx. diff --git a/api/server/controllers/ContextProjectionController.js b/api/server/controllers/ContextProjectionController.js index eaf9592e73..9c56b2ae34 100644 --- a/api/server/controllers/ContextProjectionController.js +++ b/api/server/controllers/ContextProjectionController.js @@ -18,7 +18,11 @@ async function contextProjectionController(req, res) { return; } const projection = await resolveContextProjection( - { userId: req.user?.id, getMessages: db.getMessages }, + { + userId: req.user?.id, + getMessages: db.getMessages, + getMessageTextStats: db.getMessageTextStats, + }, params, ); res.json(projection ?? null); diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index b6032530c9..ce89addce8 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -1450,11 +1450,12 @@ class AgentClient extends BaseClient { }); } } catch (err) { - logger.error( - '[api/server/controllers/agents/client.js #sendCompletion] Operation aborted', - err, - ); - if (!abortController.signal.aborted) { + if (abortController.signal.aborted) { + logger.debug( + '[api/server/controllers/agents/client.js #sendCompletion] Operation aborted by user', + { conversationId: this.conversationId, name: err?.name, code: err?.code }, + ); + } else { logger.error( '[api/server/controllers/agents/client.js #sendCompletion] Unhandled error type', err, diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js index 35bd53c579..feed002b0e 100644 --- a/api/server/middleware/abortMiddleware.js +++ b/api/server/middleware/abortMiddleware.js @@ -15,6 +15,36 @@ const { sendError } = require('~/server/middleware/error'); const { abortRun } = require('./abortRun'); const db = require('~/models'); +/** + * @param {Error | unknown} error + * @returns {boolean} + */ +const isAbortError = (error) => { + const visited = new Set(); + let current = error; + + while (current && typeof current === 'object' && !visited.has(current)) { + visited.add(current); + + const errorName = current.name; + const errorCode = current.code; + const errorMessage = typeof current.message === 'string' ? current.message : ''; + + if ( + errorName === 'AbortError' || + errorCode === 'ABORT_ERR' || + errorMessage.includes('AbortError') || + /(?:operation|request|stream) was aborted/i.test(errorMessage) + ) { + return true; + } + + current = current.cause; + } + + return false; +}; + /** * Spend tokens for all models from collected usage. * This handles both sequential and parallel agent execution. @@ -200,18 +230,26 @@ const handleAbort = function () { * @returns {Promise} */ const handleAbortError = async (res, req, error, data) => { + const { sender, conversationId, messageId, parentMessageId, userMessageId, partialText } = data; + if (error?.message?.includes('base64')) { logger.error('[handleAbortError] Error in base64 encoding', { ...error, stack: smartTruncateText(error?.stack, 1000), message: truncateText(error.message, 350), }); + } else if (isAbortError(error)) { + logger.debug('[handleAbortError] AI response aborted by user', { + conversationId, + code: error?.code, + name: error?.name, + message: truncateText(error?.message ?? 'AbortError', 350), + }); } else { logger.error('[handleAbortError] AI response error; aborting request:', error); } - const { sender, conversationId, messageId, parentMessageId, userMessageId, partialText } = data; - if (error.stack && error.stack.includes('google')) { + if (error?.stack && error.stack.includes('google')) { logger.warn( `AI Response error for conversation ${conversationId} likely caused by Google censor/filter`, ); diff --git a/api/server/middleware/abortMiddleware.spec.js b/api/server/middleware/abortMiddleware.spec.js index a4ce85674b..06e434065a 100644 --- a/api/server/middleware/abortMiddleware.spec.js +++ b/api/server/middleware/abortMiddleware.spec.js @@ -73,7 +73,18 @@ jest.mock('./abortRun', () => ({ abortRun: jest.fn(), })); -const { spendCollectedUsage } = require('./abortMiddleware'); +const { logger } = require('@librechat/data-schemas'); +const { sendError } = require('~/server/middleware/error'); +const { handleAbortError, spendCollectedUsage } = require('./abortMiddleware'); + +const buildAbortRequest = () => ({ + body: { + model: 'gpt-4', + }, + user: { + id: 'user-123', + }, +}); describe('abortMiddleware - spendCollectedUsage', () => { beforeEach(() => { @@ -237,3 +248,65 @@ describe('abortMiddleware - spendCollectedUsage', () => { }); }); }); + +describe('abortMiddleware - handleAbortError', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it.each([ + [ + 'native DOMException AbortError', + new DOMException('The operation was aborted', 'AbortError'), + 'AbortError', + ], + [ + 'wrapped AbortError message', + new Error('SSE stream disconnected: AbortError: The operation was aborted'), + 'Error', + ], + [ + 'cause-nested AbortError', + new Error('Request failed', { + cause: new DOMException('The operation was aborted', 'AbortError'), + }), + 'Error', + ], + ])('logs a %s as a debug event instead of an error', async (_label, error, name) => { + await handleAbortError({}, buildAbortRequest(), error, { + sender: 'AI', + conversationId: 'convo-123', + messageId: 'message-123', + parentMessageId: 'parent-123', + userMessageId: 'user-message-123', + }); + + expect(logger.error).not.toHaveBeenCalled(); + expect(logger.debug).toHaveBeenCalledWith('[handleAbortError] AI response aborted by user', { + conversationId: 'convo-123', + code: error.code, + name, + message: error.message, + }); + expect(sendError).toHaveBeenCalledTimes(1); + }); + + it('keeps unexpected generation errors classified as errors', async () => { + const error = new Error('Provider failed'); + + await handleAbortError({}, buildAbortRequest(), error, { + sender: 'AI', + conversationId: 'convo-123', + messageId: 'message-123', + parentMessageId: 'parent-123', + userMessageId: 'user-message-123', + }); + + expect(logger.error).toHaveBeenCalledWith( + '[handleAbortError] AI response error; aborting request:', + error, + ); + expect(logger.debug).not.toHaveBeenCalled(); + expect(sendError).toHaveBeenCalledTimes(1); + }); +}); diff --git a/api/server/middleware/limiters/contextProjectionLimiter.js b/api/server/middleware/limiters/contextProjectionLimiter.js new file mode 100644 index 0000000000..1f70c7ea8e --- /dev/null +++ b/api/server/middleware/limiters/contextProjectionLimiter.js @@ -0,0 +1,19 @@ +const rateLimit = require('express-rate-limit'); +const { limiterCache } = require('@librechat/api'); + +const { CONTEXT_PROJECTION_WINDOW = 1, CONTEXT_PROJECTION_MAX = 20 } = process.env; + +const windowMs = (parseInt(CONTEXT_PROJECTION_WINDOW, 10) || 1) * 60 * 1000; +const max = parseInt(CONTEXT_PROJECTION_MAX, 10) || 20; + +const contextProjectionLimiter = rateLimit({ + windowMs, + max, + handler: (_req, res) => { + res.status(429).json({ message: 'Too many context projection requests. Try again later' }); + }, + keyGenerator: (req) => req.user?.id, + store: limiterCache('context_projection_limiter'), +}); + +module.exports = contextProjectionLimiter; diff --git a/api/server/middleware/limiters/index.js b/api/server/middleware/limiters/index.js index 4a569e2698..19f246d039 100644 --- a/api/server/middleware/limiters/index.js +++ b/api/server/middleware/limiters/index.js @@ -9,6 +9,7 @@ const registerLimiter = require('./registerLimiter'); const toolCallLimiter = require('./toolCallLimiter'); const messageLimiters = require('./messageLimiters'); const promptUsageLimiter = require('./promptUsageLimiter'); +const contextProjectionLimiter = require('./contextProjectionLimiter'); const verifyEmailLimiter = require('./verifyEmailLimiter'); const resetPasswordLimiter = require('./resetPasswordLimiter'); const twoFactorTempLimiter = require('./twoFactorTempLimiter'); @@ -22,6 +23,7 @@ module.exports = { loginLimiter, registerLimiter, toolCallLimiter, + contextProjectionLimiter, createTTSLimiters, createSTTLimiters, verifyEmailLimiter, diff --git a/api/server/middleware/optionalShareFileAuth.js b/api/server/middleware/optionalShareFileAuth.js index f5c4ac07a5..bebf087d20 100644 --- a/api/server/middleware/optionalShareFileAuth.js +++ b/api/server/middleware/optionalShareFileAuth.js @@ -3,9 +3,9 @@ const jwt = require('jsonwebtoken'); const { isEnabled } = require('@librechat/api'); const { logger, runAsSystem } = require('@librechat/data-schemas'); const { SystemRoles } = require('librechat-data-provider'); -const { getUserById } = require('~/models'); +const { getUserById, findSession } = require('~/models'); -const verifyRefreshToken = (token) => { +const verifySignedUserId = (token) => { try { const payload = jwt.verify(token, process.env.JWT_REFRESH_SECRET); return typeof payload?.id === 'string' ? payload.id : null; @@ -14,13 +14,36 @@ const verifyRefreshToken = (token) => { } }; +const getRefreshTokenUserId = async (token) => { + const userId = verifySignedUserId(token); + if (!userId) { + return null; + } + + const session = await runAsSystem(() => findSession({ userId, refreshToken: token })); + return session ? userId : null; +}; + +const getOpenIdUserId = (parsed, req) => { + if (parsed.token_provider !== 'openid' || !isEnabled(process.env.OPENID_REUSE_TOKENS)) { + return null; + } + + const sessionRefreshToken = req.session?.openidTokens?.refreshToken; + if (!parsed.refreshToken || parsed.refreshToken !== sessionRefreshToken) { + return null; + } + + return verifySignedUserId(parsed.openid_user_id); +}; + /** * Fallback auth for share file routes that are hit by ``/anchor requests, * which can't carry the bearer access token. Resolves the viewer from the - * `refreshToken` cookie (or the signed `openid_user_id` cookie) — the same - * mechanism secure image links use — so non-public shared links can authorize - * the viewer's ACL. Never blocks: on any failure it leaves `req.user` unset and - * lets `canAccessSharedLink` decide (public access, 401, or 403). + * `refreshToken` cookie (or an active OpenID session plus signed `openid_user_id` + * cookie) so non-public shared links can authorize the viewer's ACL. Never + * blocks: on any failure it leaves `req.user` unset and lets + * `canAccessSharedLink` decide (public access, 401, or 403). */ const optionalShareFileAuth = async (req, res, next) => { if (req.user) { @@ -34,22 +57,17 @@ const optionalShareFileAuth = async (req, res, next) => { } const parsed = cookie.parse(cookieHeader); - const useOpenId = - parsed.token_provider === 'openid' && isEnabled(process.env.OPENID_REUSE_TOKENS); - const token = useOpenId ? parsed.openid_user_id : parsed.refreshToken; - if (!token) { - return next(); - } - - const userId = verifyRefreshToken(token); + const userId = + getOpenIdUserId(parsed, req) || + (parsed.refreshToken ? await getRefreshTokenUserId(parsed.refreshToken) : null); if (!userId) { return next(); } // Resolve in system context: this runs before canAccessSharedLink establishes // the share tenant, so under strict tenant isolation a tenant-scoped User - // query would otherwise throw. The viewer's id comes from their own verified - // refresh token; the share's tenant-scoped ACL check still gates access. + // query would otherwise throw. The viewer's id comes from verified, active + // cookie auth; the share's tenant-scoped ACL check still gates access. const user = await runAsSystem(() => getUserById(userId, '-password -__v -totpSecret -backupCodes'), ); diff --git a/api/server/middleware/optionalShareFileAuth.spec.js b/api/server/middleware/optionalShareFileAuth.spec.js index 96e147198b..ffc0cf5cfc 100644 --- a/api/server/middleware/optionalShareFileAuth.spec.js +++ b/api/server/middleware/optionalShareFileAuth.spec.js @@ -1,14 +1,27 @@ const mockVerify = jest.fn(); const mockGetUserById = jest.fn(); +const mockFindSession = jest.fn(); +const mockRunAsSystem = jest.fn((fn) => fn()); jest.mock('jsonwebtoken', () => ({ verify: (...args) => mockVerify(...args) })); -jest.mock('@librechat/api', () => ({ isEnabled: (v) => v === 'true' || v === true })); -jest.mock('@librechat/data-schemas', () => ({ - logger: { warn: jest.fn(), error: jest.fn() }, - runAsSystem: (fn) => fn(), +jest.mock('@librechat/api', () => ({ isEnabled: (v) => v === 'true' || v === true }), { + virtual: true, +}); +jest.mock( + '@librechat/data-schemas', + () => ({ + logger: { warn: jest.fn(), error: jest.fn() }, + runAsSystem: (...args) => mockRunAsSystem(...args), + }), + { virtual: true }, +); +jest.mock('librechat-data-provider', () => ({ SystemRoles: { USER: 'USER' } }), { + virtual: true, +}); +jest.mock('~/models', () => ({ + getUserById: (...args) => mockGetUserById(...args), + findSession: (...args) => mockFindSession(...args), })); -jest.mock('librechat-data-provider', () => ({ SystemRoles: { USER: 'USER' } })); -jest.mock('~/models', () => ({ getUserById: (...args) => mockGetUserById(...args) })); const optionalShareFileAuth = require('./optionalShareFileAuth'); @@ -30,20 +43,25 @@ describe('optionalShareFileAuth', () => { expect(next).toHaveBeenCalledTimes(1); expect(mockVerify).not.toHaveBeenCalled(); expect(mockGetUserById).not.toHaveBeenCalled(); + expect(mockFindSession).not.toHaveBeenCalled(); }); - it('resolves the viewer from a valid refreshToken cookie', async () => { + it('resolves the viewer from a valid refreshToken cookie with a live session', async () => { mockVerify.mockReturnValue({ id: 'viewer-1' }); + mockFindSession.mockResolvedValue({ _id: 'session-1' }); mockGetUserById.mockResolvedValue({ _id: 'viewer-1', role: 'USER' }); const req = { headers: { cookie: 'refreshToken=good.jwt' } }; const next = await run(req); expect(next).toHaveBeenCalledTimes(1); expect(mockVerify).toHaveBeenCalledWith('good.jwt', 'test-secret'); + expect(mockFindSession).toHaveBeenCalledWith({ userId: 'viewer-1', refreshToken: 'good.jwt' }); + expect(mockRunAsSystem).toHaveBeenCalledTimes(2); expect(req.user).toMatchObject({ id: 'viewer-1', role: 'USER' }); }); it('defaults the role to USER when the record has none', async () => { mockVerify.mockReturnValue({ id: 'viewer-2' }); + mockFindSession.mockResolvedValue({ _id: 'session-2' }); mockGetUserById.mockResolvedValue({ _id: 'viewer-2' }); const req = { headers: { cookie: 'refreshToken=good.jwt' } }; await run(req); @@ -58,6 +76,21 @@ describe('optionalShareFileAuth', () => { expect(mockGetUserById).not.toHaveBeenCalled(); }); + it('leaves req.user unset when the refresh token has no live session', async () => { + mockVerify.mockReturnValue({ id: 'viewer-3' }); + mockFindSession.mockResolvedValue(null); + const req = { headers: { cookie: 'refreshToken=revoked.jwt' } }; + const next = await run(req); + expect(next).toHaveBeenCalledTimes(1); + expect(req.user).toBeUndefined(); + expect(mockFindSession).toHaveBeenCalledWith({ + userId: 'viewer-3', + refreshToken: 'revoked.jwt', + }); + expect(mockRunAsSystem).toHaveBeenCalledTimes(1); + expect(mockGetUserById).not.toHaveBeenCalled(); + }); + it('leaves req.user unset when the token is invalid', async () => { mockVerify.mockImplementation(() => { throw new Error('bad token'); @@ -69,16 +102,35 @@ describe('optionalShareFileAuth', () => { expect(mockGetUserById).not.toHaveBeenCalled(); }); - it('uses the signed openid_user_id cookie for OpenID-reuse sessions', async () => { + it('uses the signed openid_user_id cookie only for active OpenID-reuse sessions', async () => { process.env.OPENID_REUSE_TOKENS = 'true'; mockVerify.mockReturnValue({ id: 'oidc-1' }); mockGetUserById.mockResolvedValue({ _id: 'oidc-1', role: 'USER' }); const req = { - headers: { cookie: 'token_provider=openid; openid_user_id=signed.jwt' }, + headers: { + cookie: 'token_provider=openid; refreshToken=stored-refresh; openid_user_id=signed.jwt', + }, + session: { openidTokens: { refreshToken: 'stored-refresh' } }, }; await run(req); expect(mockVerify).toHaveBeenCalledWith('signed.jwt', 'test-secret'); + expect(mockFindSession).not.toHaveBeenCalled(); expect(req.user).toMatchObject({ id: 'oidc-1' }); delete process.env.OPENID_REUSE_TOKENS; }); + + it('leaves req.user unset for OpenID-reuse cookies without an active matching session', async () => { + process.env.OPENID_REUSE_TOKENS = 'true'; + mockVerify.mockReturnValue({ id: 'oidc-2' }); + const req = { + headers: { + cookie: 'token_provider=openid; refreshToken=stale-refresh; openid_user_id=signed.jwt', + }, + session: { openidTokens: { refreshToken: 'current-refresh' } }, + }; + await run(req); + expect(req.user).toBeUndefined(); + expect(mockGetUserById).not.toHaveBeenCalled(); + delete process.env.OPENID_REUSE_TOKENS; + }); }); diff --git a/api/server/routes/__tests__/config.spec.js b/api/server/routes/__tests__/config.spec.js index e0fc200486..af82399e76 100644 --- a/api/server/routes/__tests__/config.spec.js +++ b/api/server/routes/__tests__/config.spec.js @@ -185,18 +185,22 @@ describe('GET /api/config', () => { expect(response.body).not.toHaveProperty('customFooter'); }); - it('should include public share footer fields when share context is requested', async () => { + it('should not include share-only fields when share context is requested', async () => { process.env.ANALYTICS_GTM_ID = 'GTM-XYZ'; process.env.CUSTOM_FOOTER = 'public footer text'; process.env.HELP_AND_FAQ_URL = 'https://internal.example.com/faq'; + process.env.SANDPACK_BUNDLER_URL = 'https://bundler.test'; + process.env.SANDPACK_STATIC_BUNDLER_URL = 'https://static-bundler.test'; mockGetAppConfig.mockResolvedValue(baseAppConfig); const app = createApp(null); const response = await request(app).get('/api/config?context=share'); expect(response.statusCode).toBe(200); - expect(response.body.analyticsGtmId).toBe('GTM-XYZ'); - expect(response.body.customFooter).toBe('public footer text'); + expect(response.body).not.toHaveProperty('analyticsGtmId'); + expect(response.body).not.toHaveProperty('customFooter'); + expect(response.body).not.toHaveProperty('bundlerURL'); + expect(response.body).not.toHaveProperty('staticBundlerURL'); expect(response.body).not.toHaveProperty('helpAndFaqURL'); expect(response.body).not.toHaveProperty('allowAccountDeletion'); }); diff --git a/api/server/routes/__tests__/share.spec.js b/api/server/routes/__tests__/share.spec.js index ecbd820ae6..c9d6cb0037 100644 --- a/api/server/routes/__tests__/share.spec.js +++ b/api/server/routes/__tests__/share.spec.js @@ -6,6 +6,10 @@ const mockGetSharedLinkExpiration = jest.fn(); const mockGrantCreationPermissions = jest.fn(); const mockUpdateSharedLinkPermissionsExpiration = jest.fn(); const mockSharedLinksAccess = jest.fn((_req, _res, next) => next()); +const mockBuildSharedLinkStartupPayload = jest.fn(); +const mockCanAccessSharedLink = jest.fn((_req, _res, next) => next()); +const mockGetAppConfig = jest.fn(); +const mockGetTenantId = jest.fn(() => undefined); jest.mock('@librechat/api', () => ({ isEnabled: jest.fn(() => true), @@ -16,6 +20,7 @@ jest.mock('@librechat/api', () => ({ ensureLinkPermissions: jest.fn(), isFileSnapshotEnabled: jest.fn(() => true), isFileSnapshotKillSwitchActive: jest.fn(() => false), + buildSharedLinkStartupPayload: (...args) => mockBuildSharedLinkStartupPayload(...args), deleteSharedLinkWithCleanup: jest.fn(), getSharedLinkExpiration: (...args) => mockGetSharedLinkExpiration(...args), isActiveExpirationDate: jest.fn((expiredAt) => expiredAt > new Date()), @@ -23,9 +28,11 @@ jest.mock('@librechat/api', () => ({ jest.mock('@librechat/data-schemas', () => ({ logger: { error: jest.fn(), warn: jest.fn() }, + getTenantId: (...args) => mockGetTenantId(...args), createTempChatExpirationDate: jest.fn(() => new Date('2030-01-01T00:00:00.000Z')), runAsSystem: jest.fn((fn) => fn()), tenantStorage: { run: jest.fn((_ctx, fn) => fn()) }, + SYSTEM_TENANT_ID: '__SYSTEM__', })); jest.mock('librechat-data-provider', () => ({ @@ -84,11 +91,19 @@ jest.mock('~/server/utils/files', () => ({ getContentDisposition: jest.fn((name, disposition = 'attachment') => `${disposition}; ${name}`), })); -jest.mock('~/server/middleware/canAccessSharedLink', () => (_req, _res, next) => next()); +jest.mock( + '~/server/middleware/canAccessSharedLink', + () => + (...args) => + mockCanAccessSharedLink(...args), +); jest.mock('~/server/middleware/optionalShareFileAuth', () => (_req, _res, next) => next()); jest.mock('~/server/middleware/optionalJwtAuth', () => (req, _res, next) => next()); jest.mock('~/server/middleware/requireJwtAuth', () => (req, res, next) => next()); jest.mock('~/server/middleware/config/app', () => (_req, _res, next) => next()); +jest.mock('~/server/services/Config/app', () => ({ + getAppConfig: (...args) => mockGetAppConfig(...args), +})); const { Readable } = require('stream'); const { RetentionMode } = require('librechat-data-provider'); @@ -129,9 +144,22 @@ const buildApp = ({ retentionMode = RetentionMode.TEMPORARY } = {}) => { return app; }; -describe('share routes retention', () => { +describe('share routes', () => { beforeEach(() => { jest.clearAllMocks(); + mockGetTenantId.mockReturnValue(undefined); + mockGetAppConfig.mockResolvedValue({ + interfaceConfig: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + }, + }); + mockBuildSharedLinkStartupPayload.mockReturnValue({ + appTitle: 'Shared Chat', + bundlerURL: 'https://bundler.example.com', + interface: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + }, + }); getRoleByName.mockResolvedValue({ permissions: { SHARED_LINKS: { @@ -142,6 +170,45 @@ describe('share routes retention', () => { mockGrantCreationPermissions.mockResolvedValue(undefined); }); + it('serves shared startup config after shared-link access is granted', async () => { + const response = await request(buildApp()).get('/api/share/share-123/config'); + + expect(response.status).toBe(200); + expect(response.headers['cache-control']).toBe('private, no-store'); + expect(mockCanAccessSharedLink).toHaveBeenCalled(); + expect(mockGetAppConfig).toHaveBeenCalledWith({ baseOnly: true }); + expect(mockBuildSharedLinkStartupPayload).toHaveBeenCalledWith({ + interfaceConfig: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + }, + }); + expect(response.body).toEqual({ + appTitle: 'Shared Chat', + bundlerURL: 'https://bundler.example.com', + interface: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + }, + }); + }); + + it('uses tenant-scoped app config for shared startup config when tenant context is present', async () => { + mockGetTenantId.mockReturnValue('tenant-abc'); + + const response = await request(buildApp()).get('/api/share/share-123/config'); + + expect(response.status).toBe(200); + expect(mockGetAppConfig).toHaveBeenCalledWith({ tenantId: 'tenant-abc' }); + }); + + it('uses base app config for shared startup config in system context', async () => { + mockGetTenantId.mockReturnValue('__SYSTEM__'); + + const response = await request(buildApp()).get('/api/share/share-123/config'); + + expect(response.status).toBe(200); + expect(mockGetAppConfig).toHaveBeenCalledWith({ baseOnly: true }); + }); + it('prevents successful shared message responses from being cached', async () => { getSharedMessages.mockResolvedValue({ shareId: 'share-123', messages: [] }); diff --git a/api/server/routes/config.js b/api/server/routes/config.js index 19f25b844f..f6eb66374e 100644 --- a/api/server/routes/config.js +++ b/api/server/routes/config.js @@ -107,9 +107,9 @@ function buildPreLoginPayload() { } /** - * Public share fields rendered by `client/src/components/Share/ShareView.tsx`. - * They remain off the default anonymous config used by login screens, and are - * exposed to anonymous callers only when the client asks for share context. + * Fields shared by authenticated chat and share-view config. Anonymous share + * views receive these through `/api/share/:shareId/config` after share access + * checks, not through the generic startup config endpoint. */ function buildPublicSharePayload() { /** @type {Partial} */ @@ -215,7 +215,6 @@ router.get('/', async function (req, res) { /** @type {Partial} */ const payload = { ...preLoginPayload, - ...(req.query.context === 'share' ? publicSharePayload : {}), socialLogins: baseConfig?.registration?.socialLogins ?? defaultSocialLogins, turnstile: baseConfig?.turnstileConfig, ...(rum ? { rum } : {}), diff --git a/api/server/routes/endpoints.js b/api/server/routes/endpoints.js index b11de153df..ea55a9e54a 100644 --- a/api/server/routes/endpoints.js +++ b/api/server/routes/endpoints.js @@ -4,11 +4,18 @@ const configMiddleware = require('~/server/middleware/config/app'); const endpointController = require('~/server/controllers/EndpointController'); const tokenConfigController = require('~/server/controllers/TokenConfigController'); const contextProjectionController = require('~/server/controllers/ContextProjectionController'); +const { contextProjectionLimiter } = require('~/server/middleware/limiters'); const router = express.Router(); /** Auth required for role/tenant-scoped endpoint config resolution. */ router.get('/', requireJwtAuth, endpointController); router.get('/token-config', requireJwtAuth, configMiddleware, tokenConfigController); -router.post('/context-projection', requireJwtAuth, configMiddleware, contextProjectionController); +router.post( + '/context-projection', + requireJwtAuth, + contextProjectionLimiter, + configMiddleware, + contextProjectionController, +); module.exports = router; diff --git a/api/server/routes/share.js b/api/server/routes/share.js index 09f84be357..1a15bd2f73 100644 --- a/api/server/routes/share.js +++ b/api/server/routes/share.js @@ -7,6 +7,7 @@ const { ensureLinkPermissions, isFileSnapshotEnabled, isFileSnapshotKillSwitchActive, + buildSharedLinkStartupPayload, deleteSharedLinkWithCleanup, updateSharedLinkPermissionsExpiration, isActiveExpirationDate, @@ -14,8 +15,10 @@ const { } = require('@librechat/api'); const { logger, + getTenantId, runAsSystem, tenantStorage, + SYSTEM_TENANT_ID, createTempChatExpirationDate, } = require('@librechat/data-schemas'); const { FileSources, PermissionTypes, Permissions } = require('librechat-data-provider'); @@ -38,6 +41,7 @@ const optionalShareFileAuth = require('~/server/middleware/optionalShareFileAuth const optionalJwtAuth = require('~/server/middleware/optionalJwtAuth'); const requireJwtAuth = require('~/server/middleware/requireJwtAuth'); const configMiddleware = require('~/server/middleware/config/app'); +const { getAppConfig } = require('~/server/services/Config/app'); const router = express.Router(); const checkSharedLinksAccess = generateCheckAccess({ @@ -76,6 +80,14 @@ const runWithTenant = (tenantId, fn) => * 'failed' on the next poll so the client poller terminates. */ const PREVIEW_LAZY_SWEEP_CUTOFF_MS = 2 * 60 * 1000; +const getShareStartupPayload = async () => { + const tenantId = getTenantId(); + const appConfig = await getAppConfig( + tenantId && tenantId !== SYSTEM_TENANT_ID ? { tenantId } : { baseOnly: true }, + ); + return buildSharedLinkStartupPayload(appConfig); +}; + /** * MIME types that are safe to render inline. Everything else (text/html, SVG, * and other active content) is served as an `attachment` so a public viewer @@ -212,6 +224,17 @@ const streamSharedFile = async (req, res, file, requestedDisposition) => { }; if (allowSharedLinks) { + router.get('/:shareId/config', optionalJwtAuth, canAccessSharedLink, async (_req, res) => { + try { + const payload = await getShareStartupPayload(); + res.set('Cache-Control', 'private, no-store'); + res.status(200).json(payload); + } catch (error) { + logger.error('Error getting shared startup config:', error); + res.status(500).json({ message: 'Error getting shared startup config' }); + } + }); + router.get( '/:shareId', optionalJwtAuth, diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index fd9832d6be..e311b0fd0b 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -1466,20 +1466,25 @@ async function loadToolsForExecution({ AgentConstants.PROGRAMMATIC_TOOL_CALLING, ].filter((name) => toolNames.includes(name)); const isPTCRequested = ptcToolNames.length > 0; + const isBashToolRequested = toolNames.includes(AgentConstants.BASH_TOOL); + const isLegacyExecuteCodeRequested = toolNames.includes(Tools.execute_code); + const isCodeExecutionToolRequested = isBashToolRequested || isLegacyExecuteCodeRequested; let enabledCapabilities; - if (actionsEnabled === undefined || isPTCRequested) { + if (actionsEnabled === undefined || isPTCRequested || isCodeExecutionToolRequested) { enabledCapabilities = await resolveAgentCapabilities(req, appConfig, agent?.id); } if (actionsEnabled === undefined) { actionsEnabled = enabledCapabilities.has(AgentCapabilities.actions); } + const codeExecutionEnabled = + enabledCapabilities?.has(AgentCapabilities.execute_code) === true && + agent?.tools?.includes(Tools.execute_code) === true; const isPTC = isPTCRequested && enabledCapabilities.has(AgentCapabilities.programmatic_tools) && - enabledCapabilities.has(AgentCapabilities.execute_code) && - agent?.tools?.includes(Tools.execute_code) === true; + codeExecutionEnabled; logger.debug( `[loadToolsForExecution] isToolSearch: ${isToolSearch}, toolRegistry: ${toolRegistry?.size ?? 'undefined'}`, @@ -1513,7 +1518,16 @@ async function loadToolsForExecution({ } } - const isBashTool = toolNames.includes(AgentConstants.BASH_TOOL); + const isBashTool = + isBashToolRequested && + codeExecutionEnabled && + toolRegistry?.has(AgentConstants.BASH_TOOL) === true; + if (isBashToolRequested && !isBashTool) { + logger.warn( + `[loadToolsForExecution] Skipping unregistered or unauthorized ${AgentConstants.BASH_TOOL}. ` + + `User: ${req.user.id} | Agent: ${agent?.id ?? 'unknown'}`, + ); + } if (isBashTool) { try { const bashTool = createBashExecutionTool({ @@ -1550,9 +1564,22 @@ async function loadToolsForExecution({ } const requestedNonSpecialToolNames = toolNames.filter((name) => !specialToolNames.has(name)); + const allowedNonSpecialToolNames = requestedNonSpecialToolNames.filter((name) => { + if (name !== Tools.execute_code) { + return true; + } + const allowed = codeExecutionEnabled && toolRegistry?.has(Tools.execute_code) === true; + if (!allowed) { + logger.warn( + `[loadToolsForExecution] Skipping unregistered or unauthorized ${Tools.execute_code}. ` + + `User: ${req.user.id} | Agent: ${agent?.id ?? 'unknown'}`, + ); + } + return allowed; + }); const allToolNamesToLoad = isPTC - ? [...new Set([...requestedNonSpecialToolNames, ...ptcOrchestratedToolNames])] - : requestedNonSpecialToolNames; + ? [...new Set([...allowedNonSpecialToolNames, ...ptcOrchestratedToolNames])] + : allowedNonSpecialToolNames; const actionToolNames = []; const regularToolNames = []; diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index 0d69fb92ef..9f496c5f3a 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -1,3 +1,4 @@ +const { Constants: AgentConstants } = require('@librechat/agents'); const { Tools, Constants, @@ -964,6 +965,29 @@ describe('ToolService - Action Capability Gating', () => { const actionToolName = `get_weather${actionDelimiter}api_example_com`; const regularTool = Tools.web_search; + it('does not load code execution tools that were not registered for the agent', async () => { + const capabilities = [ + AgentCapabilities.tools, + AgentCapabilities.web_search, + AgentCapabilities.execute_code, + ]; + const req = createMockReq(capabilities); + const toolRegistry = new Map([[Tools.web_search, { name: Tools.web_search }]]); + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + + const result = await loadToolsForExecution({ + req, + res: {}, + agent: { id: 'agent_without_code', tools: [Tools.web_search] }, + toolNames: [AgentConstants.BASH_TOOL, Tools.execute_code], + toolRegistry, + actionsEnabled: false, + }); + + expect(result.loadedTools.map((tool) => tool.name)).toEqual([]); + expect(mockLoadToolsUtil).not.toHaveBeenCalled(); + }); + it('loads bash PTC under the legacy programmatic tool name when code capabilities are enabled', async () => { const capabilities = [ AgentCapabilities.tools, diff --git a/client/src/components/Artifacts/ArtifactPreview.tsx b/client/src/components/Artifacts/ArtifactPreview.tsx index 8257f76887..c49c22772e 100644 --- a/client/src/components/Artifacts/ArtifactPreview.tsx +++ b/client/src/components/Artifacts/ArtifactPreview.tsx @@ -4,7 +4,7 @@ import type { SandpackProviderProps, SandpackPreviewRef, } from '@codesandbox/sandpack-react/unstyled'; -import type { TStartupConfig } from 'librechat-data-provider'; +import type { SandpackStartupConfig } from '~/utils/artifacts'; import type { ArtifactFiles } from '~/common'; import { sharedFiles, buildSandpackOptions } from '~/utils/artifacts'; @@ -23,7 +23,7 @@ export const ArtifactPreview = memo(function ({ sharedProps: Partial; previewRef: MutableRefObject; currentCode?: string; - startupConfig?: TStartupConfig; + startupConfig?: SandpackStartupConfig; }) { const artifactFiles = useMemo(() => { if (Object.keys(files).length === 0) { diff --git a/client/src/components/Artifacts/ArtifactTabs.tsx b/client/src/components/Artifacts/ArtifactTabs.tsx index 32332215f0..3ebc98a366 100644 --- a/client/src/components/Artifacts/ArtifactTabs.tsx +++ b/client/src/components/Artifacts/ArtifactTabs.tsx @@ -3,11 +3,12 @@ import * as Tabs from '@radix-ui/react-tabs'; import type { SandpackPreviewRef } from '@codesandbox/sandpack-react/unstyled'; import type { editor } from 'monaco-editor'; import type { Artifact } from '~/common'; -import { useCodeState } from '~/Providers/EditorContext'; +import { useGetSharedStartupConfig, useGetStartupConfig } from '~/data-provider'; import useArtifactProps from '~/hooks/Artifacts/useArtifactProps'; import { ArtifactCodeEditor } from './ArtifactCodeEditor'; -import { useGetStartupConfig } from '~/data-provider'; +import { useCodeState } from '~/Providers/EditorContext'; import { ArtifactPreview } from './ArtifactPreview'; +import { useShareContext } from '~/Providers'; export default function ArtifactTabs({ artifact, @@ -19,7 +20,14 @@ export default function ArtifactTabs({ isSharedConvo?: boolean; }) { const { currentCode, setCurrentCode } = useCodeState(); - const { data: startupConfig } = useGetStartupConfig(); + const { shareId } = useShareContext(); + const shouldUseSharedConfig = + isSharedConvo === true && typeof shareId === 'string' && shareId.length > 0; + const { data: startupConfig } = useGetStartupConfig({ enabled: !shouldUseSharedConfig }); + const { data: sharedStartupConfig } = useGetSharedStartupConfig(shareId, { + enabled: shouldUseSharedConfig, + }); + const resolvedStartupConfig = shouldUseSharedConfig ? sharedStartupConfig : startupConfig; const monacoRef = useRef(null); const lastIdRef = useRef(null); @@ -55,7 +63,7 @@ export default function ArtifactTabs({ previewRef={previewRef} sharedProps={sharedProps} currentCode={currentCode} - startupConfig={startupConfig} + startupConfig={resolvedStartupConfig} /> diff --git a/client/src/components/Chat/Footer.tsx b/client/src/components/Chat/Footer.tsx index b841a98819..90d78737b9 100644 --- a/client/src/components/Chat/Footer.tsx +++ b/client/src/components/Chat/Footer.tsx @@ -8,7 +8,11 @@ import { useLocalize } from '~/hooks'; type FooterProps = { className?: string; - startupConfig?: TStartupConfig | null; + startupConfig?: FooterStartupConfig | null; +}; + +type FooterStartupConfig = Pick, 'analyticsGtmId' | 'customFooter'> & { + interface?: Pick, 'privacyPolicy' | 'termsOfService'>; }; function Footer({ className, startupConfig }: FooterProps) { diff --git a/client/src/components/Share/ShareView.tsx b/client/src/components/Share/ShareView.tsx index 68c7a9cd75..20a354742a 100644 --- a/client/src/components/Share/ShareView.tsx +++ b/client/src/components/Share/ShareView.tsx @@ -18,9 +18,9 @@ import { } from '@librechat/client'; import { ThemeSelector, LangSelector } from '~/components/Nav/SettingsTabs/General/Selectors'; import { ShareMessagesProvider } from './ShareMessagesProvider'; +import { useGetSharedStartupConfig } from '~/data-provider'; import { ShareArtifactsContainer } from './ShareArtifacts'; import { useLocalize, useDocumentTitle } from '~/hooks'; -import { useGetStartupConfig } from '~/data-provider'; import { ShareContext } from '~/Providers'; import MessagesView from './MessagesView'; import Footer from '../Chat/Footer'; @@ -29,9 +29,9 @@ import store from '~/store'; function SharedView() { const localize = useLocalize(); - const { data: config } = useGetStartupConfig(undefined, { context: 'share' }); const { theme, setTheme } = useContext(ThemeContext); const { shareId } = useParams(); + const { data: config } = useGetSharedStartupConfig(shareId); const { data, isLoading } = useGetSharedMessages(shareId ?? ''); const dataTree = data && buildTree({ messages: data.messages }); const messagesTree = dataTree?.length === 0 ? null : (dataTree ?? null); diff --git a/client/src/data-provider/Endpoints/queries.ts b/client/src/data-provider/Endpoints/queries.ts index 1ea8b3c0aa..389a9c44cf 100644 --- a/client/src/data-provider/Endpoints/queries.ts +++ b/client/src/data-provider/Endpoints/queries.ts @@ -88,6 +88,9 @@ export const useContextProjectionQuery = ( export const startupConfigKey = (isAuthenticated: boolean, context?: t.StartupConfigContext) => [QueryKeys.startupConfig, isAuthenticated, context ?? 'default'] as const; +export const sharedStartupConfigKey = (shareId?: string) => + [QueryKeys.sharedStartupConfig, shareId ?? ''] as const; + export const useGetStartupConfig = ( config?: UseQueryOptions, options?: { context?: t.StartupConfigContext }, @@ -107,3 +110,26 @@ export const useGetStartupConfig = ( }, ); }; + +export const useGetSharedStartupConfig = ( + shareId?: string, + config?: UseQueryOptions, +): QueryObserverResult => { + const queriesEnabled = useRecoilValue(store.queriesEnabled); + return useQuery( + sharedStartupConfigKey(shareId), + () => dataService.getSharedStartupConfig(shareId ?? ''), + { + staleTime: Infinity, + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + ...config, + enabled: + (config?.enabled ?? true) === true && + queriesEnabled && + typeof shareId === 'string' && + shareId.length > 0, + }, + ); +}; diff --git a/client/src/utils/artifacts.ts b/client/src/utils/artifacts.ts index 264c8f9e30..b5ff9552e7 100644 --- a/client/src/utils/artifacts.ts +++ b/client/src/utils/artifacts.ts @@ -178,9 +178,14 @@ export const sharedOptions: SandpackProviderProps['options'] = { externalResources: [TAILWIND_CDN], }; +export type SandpackStartupConfig = Pick< + Partial, + 'bundlerURL' | 'staticBundlerURL' +>; + export function buildSandpackOptions( template: SandpackProviderProps['template'], - startupConfig?: TStartupConfig, + startupConfig?: SandpackStartupConfig, ): SandpackProviderProps['options'] { if (!startupConfig) { return sharedOptions; diff --git a/deploy-compose.yml b/deploy-compose.yml index 77d6161a7d..a024d97259 100644 --- a/deploy-compose.yml +++ b/deploy-compose.yml @@ -48,7 +48,9 @@ services: restart: always environment: - PORT=3000 - - SESSION_SECRET=${ADMIN_PANEL_SESSION_SECRET:-${CREDS_KEY}} + # Plain expansion (not :?) so `docker compose down`/`pull` still run when this is unset. + # The panel itself refuses to start without it; set ADMIN_PANEL_SESSION_SECRET in .env. + - SESSION_SECRET=${ADMIN_PANEL_SESSION_SECRET} - API_SERVER_URL=http://api:3080 - VITE_API_BASE_URL=${DOMAIN_CLIENT:-http://localhost} - SESSION_COOKIE_SECURE=${ADMIN_PANEL_SESSION_COOKIE_SECURE:-false} diff --git a/docker-compose.yml b/docker-compose.yml index 1476231df6..0fada4bb69 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -45,7 +45,9 @@ services: restart: always environment: - PORT=3000 - - SESSION_SECRET=${ADMIN_PANEL_SESSION_SECRET:-${CREDS_KEY}} + # Plain expansion (not :?) so `docker compose down`/`pull` still run when this is unset. + # The panel itself refuses to start without it; set ADMIN_PANEL_SESSION_SECRET in .env. + - SESSION_SECRET=${ADMIN_PANEL_SESSION_SECRET} - API_SERVER_URL=http://api:${PORT:-3080} - VITE_API_BASE_URL=${DOMAIN_CLIENT:-http://localhost:3080} - SESSION_COOKIE_SECURE=${ADMIN_PANEL_SESSION_COOKIE_SECURE:-false} diff --git a/packages/api/src/admin/auditLog.spec.ts b/packages/api/src/admin/auditLog.spec.ts index cdb221bcc5..5c5134bf6a 100644 --- a/packages/api/src/admin/auditLog.spec.ts +++ b/packages/api/src/admin/auditLog.spec.ts @@ -1,6 +1,6 @@ import { Types } from 'mongoose'; import { EventEmitter } from 'events'; -import { MAX_AUDIT_EXPORT_ROWS } from '@librechat/data-schemas'; +import { MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_VERIFY_ROWS } from '@librechat/data-schemas'; import type { AdminAuditLogEntry, AuditChainVerification, @@ -64,11 +64,20 @@ function createReqRes( user: 'user' in overrides ? overrides.user : { _id: new Types.ObjectId(), role: 'admin' }, } as unknown as ServerRequest; + const emitter = new EventEmitter(); + const reqEmitter = new EventEmitter(); const json = jest.fn(); const status = jest.fn().mockReturnValue({ json }); - const res = { status, json } as unknown as Response; + const res = Object.assign(emitter, { + status, + json, + removeListener: emitter.removeListener.bind(emitter), + }) as unknown as Response; + const eventReq = Object.assign(reqEmitter, req, { + removeListener: reqEmitter.removeListener.bind(reqEmitter), + }) as unknown as ServerRequest; - return { req, res, status, json }; + return { req: eventReq, res, status, json, emitter, reqEmitter }; } interface CsvCaptureContext { @@ -351,10 +360,29 @@ describe('createAdminAuditLogHandlers', () => { }); await handlers.verifyAuditLog(req, res); expect((deps.verifyAuditChain as jest.Mock).mock.calls[0][0]).toBe('real-tenant'); + expect((deps.verifyAuditChain as jest.Mock).mock.calls[0][1].maxRows).toBe( + MAX_AUDIT_VERIFY_ROWS, + ); expect(status).toHaveBeenCalledWith(200); expect(json).toHaveBeenCalledWith(verification); }); + it('threads cancellation into verification and skips the response after client close', async () => { + const ctx = createReqRes(); + const deps = createDeps({ + verifyAuditChain: jest.fn(async (_tenantId, options) => { + ctx.emitter.emit('close'); + options?.isCancelled?.(); + return mockVerification({ ok: false, reason: 'verification cancelled' }); + }), + }); + const handlers = createAdminAuditLogHandlers(deps); + await handlers.verifyAuditLog(ctx.req, ctx.res); + const optionsArg = (deps.verifyAuditChain as jest.Mock).mock.calls[0][1]; + expect(optionsArg.isCancelled()).toBe(true); + expect(ctx.status).not.toHaveBeenCalled(); + }); + it('returns 500 when verification throws', async () => { const deps = createDeps({ verifyAuditChain: jest.fn().mockRejectedValue(new Error('boom')), diff --git a/packages/api/src/admin/auditLog.ts b/packages/api/src/admin/auditLog.ts index 6b21e02f78..570e70ddf0 100644 --- a/packages/api/src/admin/auditLog.ts +++ b/packages/api/src/admin/auditLog.ts @@ -7,6 +7,7 @@ import { AUDIT_ACTOR_TYPES, MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_LOG_LIMIT, + MAX_AUDIT_VERIFY_ROWS, } from '@librechat/data-schemas'; import type { AdminAuditLogEntry, @@ -59,7 +60,10 @@ export interface AdminAuditLogDeps { onEntry: (entry: AdminAuditLogEntry) => void | Promise, options?: { isCancelled?: () => boolean; maxRows?: number }, ) => Promise<{ count: number; truncated: boolean }>; - verifyAuditChain: (tenantId: string | undefined) => Promise; + verifyAuditChain: ( + tenantId: string | undefined, + options?: { isCancelled?: () => boolean; maxRows?: number }, + ) => Promise; } interface CallerContext { @@ -328,8 +332,25 @@ export function createAdminAuditLogHandlers(deps: AdminAuditLogDeps): { const caller = resolveCaller(req); if (!caller) return res.status(401).json({ error: 'Authentication required' }); - const result = await verifyAuditChain(caller.tenantId); - return res.status(200).json(result); + let clientAborted = false; + const markAborted = () => { + clientAborted = true; + }; + res.once('close', markAborted); + req.once('aborted', markAborted); + + try { + const result = await verifyAuditChain(caller.tenantId, { + isCancelled: () => clientAborted, + maxRows: MAX_AUDIT_VERIFY_ROWS, + }); + + if (clientAborted) return res; + return res.status(200).json(result); + } finally { + res.removeListener('close', markAborted); + req.removeListener('aborted', markAborted); + } } catch (err) { logger.error('[adminAuditLog] verify error:', err); return res.status(500).json({ error: 'Failed to verify audit log integrity' }); diff --git a/packages/api/src/admin/config.handler.spec.ts b/packages/api/src/admin/config.handler.spec.ts index 23d595de89..b44ee085ba 100644 --- a/packages/api/src/admin/config.handler.spec.ts +++ b/packages/api/src/admin/config.handler.spec.ts @@ -994,7 +994,45 @@ describe('createAdminConfigHandlers', () => { await handlers.upsertConfigOverrides(req, res); expect(res.statusCode).toBe(201); - expect(deps.upsertConfig).toHaveBeenCalled(); + expect(deps.upsertConfig).toHaveBeenCalledWith( + 'role', + 'admin', + expect.anything(), + {}, + 10, + undefined, + { expectEmpty: true, preservePriority: true }, + ); + }); + + it('requests atomic priority preservation for ASSIGN_CONFIGS-only empty-overrides upsert', async () => { + const findConfigByPrincipal = jest + .fn() + .mockResolvedValue({ _id: 'c1', priority: 7, overrides: {} }); + const { handlers, deps } = createHandlers({ + hasConfigCapability: jest.fn().mockResolvedValue(false), + hasCapability: jest.fn().mockResolvedValue(true), + findConfigByPrincipal, + }); + const req = mockReq({ + params: { principalType: 'role', principalId: 'admin' }, + body: { overrides: {}, priority: 999 }, + }); + const res = mockRes(); + + await handlers.upsertConfigOverrides(req, res); + + expect(res.statusCode).toBe(201); + expect(deps.upsertConfig).toHaveBeenCalledWith( + 'role', + 'admin', + expect.anything(), + {}, + 10, + undefined, + { expectEmpty: true, preservePriority: true }, + ); + expect(findConfigByPrincipal).not.toHaveBeenCalled(); }); it('rejects non-empty overrides for ASSIGN_CONFIGS-only caller', async () => { @@ -1113,7 +1151,15 @@ describe('createAdminConfigHandlers', () => { await handlers.upsertConfigOverrides(req, res); expect(res.statusCode).toBe(201); - expect(deps.upsertConfig).toHaveBeenCalled(); + expect(deps.upsertConfig).toHaveBeenCalledWith( + 'role', + 'admin', + expect.anything(), + {}, + 10, + undefined, + { expectEmpty: true, preservePriority: true }, + ); }); it('rejects upsert when parameterized grant targets a different principalType', async () => { @@ -1245,7 +1291,7 @@ describe('createAdminConfigHandlers', () => { expect.anything(), expect.anything(), undefined, - { expectEmpty: true }, + { expectEmpty: true, preservePriority: true }, ); }); diff --git a/packages/api/src/admin/config.ts b/packages/api/src/admin/config.ts index f19c28940b..30d3a4d79c 100644 --- a/packages/api/src/admin/config.ts +++ b/packages/api/src/admin/config.ts @@ -79,7 +79,7 @@ export interface AdminConfigDeps { overrides: Partial, priority: number, session?: ClientSession, - options?: { expectEmpty?: boolean }, + options?: { expectEmpty?: boolean; preservePriority?: boolean }, ) => Promise; patchConfigFields: ( principalType: PrincipalType, @@ -393,14 +393,25 @@ export function createAdminConfigHandlers(deps: AdminConfigDeps): { return res.status(403).json({ error: 'Insufficient permissions' }); } + if (priority != null && !hasBroadManage) { + logger.warn( + `[adminConfig] Ignoring caller-supplied priority on assign-only scope lifecycle upsert to ${principalType}/${principalId}: only broad manage:configs may modify document priority`, + ); + } + + const requestedPriority = hasBroadManage ? (priority ?? DEFAULT_PRIORITY) : DEFAULT_PRIORITY; + const upsertOptions = hasBroadManage + ? { expectEmpty: false } + : { expectEmpty: true, preservePriority: true }; + const config = await upsertConfig( principalType, principalId, principalModel(principalType), filteredOverrides, - priority ?? DEFAULT_PRIORITY, + requestedPriority, undefined, - { expectEmpty: !hasBroadManage }, + upsertOptions, ); if (!config && !hasBroadManage) { return res.status(403).json({ error: 'Insufficient permissions' }); diff --git a/packages/api/src/endpoints/projection.spec.ts b/packages/api/src/endpoints/projection.spec.ts new file mode 100644 index 0000000000..1fb7f91a4c --- /dev/null +++ b/packages/api/src/endpoints/projection.spec.ts @@ -0,0 +1,206 @@ +import { resolveContextProjection } from './projection'; +import { QUOTE_MAX_COUNT } from '~/utils/quotes'; + +jest.mock('@librechat/agents', () => ({ + Providers: { OPENAI: 'openai' }, + createTokenCounter: jest.fn(async () => jest.fn(() => 1)), + projectAgentContextUsage: jest.fn(() => ({ tokenCount: 1, maxContextTokens: 1000 })), +})); + +const GRAPH_SELECT = 'messageId parentMessageId metadata.summaryUsedTokens'; +const BODY_SELECT = 'messageId parentMessageId tokenCount isCreatedByUser text quotes'; + +function textStats(messageId: string, textBytes = 5) { + return { + messageId, + textBytes, + quoteCount: 0, + quoteBytes: 0, + quoteLineCount: 0, + nonStringQuoteCount: 0, + }; +} + +describe('resolveContextProjection', () => { + const baseParams = { + conversationId: 'conversation-1', + messageId: 'message-1', + endpoint: 'openai', + maxContextTokens: 1000, + model: 'gpt-4o', + }; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('returns null before tokenization when the conversation is too large', async () => { + const { createTokenCounter } = jest.requireMock('@librechat/agents'); + const messages = Array.from({ length: 513 }, (_, index) => ({ + messageId: `message-${index}`, + parentMessageId: index === 0 ? null : `message-${index - 1}`, + isCreatedByUser: true, + text: 'hello', + })); + const getMessages = jest.fn(async () => messages); + const getMessageTextStats = jest.fn(); + + const result = await resolveContextProjection( + { userId: 'user-1', getMessages, getMessageTextStats }, + { ...baseParams, messageId: 'message-512' }, + ); + + expect(result).toBeNull(); + expect(getMessages).toHaveBeenCalledTimes(1); + expect(getMessages).toHaveBeenCalledWith( + { conversationId: 'conversation-1', user: 'user-1' }, + GRAPH_SELECT, + { limit: 513, sort: false }, + ); + expect(getMessageTextStats).not.toHaveBeenCalled(); + expect(createTokenCounter).not.toHaveBeenCalled(); + }); + + it('returns null before tokenization when the branch is too long', async () => { + const { createTokenCounter } = jest.requireMock('@librechat/agents'); + const messages = Array.from({ length: 257 }, (_, index) => ({ + messageId: `message-${index}`, + parentMessageId: index === 0 ? null : `message-${index - 1}`, + isCreatedByUser: true, + text: 'hello', + })); + const getMessages = jest.fn(async () => messages); + const getMessageTextStats = jest.fn(); + + const result = await resolveContextProjection( + { userId: 'user-1', getMessages, getMessageTextStats }, + { ...baseParams, messageId: 'message-256' }, + ); + + expect(result).toBeNull(); + expect(getMessages).toHaveBeenCalledTimes(1); + expect(getMessageTextStats).not.toHaveBeenCalled(); + expect(createTokenCounter).not.toHaveBeenCalled(); + }); + + it('returns null before loading bodies when the branch text is too large', async () => { + const { createTokenCounter } = jest.requireMock('@librechat/agents'); + const getMessages = jest.fn(async () => [ + { + messageId: 'message-1', + parentMessageId: null, + }, + ]); + const getMessageTextStats = jest.fn(async () => [textStats('message-1', 512 * 1024 + 1)]); + const result = await resolveContextProjection( + { + userId: 'user-1', + getMessages, + getMessageTextStats, + }, + baseParams, + ); + + expect(result).toBeNull(); + expect(getMessages).toHaveBeenCalledTimes(1); + expect(getMessageTextStats).toHaveBeenCalledWith( + { + conversationId: 'conversation-1', + user: 'user-1', + messageId: { $in: ['message-1'] }, + }, + { limit: 1 }, + ); + expect(createTokenCounter).not.toHaveBeenCalled(); + }); + + it('loads only branch message bodies after resolving the graph', async () => { + const graph = [ + { messageId: 'message-1', parentMessageId: null }, + { messageId: 'message-2', parentMessageId: 'message-1' }, + { messageId: 'off-branch', parentMessageId: null }, + ]; + const bodies = [ + { + messageId: 'message-1', + parentMessageId: null, + isCreatedByUser: true, + text: 'first', + tokenCount: 5, + }, + { + messageId: 'message-2', + parentMessageId: 'message-1', + isCreatedByUser: false, + text: 'second', + tokenCount: 6, + }, + ]; + const getMessages = jest.fn(async (_filter: object, select?: string) => + select === GRAPH_SELECT ? graph : bodies, + ); + const getMessageTextStats = jest.fn(async () => [ + textStats('message-1', 5), + textStats('message-2', 6), + ]); + + const result = await resolveContextProjection( + { userId: 'user-1', getMessages, getMessageTextStats }, + { ...baseParams, messageId: 'message-2' }, + ); + + expect(result).toEqual({ tokenCount: 1, maxContextTokens: 1000 }); + expect(getMessages).toHaveBeenNthCalledWith( + 1, + { conversationId: 'conversation-1', user: 'user-1' }, + GRAPH_SELECT, + { limit: 513, sort: false }, + ); + expect(getMessageTextStats).toHaveBeenCalledWith( + { + conversationId: 'conversation-1', + user: 'user-1', + messageId: { $in: ['message-1', 'message-2'] }, + }, + { limit: 2 }, + ); + expect(getMessages).toHaveBeenNthCalledWith( + 2, + { + conversationId: 'conversation-1', + user: 'user-1', + messageId: { $in: ['message-1', 'message-2'] }, + }, + BODY_SELECT, + { limit: 2, sort: false }, + ); + }); + + it('returns null before loading bodies when a branch message has too many quotes', async () => { + const { createTokenCounter } = jest.requireMock('@librechat/agents'); + const getMessages = jest.fn(async () => [ + { + messageId: 'message-1', + parentMessageId: null, + }, + ]); + const getMessageTextStats = jest.fn(async () => [ + { + ...textStats('message-1'), + quoteCount: QUOTE_MAX_COUNT + 1, + quoteBytes: 10, + quoteLineCount: QUOTE_MAX_COUNT + 1, + }, + ]); + + const result = await resolveContextProjection( + { userId: 'user-1', getMessages, getMessageTextStats }, + baseParams, + ); + + expect(result).toBeNull(); + expect(getMessages).toHaveBeenCalledTimes(1); + expect(getMessageTextStats).toHaveBeenCalledTimes(1); + expect(createTokenCounter).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/endpoints/projection.ts b/packages/api/src/endpoints/projection.ts index 365cccc4e8..8622e4b981 100644 --- a/packages/api/src/endpoints/projection.ts +++ b/packages/api/src/endpoints/projection.ts @@ -2,7 +2,13 @@ import { HumanMessage, AIMessage } from '@langchain/core/messages'; import { Providers, createTokenCounter, projectAgentContextUsage } from '@librechat/agents'; import type { TContextProjectionRequest, TContextUsageEvent } from 'librechat-data-provider'; import type { BaseMessage } from '@langchain/core/messages'; -import { mergeQuotedText } from '~/utils/quotes'; +import { QUOTE_MAX_COUNT, mergeQuotedText } from '~/utils/quotes'; + +const MAX_PROJECTION_MESSAGES = 512; +const MAX_PROJECTION_BRANCH_MESSAGES = 256; +const MAX_PROJECTION_BRANCH_TEXT_BYTES = 512 * 1024; +const PROJECTION_GRAPH_SELECT = 'messageId parentMessageId metadata.summaryUsedTokens'; +const PROJECTION_BODY_SELECT = 'messageId parentMessageId tokenCount isCreatedByUser text quotes'; interface ProjectionMessage { messageId: string; @@ -18,13 +24,42 @@ interface ProjectionMessage { metadata?: { summaryUsedTokens?: number }; } +interface ProjectionMessageFilter { + conversationId: string; + user?: string; + messageId?: string | { $in: string[] }; +} + +interface ProjectionMessageQueryOptions { + limit?: number; + sort?: false; +} + +interface ProjectionMessageTextStats { + messageId: string; + textBytes: number; + quoteCount: number; + quoteBytes: number; + quoteLineCount: number; + nonStringQuoteCount: number; +} + +interface ProjectionMessageTextStatsOptions { + limit?: number; +} + export interface ContextProjectionDeps { /** Authenticated requester — branch lookups are scoped to this user. */ userId?: string; getMessages: ( - filter: { conversationId: string; user?: string }, + filter: ProjectionMessageFilter, select?: string, + options?: ProjectionMessageQueryOptions, ) => Promise; + getMessageTextStats: ( + filter: ProjectionMessageFilter, + options?: ProjectionMessageTextStatsOptions, + ) => Promise; } /** @@ -51,6 +86,83 @@ function resolveBranch(messages: ProjectionMessage[], tailId: string): Projectio return branch.reverse(); } +function hasValidProjectionIds(params: TContextProjectionRequest): boolean { + return typeof params.conversationId === 'string' && typeof params.messageId === 'string'; +} + +function getProjectionText(message: ProjectionMessage): string | null { + const hasQuotes = + message.isCreatedByUser === true && Array.isArray(message.quotes) && message.quotes.length > 0; + if (!hasQuotes) { + return message.text ?? ''; + } + if (message.quotes == null || message.quotes.length > QUOTE_MAX_COUNT) { + return null; + } + for (const quote of message.quotes) { + if (typeof quote !== 'string') { + return null; + } + } + return mergeQuotedText(message.text ?? '', message.quotes); +} + +function hasExceededBranchTextLimit(branch: ProjectionMessage[]): boolean { + let bytes = 0; + for (const message of branch) { + const text = getProjectionText(message); + if (text == null) { + return true; + } + bytes += Buffer.byteLength(text, 'utf8'); + if (bytes > MAX_PROJECTION_BRANCH_TEXT_BYTES) { + return true; + } + } + return false; +} + +function getEstimatedMergedTextBytes(stats: ProjectionMessageTextStats): number | null { + if ( + stats.nonStringQuoteCount > 0 || + stats.quoteCount > QUOTE_MAX_COUNT || + stats.quoteLineCount < stats.quoteCount + ) { + return null; + } + if (stats.quoteCount === 0) { + return stats.textBytes; + } + + const quotePrefixBytes = stats.quoteLineCount * 2; + const quoteLineBreakBytes = stats.quoteLineCount - stats.quoteCount; + const quoteSeparatorBytes = (stats.quoteCount - 1) * 2; + const bodySeparatorBytes = stats.textBytes > 0 ? 2 : 0; + return ( + stats.textBytes + + stats.quoteBytes + + quotePrefixBytes + + quoteLineBreakBytes + + quoteSeparatorBytes + + bodySeparatorBytes + ); +} + +function hasExceededBranchTextStatsLimit(stats: ProjectionMessageTextStats[]): boolean { + let bytes = 0; + for (const messageStats of stats) { + const messageBytes = getEstimatedMergedTextBytes(messageStats); + if (messageBytes == null) { + return true; + } + bytes += messageBytes; + if (bytes > MAX_PROJECTION_BRANCH_TEXT_BYTES) { + return true; + } + } + return false; +} + /** Maps an endpoint/provider string to the agents `Providers` enum. */ function resolveProvider(value?: string): Providers { if (value == null || value === '') { @@ -74,6 +186,43 @@ function resolveProvider(value?: string): Providers { return Providers.OPENAI; } +async function getBranchMessages( + deps: ContextProjectionDeps, + baseFilter: ProjectionMessageFilter, + branch: ProjectionMessage[], +): Promise { + const branchIds = branch.map((message) => message.messageId); + const stats = await deps.getMessageTextStats( + { ...baseFilter, messageId: { $in: branchIds } }, + { limit: branchIds.length }, + ); + if (stats.length !== branchIds.length || hasExceededBranchTextStatsLimit(stats)) { + return null; + } + + const stored = await deps.getMessages( + { ...baseFilter, messageId: { $in: branchIds } }, + PROJECTION_BODY_SELECT, + { limit: branchIds.length, sort: false }, + ); + if (stored.length !== branchIds.length) { + return null; + } + const byId = new Map(); + for (const message of stored) { + byId.set(message.messageId, message); + } + const ordered: ProjectionMessage[] = []; + for (const messageId of branchIds) { + const message = byId.get(messageId); + if (message == null) { + return null; + } + ordered.push(message); + } + return ordered; +} + /** * Server-side context-usage projection: reconstructs the viewed branch and asks * the agents SDK what the next call's context would be, WITHOUT invoking the @@ -90,19 +239,31 @@ export async function resolveContextProjection( deps: ContextProjectionDeps, params: TContextProjectionRequest, ): Promise { + if (!hasValidProjectionIds(params)) { + return null; + } + const maxContextTokens = params.maxContextTokens; if (maxContextTokens == null || maxContextTokens <= 0) { return null; } - const stored = await deps.getMessages( - { conversationId: params.conversationId, user: deps.userId }, - 'messageId parentMessageId tokenCount isCreatedByUser text quotes metadata', - ); + const baseFilter = { conversationId: params.conversationId, user: deps.userId }; + const stored = await deps.getMessages(baseFilter, PROJECTION_GRAPH_SELECT, { + limit: MAX_PROJECTION_MESSAGES + 1, + sort: false, + }); + if (stored.length > MAX_PROJECTION_MESSAGES) { + return null; + } + const branch = resolveBranch(stored, params.messageId); if (branch.length === 0) { return null; } + if (branch.length > MAX_PROJECTION_BRANCH_MESSAGES) { + return null; + } /** A summarized/compacted branch's next call sends the saved summary + the * post-summary tail, NOT this raw parent chain — projecting from the full @@ -114,23 +275,29 @@ export async function resolveContextProjection( return null; } + const bodyBranch = await getBranchMessages(deps, baseFilter, branch); + if (bodyBranch == null || hasExceededBranchTextLimit(bodyBranch)) { + return null; + } + const model = params.model; const encoding = (model ?? '').toLowerCase().includes('claude') ? 'claude' : 'o200k_base'; const tokenCounter = await createTokenCounter(encoding); const messages: BaseMessage[] = []; const indexTokenCountMap: Record = {}; - for (let i = 0; i < branch.length; i++) { - const message = branch[i]; + for (let i = 0; i < bodyBranch.length; i++) { + const message = bodyBranch[i]; /** Mirror the live path: prepend quoted excerpts into the user text the model * receives so the gauge counts the same prompt. */ const hasQuotes = message.isCreatedByUser === true && Array.isArray(message.quotes) && message.quotes.length > 0; - const text = hasQuotes - ? mergeQuotedText(message.text ?? '', message.quotes ?? []) - : (message.text ?? ''); + const text = getProjectionText(message); + if (text == null) { + return null; + } const lcMessage = message.isCreatedByUser === true ? new HumanMessage(text) : new AIMessage(text); messages.push(lcMessage); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts index 302013824d..19c972f0c6 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts @@ -44,7 +44,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); const mockLogger = logger as jest.Mocked; diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts index 12297e7be8..e087ab9f20 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts @@ -45,6 +45,9 @@ jest.mock('~/mcp/mcpConfig', () => ({ OAUTH_HANDLING_TIMEOUT: 10 * 60 * 1000, USER_CONNECTION_IDLE_TIMEOUT: 30 * 60 * 1000, TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, }, })); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index 095c2f36ea..d0aec66fc9 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -10,6 +10,7 @@ import { logger } from '@librechat/data-schemas'; import { MCPConnection } from '~/mcp/connection'; +import { mcpConfig } from '~/mcp/mcpConfig'; jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -29,7 +30,13 @@ jest.mock('~/auth', () => ({ /** Pin the page cap to a small value so the cap path is cheap to exercise. */ jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { TOOLS_LIST_MAX_PAGES: 3, CONNECTION_CHECK_TTL: 0 }, + mcpConfig: { + TOOLS_LIST_MAX_PAGES: 3, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + CONNECTION_CHECK_TTL: 0, + }, })); const mockLogger = logger as jest.Mocked; @@ -51,9 +58,27 @@ function createConnectionWithListTools(listTools: jest.Mock): MCPConnection { return conn; } +function expectListToolsCall( + listTools: jest.Mock, + callNumber: number, + params: { cursor?: string } | undefined, +): void { + expect(listTools).toHaveBeenNthCalledWith( + callNumber, + params, + expect.objectContaining({ + timeout: expect.any(Number), + maxTotalTimeout: expect.any(Number), + }), + ); +} + describe('MCPConnection.fetchTools pagination', () => { beforeEach(() => { jest.clearAllMocks(); + mcpConfig.TOOLS_LIST_MAX_TOOLS = 1000; + mcpConfig.TOOLS_LIST_MAX_BYTES = 5 * 1024 * 1024; + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; }); it('returns the tools from a single page and makes one request when there is no nextCursor', async () => { @@ -64,7 +89,7 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b']); expect(listTools).toHaveBeenCalledTimes(1); - expect(listTools).toHaveBeenNthCalledWith(1, undefined); + expectListToolsCall(listTools, 1, undefined); expect(mockLogger.warn).not.toHaveBeenCalled(); }); @@ -87,9 +112,9 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b', 'c', 'd', 'e']); expect(listTools).toHaveBeenCalledTimes(3); - expect(listTools).toHaveBeenNthCalledWith(1, undefined); - expect(listTools).toHaveBeenNthCalledWith(2, { cursor: 'c1' }); - expect(listTools).toHaveBeenNthCalledWith(3, { cursor: 'c2' }); + expectListToolsCall(listTools, 1, undefined); + expectListToolsCall(listTools, 2, { cursor: 'c1' }); + expectListToolsCall(listTools, 3, { cursor: 'c2' }); expect(mockLogger.warn).not.toHaveBeenCalled(); }); @@ -109,6 +134,96 @@ describe('MCPConnection.fetchTools pagination', () => { expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('pagination limit')); }); + it('stops at the aggregate tool-count budget and warns', async () => { + mcpConfig.TOOLS_LIST_MAX_TOOLS = 3; + const listTools = jest.fn(async (params?: { cursor?: string }) => { + if (params?.cursor == null) { + return { tools: [makeTool('a'), makeTool('b')], nextCursor: 'c1' }; + } + return { tools: [makeTool('c'), makeTool('d')], nextCursor: 'c2' }; + }); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a', 'b', 'c']); + expect(listTools).toHaveBeenCalledTimes(2); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('tool count budget')); + }); + + it('does not request another page when the tool-count budget is exactly full', async () => { + mcpConfig.TOOLS_LIST_MAX_TOOLS = 2; + const listTools = jest.fn(async (params?: { cursor?: string }) => { + if (params?.cursor == null) { + return { tools: [makeTool('a'), makeTool('b')], nextCursor: 'c1' }; + } + return { tools: [makeTool('c')] }; + }); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a', 'b']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('tool count budget')); + }); + + it('stops at the aggregate byte budget and warns', async () => { + mcpConfig.TOOLS_LIST_MAX_BYTES = 170; + const listTools = jest.fn(async () => ({ + tools: [makeTool('a'), makeTool('b')], + nextCursor: 'c1', + })); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('size budget')); + }); + + it('stops at the elapsed-time budget before requesting another page', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 1; + const listTools = jest.fn(async () => ({ tools: [makeTool('a')], nextCursor: 'c1' })); + const conn = createConnectionWithListTools(listTools); + const dateNow = jest + .spyOn(Date, 'now') + .mockReturnValueOnce(1000) + .mockReturnValueOnce(1000) + .mockReturnValueOnce(1001); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('time budget')); + dateNow.mockRestore(); + }); + + it('passes the elapsed-time budget to the SDK request timeout', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 25; + const listTools = jest.fn( + async ( + _params?: { cursor?: string }, + _options?: { timeout: number; maxTotalTimeout: number }, + ) => { + throw new Error('Request timed out'); + }, + ); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools).toEqual([]); + expect(listTools).toHaveBeenCalledTimes(1); + const options = listTools.mock.calls[0][1]!; + expect(options.timeout).toBeGreaterThan(0); + expect(options.timeout).toBeLessThanOrEqual(25); + expect(options.maxTotalTimeout).toBe(options.timeout); + expect(mockLogger.error).toHaveBeenCalledWith(expect.stringContaining('Request timed out')); + }); + it('stops and warns when the server repeats a cursor instead of looping forever', async () => { const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('x')], nextCursor: 'same' }); const conn = createConnectionWithListTools(listTools); @@ -151,7 +266,7 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b']); expect(listTools).toHaveBeenCalledTimes(2); - expect(listTools).toHaveBeenNthCalledWith(2, { cursor: '' }); + expectListToolsCall(listTools, 2, { cursor: '' }); }); it('returns the pages already fetched when a later page fails, without throwing', async () => { diff --git a/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts index c6aec18f67..702a073721 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts @@ -71,7 +71,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); const mockedResolveHostnameSSRF = resolveHostnameSSRF as jest.MockedFunction< diff --git a/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts b/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts index 4ee078795c..981f537a57 100644 --- a/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts +++ b/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts @@ -43,7 +43,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); /** Track all Agents for cleanup */ diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index bd9884ef64..87116b6a6e 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -30,6 +30,7 @@ import { mcpConfig } from './mcpConfig'; type FetchLike = (url: string | URL, init?: RequestInit) => Promise; type ManagedDispatcher = Agent | ProxyAgent; type ParsedIP = { version: 4 | 6; bits: 32 | 128; value: bigint }; +type MCPTool = MCPListToolsResult['tools'][number]; const BIGINT_ZERO = BigInt(0); const BIGINT_ONE = BigInt(1); @@ -37,6 +38,29 @@ const BIGINT_EIGHT = BigInt(8); const BIGINT_SIXTEEN = BigInt(16); const UINT16_MASK = BigInt(0xffff); +function getApproximateToolBytes(tool: MCPTool): number { + try { + return Buffer.byteLength(JSON.stringify(tool), 'utf8'); + } catch { + return Number.POSITIVE_INFINITY; + } +} + +function getToolsListBudgetExceededReason( + toolCount: number, + totalBytes: number, + maxTools: number, + maxBytes: number, +): string | null { + if (toolCount >= maxTools) { + return 'tool count'; + } + if (totalBytes >= maxBytes) { + return 'size'; + } + return null; +} + type MCPProxyConfig = | { type: 'explicit'; @@ -2202,29 +2226,77 @@ export class MCPConnection extends EventEmitter { * server that spans multiple pages (e.g. an aggregating gateway exposing many * tools) is loaded in full instead of being truncated to the first page. * - * Pagination is bounded by {@link mcpConfig.TOOLS_LIST_MAX_PAGES} and a - * repeated-cursor guard. On error, the tools already fetched are returned rather - * than discarded, and the method never throws. + * Pagination is bounded by {@link mcpConfig.TOOLS_LIST_MAX_PAGES}, aggregate + * tool count, approximate serialized size, elapsed time, and a repeated-cursor + * guard. On error, the tools already fetched are returned rather than discarded, + * and the method never throws. */ async fetchTools(): Promise { const maxPages = mcpConfig.TOOLS_LIST_MAX_PAGES; + const maxTools = mcpConfig.TOOLS_LIST_MAX_TOOLS; + const maxBytes = mcpConfig.TOOLS_LIST_MAX_BYTES; + const deadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS; const allTools: MCPListToolsResult['tools'] = []; const seenCursors = new Set(); let cursor: string | undefined; + let totalBytes = 0; for (let page = 1; page <= maxPages; page++) { - const result = await this.listToolsPage(cursor); + const exhaustedBudget = getToolsListBudgetExceededReason( + allTools.length, + totalBytes, + maxTools, + maxBytes, + ); + if (exhaustedBudget != null) { + this.warnToolsListBudgetExceeded(exhaustedBudget, allTools.length); + return allTools; + } + + const remainingMs = deadline - Date.now(); + if (remainingMs <= 0) { + this.warnToolsListBudgetExceeded('time', allTools.length); + return allTools; + } + + const result = await this.listToolsPage(cursor, remainingMs); if (result == null) { /** Request failed mid-pagination: return the pages already fetched instead of discarding them. */ return allTools; } - allTools.push(...result.tools); + for (const tool of result.tools) { + if (allTools.length >= maxTools) { + this.warnToolsListBudgetExceeded('tool count', allTools.length); + return allTools; + } + + const toolBytes = getApproximateToolBytes(tool); + if (totalBytes + toolBytes > maxBytes) { + this.warnToolsListBudgetExceeded('size', allTools.length); + return allTools; + } + + allTools.push(tool); + totalBytes += toolBytes; + } const { nextCursor } = result; if (nextCursor == null) { return allTools; } + + const nextPageBudget = getToolsListBudgetExceededReason( + allTools.length, + totalBytes, + maxTools, + maxBytes, + ); + if (nextPageBudget != null) { + this.warnToolsListBudgetExceeded(nextPageBudget, allTools.length); + return allTools; + } + if (seenCursors.has(nextCursor)) { logger.warn( `${this.getLogPrefix()} MCP server returned a repeated tools/list cursor; stopping pagination after ${page} page(s).`, @@ -2242,10 +2314,22 @@ export class MCPConnection extends EventEmitter { return allTools; } + private warnToolsListBudgetExceeded(reason: string, toolCount: number): void { + logger.warn( + `${this.getLogPrefix()} Stopping tools/list pagination because the ${reason} budget was reached after ${toolCount} tool(s).`, + ); + } + /** Fetches a single `tools/list` page, returning null (and logging) on failure so pagination can stop gracefully. */ - private async listToolsPage(cursor: string | undefined): Promise { + private async listToolsPage( + cursor: string | undefined, + timeoutMs: number, + ): Promise { try { - return await this.client.listTools(cursor != null ? { cursor } : undefined); + return await this.client.listTools(cursor != null ? { cursor } : undefined, { + timeout: timeoutMs, + maxTotalTimeout: timeoutMs, + }); } catch (error) { this.emitError(error, 'Failed to fetch tools'); return null; diff --git a/packages/api/src/mcp/mcpConfig.ts b/packages/api/src/mcp/mcpConfig.ts index ea75220958..68ae3cdfe3 100644 --- a/packages/api/src/mcp/mcpConfig.ts +++ b/packages/api/src/mcp/mcpConfig.ts @@ -26,6 +26,12 @@ export const mcpConfig: { /** Max number of `tools/list` pages to request when an MCP server paginates its tool list. * Bounds the pagination loop so a misbehaving server cannot stall tool discovery. Default: 50 */ TOOLS_LIST_MAX_PAGES: number; + /** Max total tools to retain from paginated `tools/list` responses. Default: 1000 */ + TOOLS_LIST_MAX_TOOLS: number; + /** Max approximate JSON bytes to retain from paginated `tools/list` responses. Default: 5 MiB */ + TOOLS_LIST_MAX_BYTES: number; + /** Max elapsed time (ms) for paginated `tools/list` discovery. Default: 30000 */ + TOOLS_LIST_TIMEOUT_MS: number; /** Idle timeout (ms) after which user connections are disconnected. Default: 15 minutes */ USER_CONNECTION_IDLE_TIMEOUT: number; /** Max connect/disconnect cycles before the circuit breaker trips. Default: 7 */ @@ -52,6 +58,12 @@ export const mcpConfig: { CONNECTION_CHECK_TTL: math(process.env.MCP_CONNECTION_CHECK_TTL ?? 60000), /** Max number of `tools/list` pages to request when an MCP server paginates its tool list. Clamped to >= 1. Default: 50 */ TOOLS_LIST_MAX_PAGES: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_PAGES ?? 50)), + /** Max total tools to retain from paginated `tools/list` responses. Clamped to >= 1. Default: 1000 */ + TOOLS_LIST_MAX_TOOLS: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_TOOLS ?? 1000)), + /** Max approximate JSON bytes to retain from paginated `tools/list` responses. Clamped to >= 1. Default: 5 MiB */ + TOOLS_LIST_MAX_BYTES: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_BYTES ?? 5 * 1024 * 1024)), + /** Max elapsed time (ms) for paginated `tools/list` discovery. Clamped to >= 1. Default: 30000 */ + TOOLS_LIST_TIMEOUT_MS: Math.max(1, math(process.env.MCP_TOOLS_LIST_TIMEOUT_MS ?? 30_000)), /** Idle timeout (ms) after which user connections are disconnected. Default: 15 minutes */ USER_CONNECTION_IDLE_TIMEOUT: math( process.env.MCP_USER_CONNECTION_IDLE_TIMEOUT ?? 15 * 60 * 1000, diff --git a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts index dd73e89d43..cd21502432 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts @@ -56,7 +56,13 @@ jest.mock('~/cluster', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); jest.mock('~/mcp/registry/db/ServerConfigsDB', () => ({ diff --git a/packages/api/src/shared-links/access.test.ts b/packages/api/src/shared-links/access.test.ts index 795ef6c7f1..3f8ea77146 100644 --- a/packages/api/src/shared-links/access.test.ts +++ b/packages/api/src/shared-links/access.test.ts @@ -124,6 +124,20 @@ describe('canAccessSharedLink', () => { expect(res._status).toBe(404); expect(next).not.toHaveBeenCalled(); }); + + test('returns 404 when share is expired but ACL still exists', async () => { + const link = await createTestLink({ expiredAt: new Date('2020-01-01T00:00:00.000Z') }); + await grantPublicViewer(link._id); + process.env.ALLOW_SHARED_LINKS_PUBLIC = 'true'; + + const req = createReq({ params: { shareId: link.shareId } }); + const res = createRes(); + const next = jest.fn(); + await canAccessSharedLink(req, res, next as unknown as NextFunction); + + expect(res._status).toBe(404); + expect(next).not.toHaveBeenCalled(); + }); }); describe('public links', () => { diff --git a/packages/api/src/shared-links/access.ts b/packages/api/src/shared-links/access.ts index d83c5fd24d..3cbb971c3a 100644 --- a/packages/api/src/shared-links/access.ts +++ b/packages/api/src/shared-links/access.ts @@ -1,5 +1,10 @@ import { ResourceType, PermissionBits } from 'librechat-data-provider'; -import { getTenantId, runAsSystem, tenantStorage } from '@librechat/data-schemas'; +import { + getTenantId, + runAsSystem, + tenantStorage, + activeExpirationFilter, +} from '@librechat/data-schemas'; import type { Request, Response, NextFunction } from 'express'; import type { IUser } from '@librechat/data-schemas'; import type { Types, Model } from 'mongoose'; @@ -53,7 +58,10 @@ export function createSharedLinkAccessMiddleware(deps: SharedLinkAccessDeps) { const SharedLink = mg.models.SharedLink as Model; const findShare = async () => - (await SharedLink.findOne({ shareId }).lean()) as RawSharedLink | null; + (await SharedLink.findOne({ + shareId, + ...activeExpirationFilter(), + }).lean()) as RawSharedLink | null; const rawShare = getTenantId() ? await findShare() : await runAsSystem(findShare); if (!rawShare) { diff --git a/packages/api/src/shared-links/config.test.ts b/packages/api/src/shared-links/config.test.ts index 12d4c3d23f..5df601d541 100644 --- a/packages/api/src/shared-links/config.test.ts +++ b/packages/api/src/shared-links/config.test.ts @@ -1,5 +1,5 @@ import type { AppConfig } from '@librechat/data-schemas'; -import { isFileSnapshotEnabled } from './config'; +import { buildSharedLinkStartupPayload, isFileSnapshotEnabled } from './config'; const withSharedLinks = (sharedLinks: unknown): AppConfig => ({ interfaceConfig: { sharedLinks } }) as unknown as AppConfig; @@ -41,3 +41,45 @@ describe('isFileSnapshotEnabled', () => { expect(isFileSnapshotEnabled(withSharedLinks({ snapshotFiles: false }))).toBe(true); }); }); + +describe('buildSharedLinkStartupPayload', () => { + it('builds the share-view startup allowlist', () => { + const payload = buildSharedLinkStartupPayload( + { + interfaceConfig: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + termsOfService: { externalUrl: 'https://example.com/tos' }, + modelSelect: true, + }, + } as AppConfig, + { + ANALYTICS_GTM_ID: 'GTM-XYZ', + APP_TITLE: 'Test Chat', + CUSTOM_FOOTER: 'Shared footer', + SANDPACK_BUNDLER_URL: 'https://bundler.example.com', + SANDPACK_STATIC_BUNDLER_URL: 'https://static-bundler.example.com', + }, + ); + + expect(payload).toEqual({ + appTitle: 'Test Chat', + analyticsGtmId: 'GTM-XYZ', + bundlerURL: 'https://bundler.example.com', + staticBundlerURL: 'https://static-bundler.example.com', + customFooter: 'Shared footer', + interface: { + privacyPolicy: { externalUrl: 'https://example.com/privacy' }, + termsOfService: { externalUrl: 'https://example.com/tos' }, + }, + }); + }); + + it('defaults the app title and omits unrelated interface config', () => { + const payload = buildSharedLinkStartupPayload( + { interfaceConfig: { modelSelect: true } } as AppConfig, + {}, + ); + + expect(payload).toEqual({ appTitle: 'LibreChat' }); + }); +}); diff --git a/packages/api/src/shared-links/config.ts b/packages/api/src/shared-links/config.ts index cd9766b763..1d19c7d22c 100644 --- a/packages/api/src/shared-links/config.ts +++ b/packages/api/src/shared-links/config.ts @@ -1,6 +1,9 @@ +import type { TSharedLinkStartupConfig } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; import { isEnabled } from '~/utils'; +type SharedLinkStartupEnv = NodeJS.ProcessEnv; + /** * Whether shared links should snapshot the files referenced by the shared chat * snapshot. The `SHARED_LINKS_SNAPSHOT_FILES` env var overrides the yaml @@ -31,3 +34,35 @@ export function isFileSnapshotKillSwitchActive(): boolean { const envValue = process.env.SHARED_LINKS_SNAPSHOT_FILES; return envValue !== undefined && !isEnabled(envValue); } + +export function buildSharedLinkStartupPayload( + appConfig?: AppConfig | null, + env: SharedLinkStartupEnv = process.env, +): TSharedLinkStartupConfig { + const payload: TSharedLinkStartupConfig = { + appTitle: env.APP_TITLE || 'LibreChat', + }; + + if (typeof env.ANALYTICS_GTM_ID === 'string') { + payload.analyticsGtmId = env.ANALYTICS_GTM_ID; + } + if (typeof env.SANDPACK_BUNDLER_URL === 'string') { + payload.bundlerURL = env.SANDPACK_BUNDLER_URL; + } + if (typeof env.SANDPACK_STATIC_BUNDLER_URL === 'string') { + payload.staticBundlerURL = env.SANDPACK_STATIC_BUNDLER_URL; + } + if (typeof env.CUSTOM_FOOTER === 'string') { + payload.customFooter = env.CUSTOM_FOOTER; + } + + const { privacyPolicy, termsOfService } = appConfig?.interfaceConfig ?? {}; + if (privacyPolicy || termsOfService) { + payload.interface = { + ...(privacyPolicy ? { privacyPolicy } : {}), + ...(termsOfService ? { termsOfService } : {}), + }; + } + + return payload; +} diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index bf3d3031bd..326e60f644 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -71,6 +71,7 @@ export const messagesBranch = () => `${messagesRoot}/branch`; const shareRoot = `${BASE_URL}/api/share`; export const shareMessages = (shareId: string) => `${shareRoot}/${shareId}`; +export const sharedStartupConfig = (shareId: string) => `${shareMessages(shareId)}/config`; export const getSharedLink = (conversationId: string) => `${shareRoot}/link/${conversationId}`; export const getSharedLinks = ( pageSize: number, diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index eafcc2f97f..2e0439cc0e 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1447,6 +1447,19 @@ export type TStartupConfig = { }; }; +export type TSharedLinkStartupInterface = Pick< + Partial, + 'privacyPolicy' | 'termsOfService' +>; + +export type TSharedLinkStartupConfig = Pick & + Pick< + Partial, + 'analyticsGtmId' | 'bundlerURL' | 'customFooter' | 'staticBundlerURL' + > & { + interface?: TSharedLinkStartupInterface; + }; + export enum OCRStrategy { MISTRAL_OCR = 'mistral_ocr', CUSTOM_OCR = 'custom_ocr', diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 1ffac23dd4..80e5555170 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -64,6 +64,10 @@ export function getSharedMessages(shareId: string): Promise { + return request.get(endpoints.sharedStartupConfig(shareId)); +} + export const listSharedLinks = async ( params: q.SharedLinksListParams, ): Promise => { diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index 54c7b61c25..dde54005a6 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -1,6 +1,7 @@ export enum QueryKeys { messages = 'messages', sharedMessages = 'sharedMessages', + sharedStartupConfig = 'sharedStartupConfig', sharedLinks = 'sharedLinks', allConversations = 'allConversations', archivedConversations = 'archivedConversations', diff --git a/packages/data-schemas/src/index.ts b/packages/data-schemas/src/index.ts index e35c29e2e3..ee393b2de4 100644 --- a/packages/data-schemas/src/index.ts +++ b/packages/data-schemas/src/index.ts @@ -27,6 +27,7 @@ export { AUDIT_SCHEMA_VERSION, MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_LOG_LIMIT, + MAX_AUDIT_VERIFY_ROWS, } from './methods'; export type * from './types'; export type * from './methods'; diff --git a/packages/data-schemas/src/methods/auditLog.spec.ts b/packages/data-schemas/src/methods/auditLog.spec.ts index eb5504b738..aefdf34e81 100644 --- a/packages/data-schemas/src/methods/auditLog.spec.ts +++ b/packages/data-schemas/src/methods/auditLog.spec.ts @@ -351,6 +351,30 @@ describe('auditLog methods', () => { expect(result.checked).toBe(0); }); + it('stops verification when the caller cancels', async () => { + await seed(3); + const result = await methods.verifyAuditChain('tenant-a', { isCancelled: () => true }); + expect(result.ok).toBe(false); + expect(result.checked).toBe(0); + expect(result.reason).toBe('verification cancelled'); + }); + + it('bounds verification when rows exceed the configured cap', async () => { + await seed(3); + const result = await methods.verifyAuditChain('tenant-a', { maxRows: 2 }); + expect(result.ok).toBe(false); + expect(result.checked).toBe(2); + expect(result.brokenAt).toBe(3); + expect(result.reason).toMatch(/row limit exceeded/); + }); + + it('allows an exact-cap verification to complete', async () => { + await seed(2); + const result = await methods.verifyAuditChain('tenant-a', { maxRows: 2 }); + expect(result.ok).toBe(true); + expect(result.checked).toBe(2); + }); + it('detects a tampered field (hash mismatch)', async () => { await seed(3); // mutate a field via the raw driver, bypassing the append-only hooks diff --git a/packages/data-schemas/src/methods/auditLog.ts b/packages/data-schemas/src/methods/auditLog.ts index 6d520df07f..0ab6a7dd51 100644 --- a/packages/data-schemas/src/methods/auditLog.ts +++ b/packages/data-schemas/src/methods/auditLog.ts @@ -27,6 +27,11 @@ export const MAX_AUDIT_LOG_LIMIT = 500; * should be sliced by `from`/`to`. */ export const MAX_AUDIT_EXPORT_ROWS = 100_000; +/** + * Upper bound on rows verified in a single HTTP-triggered integrity check. Full + * offline jobs can pass a larger value explicitly when needed. + */ +export const MAX_AUDIT_VERIFY_ROWS = 100_000; /** Record-format version stamped on every new entry. */ export const AUDIT_SCHEMA_VERSION = 1; const MAX_SEARCH_LENGTH = 200; @@ -516,6 +521,8 @@ export function createAuditLogMethods(mongoose: typeof import('mongoose')): Audi .cursor({ batchSize: 500 }); const trustedCheckpoint = options?.trustedCheckpoint; + const isCancelled = options?.isCancelled; + const maxRows = options?.maxRows; let prevHash = GENESIS_HASH; let expectedSeq: number | null = null; let firstSeq: number | null = null; @@ -524,6 +531,28 @@ export function createAuditLogMethods(mongoose: typeof import('mongoose')): Audi try { for await (const doc of cursor) { + if (isCancelled?.()) { + await cursor.close(); + return { + ok: false, + chainKey, + checked, + reason: 'verification cancelled', + range: firstSeq !== null ? { firstSeq, lastSeq } : undefined, + }; + } + if (maxRows != null && checked >= maxRows) { + await cursor.close(); + return { + ok: false, + chainKey, + checked, + brokenAt: doc.seq, + reason: `verification row limit exceeded (${maxRows})`, + range: + firstSeq !== null ? { firstSeq, lastSeq } : { firstSeq: doc.seq, lastSeq: doc.seq }, + }; + } if (firstSeq === null) { firstSeq = doc.seq; expectedSeq = doc.seq; diff --git a/packages/data-schemas/src/methods/config.spec.ts b/packages/data-schemas/src/methods/config.spec.ts index 35cc3cece7..5111a915e6 100644 --- a/packages/data-schemas/src/methods/config.spec.ts +++ b/packages/data-schemas/src/methods/config.spec.ts @@ -587,6 +587,39 @@ describe('expectEmpty atomic guard', () => { expect(result!.priority).toBe(99); }); + it('upsertConfig with preservePriority inserts with the requested priority', async () => { + const result = await methods.upsertConfig( + PrincipalType.ROLE, + 'admin', + PrincipalModel.ROLE, + {}, + 10, + undefined, + { expectEmpty: true, preservePriority: true }, + ); + + expect(result).toBeTruthy(); + expect(result!.priority).toBe(10); + }); + + it('upsertConfig with preservePriority keeps an empty existing doc priority', async () => { + await methods.upsertConfig(PrincipalType.ROLE, 'admin', PrincipalModel.ROLE, {}, 5); + + const result = await methods.upsertConfig( + PrincipalType.ROLE, + 'admin', + PrincipalModel.ROLE, + {}, + 99, + undefined, + { expectEmpty: true, preservePriority: true }, + ); + + expect(result).toBeTruthy(); + expect(result!.priority).toBe(5); + expect(result!.configVersion).toBe(2); + }); + it('upsertConfig with expectEmpty returns null when existing doc has non-empty overrides', async () => { await methods.upsertConfig( PrincipalType.ROLE, diff --git a/packages/data-schemas/src/methods/config.ts b/packages/data-schemas/src/methods/config.ts index 62f6c10e96..16f94efd9e 100644 --- a/packages/data-schemas/src/methods/config.ts +++ b/packages/data-schemas/src/methods/config.ts @@ -37,7 +37,7 @@ export function createConfigMethods(mongoose: typeof import('mongoose')): { overrides: Partial, priority: number, session?: ClientSession, - options?: { expectEmpty?: boolean }, + options?: { expectEmpty?: boolean; preservePriority?: boolean }, ) => Promise; patchConfigFields: ( principalType: PrincipalType, @@ -149,7 +149,7 @@ export function createConfigMethods(mongoose: typeof import('mongoose')): { overrides: Partial, priority: number, session?: ClientSession, - options?: { expectEmpty?: boolean }, + options?: { expectEmpty?: boolean; preservePriority?: boolean }, ): Promise { const Config = mongoose.models.Config as Model; @@ -168,9 +168,10 @@ export function createConfigMethods(mongoose: typeof import('mongoose')): { $set: { principalModel, overrides, - priority, + ...(options?.preservePriority ? {} : { priority }), isActive: true, }, + ...(options?.preservePriority ? { $setOnInsert: { priority } } : {}), $inc: { configVersion: 1 }, }; diff --git a/packages/data-schemas/src/methods/index.ts b/packages/data-schemas/src/methods/index.ts index b98dab3a36..efde53c3e9 100644 --- a/packages/data-schemas/src/methods/index.ts +++ b/packages/data-schemas/src/methods/index.ts @@ -25,6 +25,7 @@ import { AUDIT_SCHEMA_VERSION, MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_LOG_LIMIT, + MAX_AUDIT_VERIFY_ROWS, type AuditLogMethods, } from './auditLog'; import { createShareMethods, type ShareMethods } from './share'; @@ -108,7 +109,7 @@ export { deriveStructuredFrontmatterFields, inferSkillFileCategory, }; -export { AUDIT_SCHEMA_VERSION, MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_LOG_LIMIT }; +export { AUDIT_SCHEMA_VERSION, MAX_AUDIT_EXPORT_ROWS, MAX_AUDIT_LOG_LIMIT, MAX_AUDIT_VERIFY_ROWS }; export type AllMethods = UserMethods & SessionMethods & diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 3c3986de93..7e3747d12c 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -21,6 +21,7 @@ let mongoServer: InstanceType; let Message: mongoose.Model; let saveMessage: ReturnType['saveMessage']; let getMessages: ReturnType['getMessages']; +let getMessageTextStats: ReturnType['getMessageTextStats']; let updateMessage: ReturnType['updateMessage']; let deleteMessages: ReturnType['deleteMessages']; let bulkSaveMessages: ReturnType['bulkSaveMessages']; @@ -39,6 +40,7 @@ beforeAll(async () => { const methods = createMessageMethods(mongoose); saveMessage = methods.saveMessage; getMessages = methods.getMessages; + getMessageTextStats = methods.getMessageTextStats; updateMessage = methods.updateMessage; deleteMessages = methods.deleteMessages; bulkSaveMessages = methods.bulkSaveMessages; @@ -240,6 +242,62 @@ describe('Message Operations', () => { expect(messages[0].text).toBe('First message'); expect(messages[1].text).toBe('Second message'); }); + + it('should limit retrieved messages when requested', async () => { + const conversationId = uuidv4(); + + await saveMessage(mockCtx, { + messageId: 'msg1', + conversationId, + text: 'First message', + user: 'user123', + }); + + await saveMessage(mockCtx, { + messageId: 'msg2', + conversationId, + text: 'Second message', + user: 'user123', + }); + + await saveMessage(mockCtx, { + messageId: 'msg3', + conversationId, + text: 'Third message', + user: 'user123', + }); + + const messages = await getMessages({ conversationId }, undefined, { limit: 2 }); + + expect(messages).toHaveLength(2); + expect(messages[0].text).toBe('First message'); + expect(messages[1].text).toBe('Second message'); + }); + + it('should retrieve message text stats without returning message bodies', async () => { + const conversationId = uuidv4(); + + await saveMessage(mockCtx, { + messageId: 'msg1', + conversationId, + text: 'hello', + quotes: ['a\nb', ''], + user: 'user123', + }); + + const stats = await getMessageTextStats({ conversationId, user: 'user123' }, { limit: 1 }); + + expect(stats).toEqual([ + { + messageId: 'msg1', + textBytes: 5, + quoteCount: 2, + quoteBytes: 3, + quoteLineCount: 3, + nonStringQuoteCount: 0, + }, + ]); + }); }); describe('deleteMessages', () => { diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 9ee15c7978..42273cbf5f 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -1,5 +1,5 @@ import { RetentionMode } from 'librechat-data-provider'; -import type { DeleteResult, FilterQuery, Model } from 'mongoose'; +import type { DeleteResult, FilterQuery, Model, PipelineStage } from 'mongoose'; import type { AppConfig, IMessage } from '~/types'; import { createTempChatExpirationDate } from '~/utils/tempChatRetention'; import { createFallbackRetentionDate } from '~/utils/retention'; @@ -9,6 +9,24 @@ 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; +interface MessageQueryOptions { + limit?: number; + sort?: Record | false; +} + +interface MessageTextStatsOptions { + limit?: number; +} + +export interface MessageTextStats { + messageId: string; + textBytes: number; + quoteCount: number; + quoteBytes: number; + quoteLineCount: number; + nonStringQuoteCount: number; +} + export interface MessageMethods { saveMessage( ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] }, @@ -37,7 +55,15 @@ export interface MessageMethods { userId: string, params: { messageId: string; conversationId: string }, ): Promise; - getMessages(filter: FilterQuery, select?: string): Promise; + getMessages( + filter: FilterQuery, + select?: string, + options?: MessageQueryOptions, + ): Promise; + getMessageTextStats( + filter: FilterQuery, + options?: MessageTextStatsOptions, + ): Promise; getMessage(params: { user: string; messageId: string }): Promise; getMessagesByCursor( filter: FilterQuery, @@ -323,20 +349,118 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa /** * Retrieves messages from the database. */ - async function getMessages(filter: FilterQuery, select?: string) { + async function getMessages( + filter: FilterQuery, + select?: string, + options: MessageQueryOptions = {}, + ) { try { const Message = mongoose.models.Message as Model; + const query = Message.find(filter); if (select) { - return await Message.find(filter).select(select).sort({ createdAt: 1 }).lean(); + query.select(select); + } + if (options.sort !== false) { + query.sort(options.sort ?? { createdAt: 1 }); + } + if (options.limit != null && options.limit > 0) { + query.limit(options.limit); } - return await Message.find(filter).sort({ createdAt: 1 }).lean(); + return await query.lean(); } catch (err) { logger.error('Error getting messages:', err); throw err; } } + async function getMessageTextStats( + filter: FilterQuery, + options: MessageTextStatsOptions = {}, + ) { + try { + const Message = mongoose.models.Message as Model; + const pipeline: PipelineStage[] = [{ $match: filter }]; + if (options.limit != null && options.limit > 0) { + pipeline.push({ $limit: options.limit }); + } + pipeline.push({ + $project: { + _id: 0, + messageId: 1, + textBytes: { + $cond: [{ $eq: [{ $type: '$text' }, 'string'] }, { $strLenBytes: '$text' }, 0], + }, + quoteCount: { + $cond: [{ $isArray: '$quotes' }, { $size: '$quotes' }, 0], + }, + quoteBytes: { + $cond: [ + { $isArray: '$quotes' }, + { + $sum: { + $map: { + input: '$quotes', + as: 'quote', + in: { + $cond: [ + { $eq: [{ $type: '$$quote' }, 'string'] }, + { $strLenBytes: '$$quote' }, + 0, + ], + }, + }, + }, + }, + 0, + ], + }, + quoteLineCount: { + $cond: [ + { $isArray: '$quotes' }, + { + $sum: { + $map: { + input: '$quotes', + as: 'quote', + in: { + $cond: [ + { $eq: [{ $type: '$$quote' }, 'string'] }, + { $size: { $split: ['$$quote', '\n'] } }, + 0, + ], + }, + }, + }, + }, + 0, + ], + }, + nonStringQuoteCount: { + $cond: [ + { $isArray: '$quotes' }, + { + $size: { + $filter: { + input: '$quotes', + as: 'quote', + cond: { $ne: [{ $type: '$$quote' }, 'string'] }, + }, + }, + }, + 0, + ], + }, + }, + }); + + return await Message.aggregate(pipeline); + } catch (err) { + logger.error('Error getting message text stats:', err); + throw err; + } + } + /** * Retrieves a single message from the database. */ @@ -423,6 +547,7 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateMessage, deleteMessagesSince, getMessages, + getMessageTextStats, getMessage, getMessagesByCursor, searchMessages, diff --git a/packages/data-schemas/src/types/auditLog.ts b/packages/data-schemas/src/types/auditLog.ts index f1f6a94dda..0a851fec58 100644 --- a/packages/data-schemas/src/types/auditLog.ts +++ b/packages/data-schemas/src/types/auditLog.ts @@ -166,6 +166,10 @@ export interface PurgeAuditLogResult { } export interface VerifyAuditChainOptions { + /** Stop verification early when the caller has gone away. */ + isCancelled?: () => boolean; + /** Maximum rows to inspect before returning a bounded, non-OK result. */ + maxRows?: number; /** When the chain no longer starts at `seq: 1` (a prefix was purged), the * verifier requires this boundary to distinguish an authorized retention purge * from an attacker deleting the oldest rows. Without it, a non-genesis start