diff --git a/packages/data-schemas/src/methods/aclEntry.tenant.spec.ts b/packages/data-schemas/src/methods/aclEntry.tenant.spec.ts new file mode 100644 index 0000000000..a216ce9e63 --- /dev/null +++ b/packages/data-schemas/src/methods/aclEntry.tenant.spec.ts @@ -0,0 +1,114 @@ +import mongoose from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import { PrincipalType, PrincipalModel, PermissionBits } from 'librechat-data-provider'; +import type { IAclEntry } from '..'; +import { createAclEntryMethods } from './aclEntry'; +import { createModels } from '../models'; +import { tenantStorage } from '../config/tenantContext'; + +jest.mock('~/config/winston', () => ({ + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), +})); + +const TENANT_A = 'tenant-aaaaaaaaaaaaaaaaaaaa'; +const TENANT_B = 'tenant-bbbbbbbbbbbbbbbbbbbb'; +const RESOURCE_TYPE = 'agent'; + +let mongoServer: MongoMemoryServer; +let AclEntry: mongoose.Model; +let methods: ReturnType; + +function runAs(tenantId: string, fn: () => Promise): Promise { + return tenantStorage.run({ tenantId }, fn); +} + +async function seedAcl(tenantId: string, doc: Record): Promise { + await runAs(tenantId, async () => { + await new AclEntry(doc).save(); + }); +} + +beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + AclEntry = mongoose.models.AclEntry as mongoose.Model; + methods = createAclEntryMethods(mongoose); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); +}); + +beforeEach(async () => { + await AclEntry.deleteMany({}); +}); + +describe('AclEntry resource lookups are scoped to the active tenant', () => { + it('findAccessibleResources returns only the current tenant resources', async () => { + const principalId = new mongoose.Types.ObjectId(); + const resourceA = new mongoose.Types.ObjectId(); + const resourceB = new mongoose.Types.ObjectId(); + + await seedAcl(TENANT_A, { + principalType: PrincipalType.USER, + principalModel: PrincipalModel.USER, + principalId, + resourceType: RESOURCE_TYPE, + resourceId: resourceA, + permBits: PermissionBits.VIEW, + }); + await seedAcl(TENANT_B, { + principalType: PrincipalType.USER, + principalModel: PrincipalModel.USER, + principalId, + resourceType: RESOURCE_TYPE, + resourceId: resourceB, + permBits: PermissionBits.VIEW, + }); + + const principals = [{ principalType: PrincipalType.USER, principalId }]; + + const aResources = await runAs(TENANT_A, () => + methods.findAccessibleResources(principals, RESOURCE_TYPE, PermissionBits.VIEW), + ); + expect(aResources.map(String)).toEqual([String(resourceA)]); + + const bResources = await runAs(TENANT_B, () => + methods.findAccessibleResources(principals, RESOURCE_TYPE, PermissionBits.VIEW), + ); + expect(bResources.map(String)).toEqual([String(resourceB)]); + }); + + it('findPublicResourceIds returns only the current tenant public resources', async () => { + const publicA = new mongoose.Types.ObjectId(); + const publicB = new mongoose.Types.ObjectId(); + + await seedAcl(TENANT_A, { + principalType: PrincipalType.PUBLIC, + resourceType: RESOURCE_TYPE, + resourceId: publicA, + permBits: PermissionBits.VIEW, + }); + await seedAcl(TENANT_B, { + principalType: PrincipalType.PUBLIC, + resourceType: RESOURCE_TYPE, + resourceId: publicB, + permBits: PermissionBits.VIEW, + }); + + const aPublic = await runAs(TENANT_A, () => + methods.findPublicResourceIds(RESOURCE_TYPE, PermissionBits.VIEW), + ); + expect(aPublic.map(String)).toEqual([String(publicA)]); + + const bPublic = await runAs(TENANT_B, () => + methods.findPublicResourceIds(RESOURCE_TYPE, PermissionBits.VIEW), + ); + expect(bPublic.map(String)).toEqual([String(publicB)]); + }); +}); diff --git a/packages/data-schemas/src/methods/agentCategory.tenant.spec.ts b/packages/data-schemas/src/methods/agentCategory.tenant.spec.ts new file mode 100644 index 0000000000..cccae345b9 --- /dev/null +++ b/packages/data-schemas/src/methods/agentCategory.tenant.spec.ts @@ -0,0 +1,60 @@ +import mongoose from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import type { IAgentCategory } from '..'; +import { createAgentCategoryMethods } from './agentCategory'; +import { createModels } from '../models'; +import { tenantStorage } from '../config/tenantContext'; + +jest.mock('~/config/winston', () => ({ + error: jest.fn(), + warn: jest.fn(), + info: jest.fn(), + debug: jest.fn(), +})); + +const TENANT_A = 'tenant-aaaaaaaaaaaaaaaaaaaa'; +const TENANT_B = 'tenant-bbbbbbbbbbbbbbbbbbbb'; + +let mongoServer: MongoMemoryServer; +let AgentCategory: mongoose.Model; +let methods: ReturnType; + +function runAs(tenantId: string, fn: () => Promise): Promise { + return tenantStorage.run({ tenantId }, fn); +} + +async function seedCategory(tenantId: string, value: string): Promise { + await runAs(tenantId, async () => { + await new AgentCategory({ value, label: value, isActive: true }).save(); + }); +} + +beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + AgentCategory = mongoose.models.AgentCategory as mongoose.Model; + methods = createAgentCategoryMethods(mongoose); +}); + +afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); +}); + +beforeEach(async () => { + await AgentCategory.deleteMany({}); +}); + +describe('getValidCategoryValues is scoped to the active tenant', () => { + it('returns only the current tenant category values', async () => { + await seedCategory(TENANT_A, 'alpha'); + await seedCategory(TENANT_B, 'beta'); + + const aValues = await runAs(TENANT_A, () => methods.getValidCategoryValues()); + expect(aValues).toEqual(['alpha']); + + const bValues = await runAs(TENANT_B, () => methods.getValidCategoryValues()); + expect(bValues).toEqual(['beta']); + }); +}); diff --git a/packages/data-schemas/src/models/plugins/tenantIsolation.spec.ts b/packages/data-schemas/src/models/plugins/tenantIsolation.spec.ts index 82e74d2cea..ba65e17be8 100644 --- a/packages/data-schemas/src/models/plugins/tenantIsolation.spec.ts +++ b/packages/data-schemas/src/models/plugins/tenantIsolation.spec.ts @@ -154,6 +154,35 @@ describe('applyTenantIsolation', () => { const tenantADoc = await TestModel.findOne({ tenantId: 'tenant-a' }).lean(); expect(tenantADoc!.name).toBe('updated'); }); + + it('injects tenantId filter into distinct', async () => { + const names = await tenantStorage.run({ tenantId: 'tenant-a' }, async () => + TestModel.distinct('name'), + ); + + expect(names).toEqual(['tenant-a-doc']); + }); + + it('injects tenantId filter into find().distinct() (op switches to distinct)', async () => { + const names = await tenantStorage.run({ tenantId: 'tenant-b' }, async () => + TestModel.find().distinct('name'), + ); + + expect(names).toEqual(['tenant-b-doc']); + }); + + it('does not scope distinct when context is absent (non-strict)', async () => { + const names = await TestModel.distinct('name'); + expect(names.sort()).toEqual(['no-tenant-doc', 'tenant-a-doc', 'tenant-b-doc']); + }); + + it('bypasses distinct filter for SYSTEM_TENANT_ID', async () => { + const names = await tenantStorage.run({ tenantId: SYSTEM_TENANT_ID }, async () => + TestModel.distinct('name'), + ); + + expect(names).toHaveLength(3); + }); }); describe('aggregate filtering', () => { diff --git a/packages/data-schemas/src/models/plugins/tenantIsolation.ts b/packages/data-schemas/src/models/plugins/tenantIsolation.ts index 06f8e035fb..1450610345 100644 --- a/packages/data-schemas/src/models/plugins/tenantIsolation.ts +++ b/packages/data-schemas/src/models/plugins/tenantIsolation.ts @@ -137,6 +137,7 @@ export function applyTenantIsolation(schema: Schema): void { schema.pre('find', queryMiddleware); schema.pre('findOne', queryMiddleware); + schema.pre('distinct', queryMiddleware); schema.pre('findOneAndUpdate', queryMiddleware); schema.pre('findOneAndDelete', queryMiddleware); schema.pre('findOneAndReplace', queryMiddleware);