diff --git a/api/server/routes/memories.js b/api/server/routes/memories.js index 1ea73309fc..d19819dd80 100644 --- a/api/server/routes/memories.js +++ b/api/server/routes/memories.js @@ -2,6 +2,8 @@ const express = require('express'); const { Tokenizer, generateCheckAccess, + getMemoryAgentIdParam, + createAgentMemoryPartitionMiddleware, blockFilteredMemoryContent, projectStoredMemories, createMemoryManagementHandlers, @@ -12,7 +14,8 @@ const { ResourceType, Permissions, } = require('librechat-data-provider'); -const { findAccessibleResources } = require('~/server/services/PermissionService'); +const { checkPermission, findAccessibleResources } = require('~/server/services/PermissionService'); +const { hasCapability } = require('~/server/middleware/roles/capabilities'); const { getAllUserMemories, getUserMemories, @@ -23,6 +26,7 @@ const { setMemory, setMemoryById, deleteMemoryById, + getAgent, getAgents, } = require('~/models'); const { requireJwtAuth, configMiddleware } = require('~/server/middleware'); @@ -65,9 +69,37 @@ const opaqueMemoryHandlers = createMemoryManagementHandlers({ router.use(requireJwtAuth); -/** Normalizes the optional agent partition param; undefined = shared personal pool */ -const getAgentIdParam = (value) => - typeof value === 'string' && value.trim() !== '' ? value.trim() : undefined; +const agentPartitionDependencies = { + getAgent, + getRoleByName, + hasCapability, + checkPermission, +}; +const validateBodyAgentPartition = createAgentMemoryPartitionMiddleware({ + source: 'body', + ...agentPartitionDependencies, +}); +const validateQueryAgentPartition = createAgentMemoryPartitionMiddleware({ + source: 'query', + ...agentPartitionDependencies, +}); +const validateDeletedAgentPartition = createAgentMemoryPartitionMiddleware({ + source: 'query', + allowMissingAgent: true, + ...agentPartitionDependencies, +}); +const createMemoryMiddleware = [ + memoryPayloadLimit, + checkMemoryCreate, + validateBodyAgentPartition, + configMiddleware, +]; +const updateMemoryMiddleware = [ + memoryPayloadLimit, + checkMemoryUpdate, + validateQueryAgentPartition, + configMiddleware, +]; /** Resolves agent display names for agent-partitioned memories, restricted * to agents the requester can VIEW — `agentId` is caller-supplied on write, @@ -144,9 +176,9 @@ router.get('/', checkMemoryRead, configMiddleware, async (req, res) => { * Body: { key: string, value: string } * Returns 201 and { created: true, memory: } when successful. */ -router.post('/', memoryPayloadLimit, checkMemoryCreate, configMiddleware, async (req, res) => { +router.post('/', createMemoryMiddleware, async (req, res) => { const { key, value } = req.body; - const agentId = getAgentIdParam(req.body.agentId); + const agentId = getMemoryAgentIdParam(req.body.agentId); if (typeof key !== 'string' || key.trim() === '') { return res.status(400).json({ error: 'Key is required and must be a non-empty string.' }); @@ -257,10 +289,16 @@ router.patch( '/id/:id', memoryPayloadLimit, checkMemoryUpdate, + validateQueryAgentPartition, configMiddleware, opaqueMemoryHandlers.updateById, ); -router.delete('/id/:id', checkMemoryDelete, opaqueMemoryHandlers.deleteById); +router.delete( + '/id/:id', + checkMemoryDelete, + validateDeletedAgentPartition, + opaqueMemoryHandlers.deleteById, +); /** * PATCH /memories/:key @@ -268,10 +306,10 @@ router.delete('/id/:id', checkMemoryDelete, opaqueMemoryHandlers.deleteById); * Body: { key?: string, value: string } * Returns 200 and { updated: true, memory: } when successful. */ -router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, async (req, res) => { +router.patch('/:key', updateMemoryMiddleware, async (req, res) => { const { key: urlKey } = req.params; const { key: bodyKey, value } = req.body || {}; - const agentId = getAgentIdParam(req.query.agentId); + const agentId = getMemoryAgentIdParam(req.query.agentId); if (typeof value !== 'string' || value.trim() === '') { return res.status(400).json({ error: 'Value is required and must be a non-empty string.' }); @@ -358,9 +396,9 @@ router.patch('/:key', memoryPayloadLimit, checkMemoryUpdate, configMiddleware, a * Deletes a memory entry for the authenticated user. * Returns 200 and { deleted: true } when successful. */ -router.delete('/:key', checkMemoryDelete, async (req, res) => { +router.delete('/:key', checkMemoryDelete, validateDeletedAgentPartition, async (req, res) => { const { key } = req.params; - const agentId = getAgentIdParam(req.query.agentId); + const agentId = getMemoryAgentIdParam(req.query.agentId); try { const result = await deleteMemory({ userId: req.user.id, key, agentId }); diff --git a/api/server/routes/memories.test.js b/api/server/routes/memories.test.js index f7130afc72..54727b6490 100644 --- a/api/server/routes/memories.test.js +++ b/api/server/routes/memories.test.js @@ -2,9 +2,15 @@ const express = require('express'); const request = require('supertest'); const mockGetAllUserMemories = jest.fn(); +const mockGetUserMemories = jest.fn(); +const mockCreateMemory = jest.fn(); +const mockDeleteMemory = jest.fn(); const mockProjectStoredMemories = jest.fn((memories) => memories); const mockSetMemoryById = jest.fn(); const mockDeleteMemoryById = jest.fn(); +const mockGetAgent = jest.fn(); +const mockCheckPermission = jest.fn(); +const mockHasCapability = jest.fn(); const mockOpaqueUpdateById = jest.fn((_req, res) => res.status(202).json({ delegated: 'update' })); const mockOpaqueDeleteById = jest.fn((_req, res) => res.status(202).json({ delegated: 'delete' })); let mockFilters; @@ -15,6 +21,33 @@ jest.mock('@librechat/api', () => ({ inspectContent: jest.fn(() => null), extractMemoryContent: jest.fn(() => []), projectStoredMemories: (...args) => mockProjectStoredMemories(...args), + blockFilteredMemoryContent: jest.fn(() => false), + getMemoryAgentIdParam: (value) => + typeof value === 'string' && value.trim() !== '' ? value.trim() : undefined, + createAgentMemoryPartitionMiddleware: + ({ source, allowMissingAgent = false }) => + async (req, res, next) => { + const value = req[source]?.agentId; + const agentId = typeof value === 'string' && value.trim() !== '' ? value.trim() : undefined; + if (!agentId) { + return next(); + } + const agent = await mockGetAgent({ id: agentId }, { _id: 1 }); + if (!agent) { + return allowMissingAgent ? next() : res.status(404).json({ error: 'Agent not found.' }); + } + if (await mockHasCapability(req.user, 'manage:agents')) { + return next(); + } + const allowed = await mockCheckPermission({ + userId: req.user.id, + role: req.user.role, + resourceType: 'agent', + resourceId: agent._id, + requiredPermission: 1, + }); + return allowed ? next() : res.status(403).json({ error: 'Agent access denied.' }); + }, createMemoryManagementHandlers: jest.fn(() => ({ updateById: (...args) => mockOpaqueUpdateById(...args), deleteById: (...args) => mockOpaqueDeleteById(...args), @@ -36,19 +69,25 @@ jest.mock('librechat-data-provider', () => ({ })); jest.mock('~/server/services/PermissionService', () => ({ + checkPermission: (...args) => mockCheckPermission(...args), findAccessibleResources: jest.fn().mockResolvedValue([]), })); +jest.mock('~/server/middleware/roles/capabilities', () => ({ + hasCapability: (...args) => mockHasCapability(...args), +})); + jest.mock('~/models', () => ({ getAllUserMemories: (...args) => mockGetAllUserMemories(...args), - getUserMemories: jest.fn(), + getUserMemories: (...args) => mockGetUserMemories(...args), toggleUserMemories: jest.fn(), getRoleByName: jest.fn(), - createMemory: jest.fn(), - deleteMemory: jest.fn(), + createMemory: (...args) => mockCreateMemory(...args), + deleteMemory: (...args) => mockDeleteMemory(...args), setMemory: jest.fn(), setMemoryById: (...args) => mockSetMemoryById(...args), deleteMemoryById: (...args) => mockDeleteMemoryById(...args), + getAgent: (...args) => mockGetAgent(...args), getAgents: jest.fn(), })); @@ -75,6 +114,130 @@ const buildApp = () => { return app; }; +describe('agent memory partition authorization', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockFilters = undefined; + mockGetAgent.mockResolvedValue({ _id: 'agent-object-id', id: 'agent-1' }); + mockCheckPermission.mockResolvedValue(true); + mockHasCapability.mockResolvedValue(false); + }); + + it('rejects a nonexistent client-selected agent partition', async () => { + mockGetAgent.mockResolvedValue(null); + + const response = await request(buildApp()) + .post('/api/memories') + .send({ key: 'timezone', value: 'UTC', agentId: ' attacker-partition ' }); + + expect(response.status).toBe(404); + expect(response.body).toEqual({ error: 'Agent not found.' }); + expect(mockGetAgent).toHaveBeenCalledWith({ id: 'attacker-partition' }, { _id: 1 }); + expect(mockGetUserMemories).not.toHaveBeenCalled(); + expect(mockCreateMemory).not.toHaveBeenCalled(); + }); + + it('rejects an agent partition the requester cannot view', async () => { + mockCheckPermission.mockResolvedValue(false); + + const response = await request(buildApp()) + .post('/api/memories') + .send({ key: 'timezone', value: 'UTC', agentId: 'agent-1' }); + + expect(response.status).toBe(403); + expect(response.body).toEqual({ error: 'Agent access denied.' }); + expect(mockCheckPermission).toHaveBeenCalledWith({ + userId: 'user-1', + role: 'USER', + resourceType: 'agent', + resourceId: 'agent-object-id', + requiredPermission: 1, + }); + expect(mockGetUserMemories).not.toHaveBeenCalled(); + expect(mockCreateMemory).not.toHaveBeenCalled(); + }); + + it('allows a partition for an agent the requester can view', async () => { + mockGetUserMemories + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ key: 'timezone', value: 'UTC', agentId: 'agent-1' }]); + mockCreateMemory.mockResolvedValue({ ok: true }); + + const response = await request(buildApp()) + .post('/api/memories') + .send({ key: 'timezone', value: 'UTC', agentId: 'agent-1' }); + + expect(response.status).toBe(201); + expect(mockCreateMemory).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'user-1', agentId: 'agent-1' }), + ); + }); + + it('allows agent managers without consulting the resource ACL', async () => { + mockHasCapability.mockResolvedValue(true); + mockCheckPermission.mockResolvedValue(false); + mockGetUserMemories + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ key: 'timezone', value: 'UTC', agentId: 'agent-1' }]); + mockCreateMemory.mockResolvedValue({ ok: true }); + + const response = await request(buildApp()) + .post('/api/memories') + .send({ key: 'timezone', value: 'UTC', agentId: 'agent-1' }); + + expect(response.status).toBe(201); + expect(mockCheckPermission).not.toHaveBeenCalled(); + expect(mockCreateMemory).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'user-1', agentId: 'agent-1' }), + ); + }); + + it.each([ + ['patch', '/api/memories/id/507f1f77bcf86cd799439011?agentId=agent-1'], + ['delete', '/api/memories/id/507f1f77bcf86cd799439011?agentId=agent-1'], + ['patch', '/api/memories/timezone?agentId=agent-1'], + ['delete', '/api/memories/timezone?agentId=agent-1'], + ])('rejects unauthorized %s mutations of an agent partition', async (method, path) => { + mockCheckPermission.mockResolvedValue(false); + + const pendingRequest = request(buildApp())[method](path); + const response = + method === 'patch' ? await pendingRequest.send({ value: 'UTC' }) : await pendingRequest; + + expect(response.status).toBe(403); + expect(response.body).toEqual({ error: 'Agent access denied.' }); + expect(mockOpaqueUpdateById).not.toHaveBeenCalled(); + expect(mockOpaqueDeleteById).not.toHaveBeenCalled(); + }); + + it('allows deleting a user-owned partition after its agent is removed', async () => { + mockGetAgent.mockResolvedValue(null); + mockDeleteMemory.mockResolvedValue({ ok: true }); + + const response = await request(buildApp()).delete( + '/api/memories/timezone?agentId=removed-agent', + ); + + expect(response.status).toBe(200); + expect(mockDeleteMemory).toHaveBeenCalledWith({ + userId: 'user-1', + key: 'timezone', + agentId: 'removed-agent', + }); + }); + + it('allows deleting a user-owned opaque memory after its agent is removed', async () => { + mockGetAgent.mockResolvedValue(null); + + const response = await request(buildApp()).delete( + '/api/memories/id/507f1f77bcf86cd799439011?agentId=removed-agent', + ); + + expect(response.status).toBe(202); + expect(mockOpaqueDeleteById).toHaveBeenCalledTimes(1); + }); +}); + describe('GET /api/memories content policy', () => { beforeEach(() => { jest.clearAllMocks(); @@ -169,6 +332,9 @@ describe('GET /api/memories content policy', () => { describe('opaque memory management routes', () => { beforeEach(() => { jest.clearAllMocks(); + mockGetAgent.mockResolvedValue({ _id: 'agent-object-id', id: 'agent-1' }); + mockCheckPermission.mockResolvedValue(true); + mockHasCapability.mockResolvedValue(false); mockFilters = { memories: { pii: { diff --git a/packages/api/src/memory/authorization.spec.ts b/packages/api/src/memory/authorization.spec.ts new file mode 100644 index 0000000000..a7d96e0d96 --- /dev/null +++ b/packages/api/src/memory/authorization.spec.ts @@ -0,0 +1,187 @@ +import { Types } from 'mongoose'; +import { Permissions, PermissionTypes } from 'librechat-data-provider'; +import type { IRole, IUser } from '@librechat/data-schemas'; +import type { Request, Response } from 'express'; +import { + getMemoryAgentIdParam, + getAgentMemoryPartitionAccess, + createAgentMemoryPartitionMiddleware, +} from './authorization'; + +const user = { id: new Types.ObjectId().toString(), role: 'USER' } as IUser; +const agent = { _id: new Types.ObjectId() }; + +function createRequest(): Request { + return {} as Request; +} + +function createRole(canUseAgents: boolean): IRole { + return { + permissions: { + [PermissionTypes.AGENTS]: { + [Permissions.USE]: canUseAgents, + }, + }, + } as IRole; +} + +function createDependencies({ + canUseAgents = true, + canManageAgents = false, + canViewAgent = true, +} = {}) { + return { + getAgent: jest.fn().mockResolvedValue(agent), + getRoleByName: jest.fn().mockResolvedValue(createRole(canUseAgents)), + hasCapability: jest.fn().mockResolvedValue(canManageAgents), + checkPermission: jest.fn().mockResolvedValue(canViewAgent), + }; +} + +function createResponse(): Response { + const response = {} as Response; + response.status = jest.fn(() => response); + response.json = jest.fn(() => response); + return response; +} + +describe('getAgentMemoryPartitionAccess', () => { + it('normalizes only non-empty string partition IDs', () => { + expect(getMemoryAgentIdParam(' agent-1 ')).toBe('agent-1'); + expect(getMemoryAgentIdParam(' ')).toBeUndefined(); + expect(getMemoryAgentIdParam(['agent-1'])).toBeUndefined(); + }); + + it('starts independent agent, capability, and role reads concurrently', async () => { + let resolveAgent!: (value: typeof agent) => void; + const agentPromise = new Promise((resolve) => { + resolveAgent = resolve; + }); + const dependencies = createDependencies(); + dependencies.getAgent.mockReturnValue(agentPromise); + + const accessPromise = getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'agent-1', + ...dependencies, + }); + + expect(dependencies.getAgent).toHaveBeenCalledTimes(1); + expect(dependencies.getAgent).toHaveBeenCalledWith({ id: 'agent-1' }, { _id: 1 }); + expect(dependencies.hasCapability).toHaveBeenCalledTimes(1); + expect(dependencies.getRoleByName).toHaveBeenCalledTimes(1); + + resolveAgent(agent); + await expect(accessPromise).resolves.toBe('allowed'); + }); + + it('denies access when the role cannot use agents', async () => { + const dependencies = createDependencies({ canUseAgents: false }); + + await expect( + getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'agent-1', + ...dependencies, + }), + ).resolves.toBe('denied'); + expect(dependencies.checkPermission).not.toHaveBeenCalled(); + }); + + it('allows agent managers without consulting the resource ACL', async () => { + const dependencies = createDependencies({ canManageAgents: true, canViewAgent: false }); + + await expect( + getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'agent-1', + ...dependencies, + }), + ).resolves.toBe('allowed'); + expect(dependencies.checkPermission).not.toHaveBeenCalled(); + }); + + it('uses the agent VIEW ACL for non-managers', async () => { + const dependencies = createDependencies({ canViewAgent: false }); + + await expect( + getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'agent-1', + ...dependencies, + }), + ).resolves.toBe('denied'); + expect(dependencies.checkPermission).toHaveBeenCalledWith( + expect.objectContaining({ resourceId: agent._id }), + ); + }); + + it('reports a missing agent without consulting its resource ACL', async () => { + const dependencies = createDependencies(); + dependencies.getAgent.mockResolvedValue(null); + + await expect( + getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'missing-agent', + ...dependencies, + }), + ).resolves.toBe('not_found'); + expect(dependencies.checkPermission).not.toHaveBeenCalled(); + }); + + it('does not disclose missing agents when the role cannot use agents', async () => { + const dependencies = createDependencies({ canUseAgents: false }); + dependencies.getAgent.mockResolvedValue(null); + + await expect( + getAgentMemoryPartitionAccess({ + req: createRequest(), + user, + agentId: 'missing-agent', + ...dependencies, + }), + ).resolves.toBe('denied'); + expect(dependencies.checkPermission).not.toHaveBeenCalled(); + }); + + it('allows deletion middleware to clean up a missing agent partition', async () => { + const dependencies = createDependencies({ canUseAgents: false }); + dependencies.getAgent.mockResolvedValue(null); + const middleware = createAgentMemoryPartitionMiddleware({ + source: 'query', + allowMissingAgent: true, + ...dependencies, + }); + const request = { user, query: { agentId: 'missing-agent' } } as unknown as Request; + const response = createResponse(); + const next = jest.fn(); + + await middleware(request, response, next); + + expect(next).toHaveBeenCalledTimes(1); + expect(response.status).not.toHaveBeenCalled(); + }); + + it('returns forbidden before mutation middleware proceeds when agent use is disabled', async () => { + const dependencies = createDependencies({ canUseAgents: false }); + const middleware = createAgentMemoryPartitionMiddleware({ + source: 'body', + ...dependencies, + }); + const request = { user, body: { agentId: 'agent-1' } } as unknown as Request; + const response = createResponse(); + const next = jest.fn(); + + await middleware(request, response, next); + + expect(next).not.toHaveBeenCalled(); + expect(response.status).toHaveBeenCalledWith(403); + expect(response.json).toHaveBeenCalledWith({ error: 'Agent access denied.' }); + }); +}); diff --git a/packages/api/src/memory/authorization.ts b/packages/api/src/memory/authorization.ts new file mode 100644 index 0000000000..8bb5dc3d1e --- /dev/null +++ b/packages/api/src/memory/authorization.ts @@ -0,0 +1,126 @@ +import { ResourceCapabilityMap } from '@librechat/data-schemas'; +import { + Permissions, + PermissionBits, + ResourceType, + PermissionTypes, +} from 'librechat-data-provider'; +import type { NextFunction, Request, Response } from 'express'; +import type { IAgent, IUser } from '@librechat/data-schemas'; +import type { CheckAccessParams } from '../middleware/access'; +import { checkAccess } from '../middleware/access'; + +export type AgentMemoryPartitionAccess = 'allowed' | 'denied' | 'not_found'; + +interface CheckPermissionParams { + userId: string; + role?: string; + resourceType: ResourceType; + resourceId: IAgent['_id']; + requiredPermission: number; +} + +interface AgentMemoryPartitionAccessDependencies { + getAgent: (query: { id: string }, projection: { _id: 1 }) => Promise | null>; + getRoleByName: CheckAccessParams['getRoleByName']; + hasCapability: ( + user: IUser, + capability: (typeof ResourceCapabilityMap)[ResourceType.AGENT], + ) => Promise; + checkPermission: (params: CheckPermissionParams) => Promise; +} + +interface AgentMemoryPartitionAccessParams extends AgentMemoryPartitionAccessDependencies { + req: Request; + user: IUser; + agentId: string; + allowMissingAgent?: boolean; +} + +interface AgentMemoryPartitionMiddlewareParams extends AgentMemoryPartitionAccessDependencies { + source: 'body' | 'query'; + allowMissingAgent?: boolean; +} + +export function getMemoryAgentIdParam(value: unknown): string | undefined { + return typeof value === 'string' && value.trim() !== '' ? value.trim() : undefined; +} + +export async function getAgentMemoryPartitionAccess({ + req, + user, + agentId, + allowMissingAgent = false, + getAgent, + getRoleByName, + hasCapability, + checkPermission, +}: AgentMemoryPartitionAccessParams): Promise { + const [agent, canManageAgents, canUseAgents] = await Promise.all([ + getAgent({ id: agentId }, { _id: 1 }), + hasCapability(user, ResourceCapabilityMap[ResourceType.AGENT]).catch(() => false), + checkAccess({ + req, + user, + permissionType: PermissionTypes.AGENTS, + permissions: [Permissions.USE], + getRoleByName, + }).catch(() => false), + ]); + + if (!canUseAgents && !(allowMissingAgent && !agent)) { + return 'denied'; + } + if (!agent) { + return 'not_found'; + } + if (canManageAgents) { + return 'allowed'; + } + + const canViewAgent = await checkPermission({ + userId: user.id, + role: user.role, + resourceType: ResourceType.AGENT, + resourceId: agent._id, + requiredPermission: PermissionBits.VIEW, + }); + return canViewAgent ? 'allowed' : 'denied'; +} + +export function createAgentMemoryPartitionMiddleware({ + source, + allowMissingAgent = false, + ...dependencies +}: AgentMemoryPartitionMiddlewareParams): ( + req: Request, + res: Response, + next: NextFunction, +) => Promise { + return async (req, res, next) => { + const value = source === 'body' ? req.body?.agentId : req.query.agentId; + const agentId = getMemoryAgentIdParam(value); + if (!agentId) { + return next(); + } + + try { + const agentAccess = await getAgentMemoryPartitionAccess({ + req, + user: req.user as IUser, + agentId, + allowMissingAgent, + ...dependencies, + }); + if (agentAccess === 'not_found') { + return allowMissingAgent ? next() : res.status(404).json({ error: 'Agent not found.' }); + } + if (agentAccess === 'denied') { + return res.status(403).json({ error: 'Agent access denied.' }); + } + return next(); + } catch (_error) { + return res.status(500).json({ error: 'Failed to validate agent access.' }); + } + }; +} diff --git a/packages/api/src/memory/index.ts b/packages/api/src/memory/index.ts index ff4caf384f..c916eded62 100644 --- a/packages/api/src/memory/index.ts +++ b/packages/api/src/memory/index.ts @@ -1,3 +1,4 @@ export * from './config'; +export * from './authorization'; export * from './handlers'; export * from './protection'; diff --git a/packages/api/src/middleware/access.spec.ts b/packages/api/src/middleware/access.spec.ts index c77508b153..678e7d3db4 100644 --- a/packages/api/src/middleware/access.spec.ts +++ b/packages/api/src/middleware/access.spec.ts @@ -364,7 +364,36 @@ describe('access middleware', () => { }), ).resolves.toBe(true); - expect(defaultParams.getRoleByName).toHaveBeenCalledTimes(2); + expect(defaultParams.getRoleByName).toHaveBeenCalledTimes(1); + }); + + it('should reuse a synchronous role lookup across different permission types', async () => { + const role = { + name: 'user', + permissions: { + [PermissionTypes.MEMORIES]: { [Permissions.USE]: true }, + [PermissionTypes.AGENTS]: { [Permissions.USE]: true }, + }, + } as unknown as IRole; + defaultParams.getRoleByName.mockReturnValue(role); + const req = mockReq as Request; + + await expect( + checkAccess({ + ...defaultParams, + req, + permissionType: PermissionTypes.MEMORIES, + }), + ).resolves.toBe(true); + await expect( + checkAccess({ + ...defaultParams, + req, + permissionType: PermissionTypes.AGENTS, + }), + ).resolves.toBe(true); + + expect(defaultParams.getRoleByName).toHaveBeenCalledTimes(1); }); }); diff --git a/packages/api/src/middleware/access.ts b/packages/api/src/middleware/access.ts index 3b90d49d40..554ed34481 100644 --- a/packages/api/src/middleware/access.ts +++ b/packages/api/src/middleware/access.ts @@ -42,8 +42,60 @@ export type CheckAccessWithRequestCacheParams = Omit< >; type RequestPermissionCache = Map>; +type RequestRoleCache = Map>; const requestPermissionCacheKey = '__librechatRequestPermissionCache'; +const requestRoleCacheKey = '__librechatRequestRoleCache'; + +function getRequestRoleCache(req?: ServerRequest): RequestRoleCache | null { + if (!req) { + return null; + } + + const reqWithCache = req as ServerRequest & { + [requestRoleCacheKey]?: RequestRoleCache; + }; + + if (!reqWithCache[requestRoleCacheKey]) { + Object.defineProperty(reqWithCache, requestRoleCacheKey, { + value: new Map>(), + enumerable: false, + }); + } + + return reqWithCache[requestRoleCacheKey] ?? null; +} + +async function getRoleForAccess({ + req, + roleName, + getRoleByName, +}: { + req?: ServerRequest; + roleName: string; + getRoleByName: CheckAccessParams['getRoleByName']; +}): Promise { + const cache = getRequestRoleCache(req); + if (!cache) { + return await getRoleByName(roleName); + } + + let cachedRole = cache.get(roleName); + if (!cachedRole) { + try { + cachedRole = Promise.resolve(getRoleByName(roleName)); + } catch (error) { + cachedRole = Promise.reject(error); + } + cachedRole = cachedRole.catch((error) => { + cache.delete(roleName); + throw error; + }); + cache.set(roleName, cachedRole); + } + + return await cachedRole; +} function getRequestPermissionCache(req?: ServerRequest): RequestPermissionCache | null { if (!req) { @@ -100,7 +152,7 @@ export const checkAccess = async ({ return false; } - const role = await getRoleByName(user.role); + const role = await getRoleForAccess({ req, roleName: user.role, getRoleByName }); const permissionValue = role?.permissions?.[permissionType as keyof typeof role.permissions]; if (role && role.permissions && permissionValue) { const hasAnyPermission = permissions.every((permission) => {