🔧 refactor: Centralize Collection Checks for Permissions Migration (#9565)

* 🔧 refactor: Centralize Collection Existence Checks for Permissions Migration

* Replace individual collection existence checks with a unified function `ensureRequiredCollectionsExist` in the database utility module.
* Update migration scripts for agents and prompts to utilize the new function, ensuring all required collections are verified for existence in a single call.
* Remove redundant collection existence logic from migration files, improving code maintainability and clarity.

* chore: import order in migration scripts

* 🔧 test: Update Token Test Cases for Realistic Scenarios

* Changed email in test data to 'user1-alt@example.com' for a more realistic scenario.
* Clarified expectation comment for token retrieval to indicate it finds the only matching token based on criteria.

config/migrate-agent-permissions.js

@@ -1,5 +1,6 @@
 const path = require('path');
 const { logger } = require('@librechat/data-schemas');
+const { ensureRequiredCollectionsExist } = require('@librechat/api');
 const { AccessRoleIds, ResourceType, PrincipalType } = require('librechat-data-provider');
 const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants;
 
@@ -16,36 +17,11 @@ async function migrateAgentPermissionsEnhanced({ dryRun = true, batchSize = 100
 
   logger.info('Starting Enhanced Agent Permissions Migration', { dryRun, batchSize });
 
-  /** Ensurse `aclentries` collection exists for DocumentDB compatibility
-   * @param {import('mongoose').mongo.Db} db
-   * @param {string} collectionName
-   */
-  async function ensureCollectionExists(db, collectionName) {
-    try {
-      const collections = await db.listCollections({ name: collectionName }).toArray();
-      if (collections.length === 0) {
-        await db.createCollection(collectionName);
-        logger.info(`Created collection: ${collectionName}`);
-      } else {
-        logger.info(`Collection already exists: ${collectionName}`);
-      }
-    } catch (error) {
-      logger.error(`'Failed to check/create "${collectionName}" collection:`, error);
-      // If listCollections fails, try alternative approach
-      try {
-        // Try to access the collection directly - this will create it in MongoDB if it doesn't exist
-        await db.collection(collectionName).findOne({}, { projection: { _id: 1 } });
-      } catch (createError) {
-        logger.error(`Could not ensure collection ${collectionName} exists:`, createError);
-      }
-    }
-  }
-
   const mongoose = require('mongoose');
   /** @type {import('mongoose').mongo.Db | undefined} */
   const db = mongoose.connection.db;
   if (db) {
-    await ensureCollectionExists(db, 'aclentries');
+    await ensureRequiredCollectionsExist(db);
   }
 
   // Verify required roles exist

config/migrate-prompt-permissions.js

@@ -1,5 +1,6 @@
 const path = require('path');
 const { logger } = require('@librechat/data-schemas');
+const { ensureRequiredCollectionsExist } = require('@librechat/api');
 const { AccessRoleIds, ResourceType, PrincipalType } = require('librechat-data-provider');
 const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants;
 
@@ -16,36 +17,11 @@ async function migrateToPromptGroupPermissions({ dryRun = true, batchSize = 100
 
   logger.info('Starting PromptGroup Permissions Migration', { dryRun, batchSize });
 
-  /** Ensurse `aclentries` collection exists for DocumentDB compatibility
-   * @param {import('mongoose').mongo.Db} db
-   * @param {string} collectionName
-   */
-  async function ensureCollectionExists(db, collectionName) {
-    try {
-      const collections = await db.listCollections({ name: collectionName }).toArray();
-      if (collections.length === 0) {
-        await db.createCollection(collectionName);
-        logger.info(`Created collection: ${collectionName}`);
-      } else {
-        logger.info(`Collection already exists: ${collectionName}`);
-      }
-    } catch (error) {
-      logger.error(`'Failed to check/create "${collectionName}" collection:`, error);
-      // If listCollections fails, try alternative approach
-      try {
-        // Try to access the collection directly - this will create it in MongoDB if it doesn't exist
-        await db.collection(collectionName).findOne({}, { projection: { _id: 1 } });
-      } catch (createError) {
-        logger.error(`Could not ensure collection ${collectionName} exists:`, createError);
-      }
-    }
-  }
-
   const mongoose = require('mongoose');
   /** @type {import('mongoose').mongo.Db | undefined} */
   const db = mongoose.connection.db;
   if (db) {
-    await ensureCollectionExists(db, 'aclentries');
+    await ensureRequiredCollectionsExist(db);
   }
 
   // Verify required roles exist

packages/api/src/agents/migration.ts

@@ -1,7 +1,8 @@
 import { logger } from '@librechat/data-schemas';
 import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider';
+import { ensureRequiredCollectionsExist } from '../db/utils';
 import type { AccessRoleMethods, IAgent } from '@librechat/data-schemas';
-import type { Model, Mongoose, mongo } from 'mongoose';
+import type { Model, Mongoose } from 'mongoose';
 
 const { GLOBAL_PROJECT_NAME } = Constants;
 
@@ -54,31 +55,9 @@ export async function checkAgentPermissionsMigration({
   logger.debug('Checking if agent permissions migration is needed');
 
   try {
-    /** Ensurse `aclentries` collection exists for DocumentDB compatibility */
-    async function ensureCollectionExists(db: mongo.Db, collectionName: string) {
-      try {
-        const collections = await db.listCollections({ name: collectionName }).toArray();
-        if (collections.length === 0) {
-          await db.createCollection(collectionName);
-          logger.info(`Created collection: ${collectionName}`);
-        } else {
-          logger.debug(`Collection already exists: ${collectionName}`);
-        }
-      } catch (error) {
-        logger.error(`'Failed to check/create "${collectionName}" collection:`, error);
-        // If listCollections fails, try alternative approach
-        try {
-          // Try to access the collection directly - this will create it in MongoDB if it doesn't exist
-          await db.collection(collectionName).findOne({}, { projection: { _id: 1 } });
-        } catch (createError) {
-          logger.error(`Could not ensure collection ${collectionName} exists:`, createError);
-        }
-      }
-    }
-
     const db = mongoose.connection.db;
     if (db) {
-      await ensureCollectionExists(db, 'aclentries');
+      await ensureRequiredCollectionsExist(db);
     }
 
     // Verify required roles exist

packages/api/src/db/utils.ts

@@ -0,0 +1,54 @@
+import { logger } from '@librechat/data-schemas';
+import type { mongo } from 'mongoose';
+
+/**
+ * Ensures that a collection exists in the database.
+ * For DocumentDB compatibility, it tries multiple approaches.
+ * @param db - The MongoDB database instance
+ * @param collectionName - The name of the collection to ensure exists
+ */
+export async function ensureCollectionExists(db: mongo.Db, collectionName: string): Promise<void> {
+  try {
+    const collections = await db.listCollections({ name: collectionName }).toArray();
+    if (collections.length === 0) {
+      await db.createCollection(collectionName);
+      logger.info(`Created collection: ${collectionName}`);
+    } else {
+      logger.debug(`Collection already exists: ${collectionName}`);
+    }
+  } catch (error) {
+    logger.error(`Failed to check/create "${collectionName}" collection:`, error);
+    // If listCollections fails, try alternative approach
+    try {
+      // Try to access the collection directly - this will create it in MongoDB if it doesn't exist
+      await db.collection(collectionName).findOne({}, { projection: { _id: 1 } });
+    } catch (createError) {
+      logger.error(`Could not ensure collection ${collectionName} exists:`, createError);
+    }
+  }
+}
+
+/**
+ * Ensures that all required collections exist for the permission system.
+ * This includes aclentries, groups, accessroles, and any other collections
+ * needed for migrations and permission checks.
+ * @param db - The MongoDB database instance
+ */
+export async function ensureRequiredCollectionsExist(db: mongo.Db): Promise<void> {
+  const requiredCollections = [
+    'aclentries', // ACL permission entries
+    'groups', // User groups
+    'accessroles', // Access roles for permissions
+    'agents', // Agents collection
+    'promptgroups', // Prompt groups collection
+    'projects', // Projects collection
+  ];
+
+  logger.debug('Ensuring required collections exist for permission system');
+
+  for (const collectionName of requiredCollections) {
+    await ensureCollectionExists(db, collectionName);
+  }
+
+  logger.debug('All required collections have been checked/created');
+}

packages/api/src/index.ts

@@ -11,6 +11,7 @@ export * from './mcp/zod';
 export * from './format';
 export * from './mcp/utils';
 export * from './utils';
+export * from './db/utils';
 /* OAuth */
 export * from './oauth';
 /* Crypto */

packages/api/src/prompts/migration.ts

@@ -1,7 +1,8 @@
 import { logger } from '@librechat/data-schemas';
 import { AccessRoleIds, ResourceType, PrincipalType, Constants } from 'librechat-data-provider';
+import { ensureRequiredCollectionsExist } from '../db/utils';
 import type { AccessRoleMethods, IPromptGroupDocument } from '@librechat/data-schemas';
-import type { Model, Mongoose, mongo } from 'mongoose';
+import type { Model, Mongoose } from 'mongoose';
 
 const { GLOBAL_PROJECT_NAME } = Constants;
 
@@ -52,32 +53,10 @@ export async function checkPromptPermissionsMigration({
   logger.debug('Checking if prompt permissions migration is needed');
 
   try {
-    /** Ensurse `aclentries` collection exists for DocumentDB compatibility */
-    async function ensureCollectionExists(db: mongo.Db, collectionName: string) {
-      try {
-        const collections = await db.listCollections({ name: collectionName }).toArray();
-        if (collections.length === 0) {
-          await db.createCollection(collectionName);
-          logger.info(`Created collection: ${collectionName}`);
-        } else {
-          logger.debug(`Collection already exists: ${collectionName}`);
-        }
-      } catch (error) {
-        logger.error(`'Failed to check/create "${collectionName}" collection:`, error);
-        // If listCollections fails, try alternative approach
-        try {
-          // Try to access the collection directly - this will create it in MongoDB if it doesn't exist
-          await db.collection(collectionName).findOne({}, { projection: { _id: 1 } });
-        } catch (createError) {
-          logger.error(`Could not ensure collection ${collectionName} exists:`, createError);
-        }
-      }
-    }
-
     /** Native MongoDB database instance */
     const db = mongoose.connection.db;
     if (db) {
-      await ensureCollectionExists(db, 'aclentries');
+      await ensureRequiredCollectionsExist(db);
     }
 
     // Verify required roles exist

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

@@ -118,7 +118,7 @@ describe('Token Methods - Detailed Tests', () => {
         {
           token: 'token-3',
           userId: user1Id,
-          email: 'user1@example.com',
+          email: 'user1-alt@example.com', // Different email for realistic scenario
           createdAt: new Date(),
           expiresAt: new Date(Date.now() + 3600000),
         },
@@ -164,7 +164,7 @@ describe('Token Methods - Detailed Tests', () => {
       });
 
       expect(found).toBeDefined();
-      expect(found?.token).toBe('token-1'); // Should find first matching
+      expect(found?.token).toBe('token-1'); // Should find the only token matching both criteria
     });
 
     test('should return null for non-existent token', async () => {