⚡ perf: Minimize group membership query in principal resolution (#14055)

getUserPrincipals resolves a user's ACL principals on nearly every
authenticated request. It fetched full group documents (including entire
memberIds arrays) only to read each group _id, and always issued a
separate User lookup for idOnTheSource.

- Project { _id: 1 } on the memberIds group query so it returns only ids
  and can be served from the { memberIds: 1 } index instead of fetching
  and decoding whole group docs.
- Accept role and idOnTheSource from the already-loaded request user and
  thread them from the capability middleware, collapsing the hot path to
  a single indexed group query (idOnTheSource: null means known-local).

packages/api/src/middleware/capabilities.ts

@@ -14,7 +14,11 @@ import type { ServerRequest } from '~/types/http';
 
 interface CapabilityDeps {
   getUserPrincipals: (
-    params: { userId: string | Types.ObjectId; role?: string | null },
+    params: {
+      userId: string | Types.ObjectId;
+      role?: string | null;
+      idOnTheSource?: string | null;
+    },
     session?: ClientSession,
   ) => Promise<ResolvedPrincipal[]>;
   hasCapabilityForPrincipals: (params: {
@@ -28,6 +32,8 @@ export interface CapabilityUser {
   id: string;
   role: string;
   tenantId?: string;
+  /** External member id; pass `null` for local users to skip the fallback lookup. */
+  idOnTheSource?: string | null;
 }
 
 interface CapabilityStore {
@@ -153,7 +159,11 @@ export function generateCapabilityCheck(deps: CapabilityDeps): {
     if (cachedPrincipals) {
       principals = cachedPrincipals;
     } else {
-      principals = await getUserPrincipals({ userId: user.id, role: user.role });
+      principals = await getUserPrincipals({
+        userId: user.id,
+        role: user.role,
+        idOnTheSource: user.idOnTheSource,
+      });
       store?.principals.set(principalKey, principals);
     }
 
@@ -207,6 +217,7 @@ export function generateCapabilityCheck(deps: CapabilityDeps): {
           id,
           role: req.user.role ?? '',
           tenantId: (req.user as CapabilityUser).tenantId,
+          idOnTheSource: req.user.idOnTheSource ?? null,
         };
 
         if (await hasCapability(user, capability)) {

packages/data-schemas/src/methods/userGroup.spec.ts

@@ -1,6 +1,6 @@
 import mongoose, { Types } from 'mongoose';
-import { PrincipalType, SystemRoles } from 'librechat-data-provider';
 import { MongoMemoryServer } from 'mongodb-memory-server';
+import { PrincipalType, SystemRoles } from 'librechat-data-provider';
 import type * as t from '~/types';
 import { createUserGroupMethods } from './userGroup';
 import groupSchema from '~/schema/group';
@@ -382,6 +382,103 @@ describe('userGroup methods', () => {
       const rolePrincipal = principals.find((p) => p.principalType === PrincipalType.ROLE);
       expect(rolePrincipal).toBeUndefined();
     });
+
+    it('uses a supplied idOnTheSource authoritatively over the stored value', async () => {
+      const user = await createTestUser({ idOnTheSource: 'stored-ext-id' });
+      const group = await Group.create({
+        name: 'Team',
+        source: 'entra',
+        idOnTheSource: 'grp-ext',
+        memberIds: ['passed-ext-id'],
+      });
+
+      const principals = await methods.getUserPrincipals({
+        userId: user._id.toString(),
+        role: SystemRoles.USER,
+        idOnTheSource: 'passed-ext-id',
+      });
+
+      const groupPrincipal = principals.find((p) => p.principalType === PrincipalType.GROUP);
+      expect(groupPrincipal!.principalId!.toString()).toBe(group._id.toString());
+    });
+
+    it('resolves groups without a user document when idOnTheSource is supplied', async () => {
+      const missingUserId = new Types.ObjectId().toString();
+      const group = await Group.create({
+        name: 'Orphan Team',
+        source: 'entra',
+        idOnTheSource: 'grp-orphan',
+        memberIds: ['ext-orphan'],
+      });
+
+      const principals = await methods.getUserPrincipals({
+        userId: missingUserId,
+        role: SystemRoles.USER,
+        idOnTheSource: 'ext-orphan',
+      });
+
+      const groupPrincipal = principals.find((p) => p.principalType === PrincipalType.GROUP);
+      expect(groupPrincipal!.principalId!.toString()).toBe(group._id.toString());
+    });
+
+    it('treats idOnTheSource null as a local user keyed by user id', async () => {
+      const user = await createTestUser();
+      const group = await Group.create({
+        name: 'Local Team',
+        source: 'local',
+        memberIds: [user._id.toString()],
+      });
+
+      const principals = await methods.getUserPrincipals({
+        userId: user._id.toString(),
+        role: SystemRoles.USER,
+        idOnTheSource: null,
+      });
+
+      const groupPrincipal = principals.find((p) => p.principalType === PrincipalType.GROUP);
+      expect(groupPrincipal!.principalId!.toString()).toBe(group._id.toString());
+    });
+
+    it('falls back to resolving role and idOnTheSource from the DB when both are omitted', async () => {
+      const user = await createTestUser({ role: SystemRoles.ADMIN, idOnTheSource: 'ext-99' });
+      const group = await Group.create({
+        name: 'Entra Team',
+        source: 'entra',
+        idOnTheSource: 'grp-99',
+        memberIds: ['ext-99'],
+      });
+
+      const principals = await methods.getUserPrincipals({ userId: user._id.toString() });
+
+      const rolePrincipal = principals.find((p) => p.principalType === PrincipalType.ROLE);
+      expect(rolePrincipal!.principalId).toBe(SystemRoles.ADMIN);
+      const groupPrincipal = principals.find((p) => p.principalType === PrincipalType.GROUP);
+      expect(groupPrincipal!.principalId!.toString()).toBe(group._id.toString());
+    });
+
+    it('returns only the group id from the projected group query', async () => {
+      const user = await createTestUser({ idOnTheSource: 'ext-proj' });
+      await Group.create({
+        name: 'Projected Team',
+        source: 'entra',
+        idOnTheSource: 'grp-proj',
+        description: 'should not leak',
+        memberIds: ['ext-proj'],
+      });
+
+      const principals = await methods.getUserPrincipals({
+        userId: user._id.toString(),
+        role: SystemRoles.USER,
+        idOnTheSource: 'ext-proj',
+      });
+
+      const groupPrincipal = principals.find((p) => p.principalType === PrincipalType.GROUP);
+      expect(groupPrincipal).toEqual({
+        principalType: PrincipalType.GROUP,
+        principalId: expect.anything(),
+      });
+      expect(Object.keys(groupPrincipal!)).toEqual(['principalType', 'principalId']);
+    });
   });
 
   describe('syncUserEntraGroups', () => {

packages/data-schemas/src/methods/userGroup.ts

@@ -69,6 +69,7 @@ export function createUserGroupMethods(mongoose: typeof import('mongoose')): {
     params: {
       userId: string | Types.ObjectId;
       role?: string | null;
+      idOnTheSource?: string | null;
     },
     session?: ClientSession,
   ) => Promise<Array<{ principalType: PrincipalType; principalId?: string | Types.ObjectId }>>;
@@ -376,8 +377,8 @@ export function createUserGroupMethods(mongoose: typeof import('mongoose')): {
    * Tenant filtering for group memberships is handled automatically by the
    * `applyTenantIsolation` Mongoose plugin on the Group schema. The
    * `tenantContextMiddleware` (chained by `requireJwtAuth` after passport auth)
-   * sets the ALS context, so `getUserGroups()` → `findGroupsByMemberId()` queries
-   * are scoped to the requesting tenant. No explicit tenantId parameter is needed.
+   * sets the ALS context, so the `memberIds` group query below is scoped to the
+   * requesting tenant. No explicit tenantId parameter is needed.
    *
    * IMPORTANT: This relies on the ALS tenant context being active. If this
    * function is called outside a request context (e.g. startup, background jobs),
@@ -386,9 +387,15 @@ export function createUserGroupMethods(mongoose: typeof import('mongoose')): {
    *
    * Ref: #12091 (resolved by tenant context middleware in requireJwtAuth)
    *
+   * Pass `role` and `idOnTheSource` from the already-loaded request user to skip
+   * the fallback user lookup entirely, reducing the hot path to a single indexed,
+   * `_id`-projected group query. `idOnTheSource: null` means "known to be absent"
+   * (local user) and also avoids the lookup; only `undefined` triggers it.
+   *
    * @param params - Parameters object
    * @param params.userId - The user ID
-   * @param params.role - Optional user role (if not provided, will query from DB)
+   * @param params.role - Optional user role (looked up when `undefined`)
+   * @param params.idOnTheSource - Optional external member id (looked up when `undefined`)
    * @param session - Optional MongoDB session for transactions
    * @returns Array of principal objects with type and id
    */
@@ -396,10 +403,11 @@ export function createUserGroupMethods(mongoose: typeof import('mongoose')): {
     params: {
       userId: string | Types.ObjectId;
       role?: string | null;
+      idOnTheSource?: string | null;
     },
     session?: ClientSession,
   ): Promise<Array<{ principalType: PrincipalType; principalId?: string | Types.ObjectId }>> {
-    const { userId, role } = params;
+    const { userId, role, idOnTheSource } = params;
     /** `userId` must be an `ObjectId` for USER principal since ACL entries store `ObjectId`s */
     const userObjectId = typeof userId === 'string' ? new Types.ObjectId(userId) : userId;
     const principals: Array<{
@@ -407,28 +415,39 @@ export function createUserGroupMethods(mongoose: typeof import('mongoose')): {
       principalId?: string | Types.ObjectId;
     }> = [{ principalType: PrincipalType.USER, principalId: userObjectId }];
 
-    // If role is not provided, query user to get it
     let userRole = role;
-    if (userRole === undefined) {
+    let memberIdOnTheSource = idOnTheSource;
+
+    /** Single fallback lookup, only for whichever identity fields the caller omitted. */
+    if (userRole === undefined || memberIdOnTheSource === undefined) {
       const User = mongoose.models.User as Model<IUser>;
-      const query = User.findById(userId).select('role');
+      const query = User.findById(userId).select('role idOnTheSource');
       if (session) {
         query.session(session);
       }
-      const user = await query.lean<IUser>();
-      userRole = user?.role;
+      const user = await query.lean<Pick<IUser, 'role' | 'idOnTheSource'>>();
+      if (userRole === undefined) {
+        userRole = user?.role;
+      }
+      if (memberIdOnTheSource === undefined) {
+        memberIdOnTheSource = user?.idOnTheSource ?? null;
+      }
     }
 
-    // Add role as a principal if user has one
     if (userRole && userRole.trim()) {
       principals.push({ principalType: PrincipalType.ROLE, principalId: userRole });
     }
 
-    const userGroups = await getUserGroups(userId, session);
-    if (userGroups && userGroups.length > 0) {
-      userGroups.forEach((group) => {
-        principals.push({ principalType: PrincipalType.GROUP, principalId: group._id });
-      });
+    /** `memberIds` stores `idOnTheSource` for external users, else the raw user id. */
+    const memberId = memberIdOnTheSource || userId.toString();
+    const Group = mongoose.models.Group as Model<IGroup>;
+    const groupsQuery = Group.find({ memberIds: memberId }, { _id: 1 });
+    if (session) {
+      groupsQuery.session(session);
+    }
+    const userGroups = await groupsQuery.lean<Array<Pick<IGroup, '_id'>>>();
+    for (const group of userGroups) {
+      principals.push({ principalType: PrincipalType.GROUP, principalId: group._id });
     }
 
     principals.push({ principalType: PrincipalType.PUBLIC });