📦 refactor: Consolidate DB models, encapsulating Mongoose usage in data-schemas (#11830)

* chore: move database model methods to /packages/data-schemas

* chore: add TypeScript ESLint rule to warn on unused variables

* refactor: model imports to streamline access

- Consolidated model imports across various files to improve code organization and reduce redundancy.
- Updated imports for models such as Assistant, Message, Conversation, and others to a unified import path.
- Adjusted middleware and service files to reflect the new import structure, ensuring functionality remains intact.
- Enhanced test files to align with the new import paths, maintaining test coverage and integrity.

* chore: migrate database models to packages/data-schemas and refactor all direct Mongoose Model usage outside of data-schemas

* test: update agent model mocks in unit tests

- Added `getAgent` mock to `client.test.js` to enhance test coverage for agent-related functionality.
- Removed redundant `getAgent` and `getAgents` mocks from `openai.spec.js` and `responses.unit.spec.js` to streamline test setup and reduce duplication.
- Ensured consistency in agent mock implementations across test files.

* fix: update types in data-schemas

* refactor: enhance type definitions in transaction and spending methods

- Updated type definitions in `checkBalance.ts` to use specific request and response types.
- Refined `spendTokens.ts` to utilize a new `SpendTxData` interface for better clarity and type safety.
- Improved transaction handling in `transaction.ts` by introducing `TransactionResult` and `TxData` interfaces, ensuring consistent data structures across methods.
- Adjusted unit tests in `transaction.spec.ts` to accommodate new type definitions and enhance robustness.

* refactor: streamline model imports and enhance code organization

- Consolidated model imports across various controllers and services to a unified import path, improving code clarity and reducing redundancy.
- Updated multiple files to reflect the new import structure, ensuring all functionalities remain intact.
- Enhanced overall code organization by removing duplicate import statements and optimizing the usage of model methods.

* feat: implement loadAddedAgent and refactor agent loading logic

- Introduced `loadAddedAgent` function to handle loading agents from added conversations, supporting multi-convo parallel execution.
- Created a new `load.ts` file to encapsulate agent loading functionalities, including `loadEphemeralAgent` and `loadAgent`.
- Updated the `index.ts` file to export the new `load` module instead of the deprecated `loadAgent`.
- Enhanced type definitions and improved error handling in the agent loading process.
- Adjusted unit tests to reflect changes in the agent loading structure and ensure comprehensive coverage.

* refactor: enhance balance handling with new update interface

- Introduced `IBalanceUpdate` interface to streamline balance update operations across the codebase.
- Updated `upsertBalanceFields` method signatures in `balance.ts`, `transaction.ts`, and related tests to utilize the new interface for improved type safety.
- Adjusted type imports in `balance.spec.ts` to include `IBalanceUpdate`, ensuring consistency in balance management functionalities.
- Enhanced overall code clarity and maintainability by refining type definitions related to balance operations.

* feat: add unit tests for loadAgent functionality and enhance agent loading logic

- Introduced comprehensive unit tests for the `loadAgent` function, covering various scenarios including null and empty agent IDs, loading of ephemeral agents, and permission checks.
- Enhanced the `initializeClient` function by moving `getConvoFiles` to the correct position in the database method exports, ensuring proper functionality.
- Improved test coverage for agent loading, including handling of non-existent agents and user permissions.

* chore: reorder memory method exports for consistency

- Moved `deleteAllUserMemories` to the correct position in the exported memory methods, ensuring a consistent and logical order of method exports in `memory.ts`.
This commit is contained in:
Danny Avila 2026-02-17 18:23:44 -05:00
parent a85e99ff45
commit a6fb257bcf
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
182 changed files with 8548 additions and 8105 deletions

View file

@ -1,24 +1,22 @@
const { Balance } = require('~/db/models');
const { findBalanceByUser } = require('~/models');
async function balanceController(req, res) {
const balanceData = await Balance.findOne(
{ user: req.user.id },
'-_id tokenCredits autoRefillEnabled refillIntervalValue refillIntervalUnit lastRefill refillAmount',
).lean();
const balanceData = await findBalanceByUser(req.user.id);
if (!balanceData) {
return res.status(404).json({ error: 'Balance not found' });
}
// If auto-refill is not enabled, remove auto-refill related fields from the response
if (!balanceData.autoRefillEnabled) {
delete balanceData.refillIntervalValue;
delete balanceData.refillIntervalUnit;
delete balanceData.lastRefill;
delete balanceData.refillAmount;
const { _id: _, ...result } = balanceData;
if (!result.autoRefillEnabled) {
delete result.refillIntervalValue;
delete result.refillIntervalUnit;
delete result.lastRefill;
delete result.refillAmount;
}
res.status(200).json(balanceData);
res.status(200).json(result);
}
module.exports = balanceController;

View file

@ -9,22 +9,24 @@ const { enrichRemoteAgentPrincipals, backfillRemoteAgentPermissions } = require(
const {
bulkUpdateResourcePermissions,
ensureGroupPrincipalExists,
getResourcePermissionsMap,
findAccessibleResources,
getEffectivePermissions,
ensurePrincipalExists,
getAvailableRoles,
findAccessibleResources,
getResourcePermissionsMap,
} = require('~/server/services/PermissionService');
const {
searchPrincipals: searchLocalPrincipals,
sortPrincipalsByRelevance,
calculateRelevanceScore,
findRoleByIdentifier,
aggregateAclEntries,
bulkWriteAclEntries,
} = require('~/models');
const {
entraIdPrincipalFeatureEnabled,
searchEntraIdPrincipals,
} = require('~/server/services/GraphApiService');
const { AclEntry, AccessRole } = require('~/db/models');
/**
* Generic controller for resource permission endpoints
@ -185,8 +187,7 @@ const getResourcePermissions = async (req, res) => {
const { resourceType, resourceId } = req.params;
validateResourceType(resourceType);
// Use aggregation pipeline for efficient single-query data retrieval
const results = await AclEntry.aggregate([
const results = await aggregateAclEntries([
// Match ACL entries for this resource
{
$match: {
@ -282,7 +283,12 @@ const getResourcePermissions = async (req, res) => {
}
if (resourceType === ResourceType.REMOTE_AGENT) {
const enricherDeps = { AclEntry, AccessRole, logger };
const enricherDeps = {
aggregateAclEntries,
bulkWriteAclEntries,
findRoleByIdentifier,
logger,
};
const enrichResult = await enrichRemoteAgentPrincipals(enricherDeps, resourceId, principals);
principals = enrichResult.principals;
backfillRemoteAgentPermissions(enricherDeps, resourceId, enrichResult.entriesToBackfill);

View file

@ -6,33 +6,6 @@ const {
normalizeHttpError,
extractWebSearchEnvVars,
} = require('@librechat/api');
const {
deleteAllUserSessions,
deleteAllSharedLinks,
updateUserPlugins,
deleteUserById,
deleteMessages,
deletePresets,
deleteUserKey,
deleteConvos,
deleteFiles,
updateUser,
findToken,
getFiles,
} = require('~/models');
const {
ConversationTag,
AgentApiKey,
Transaction,
MemoryEntry,
Assistant,
AclEntry,
Balance,
Action,
Group,
Token,
User,
} = require('~/db/models');
const { updateUserPluginAuth, deleteUserPluginAuth } = require('~/server/services/PluginService');
const { verifyEmail, resendVerificationEmail } = require('~/server/services/AuthService');
const { getMCPManager, getFlowStateManager, getMCPServersRegistry } = require('~/config');
@ -40,10 +13,8 @@ const { invalidateCachedTools } = require('~/server/services/Config/getCachedToo
const { needsRefresh, getNewS3URL } = require('~/server/services/Files/S3/crud');
const { processDeleteRequest } = require('~/server/services/Files/process');
const { getAppConfig } = require('~/server/services/Config');
const { deleteToolCalls } = require('~/models/ToolCall');
const { deleteUserPrompts } = require('~/models/Prompt');
const { deleteUserAgents } = require('~/models/Agent');
const { getLogStores } = require('~/cache');
const db = require('~/models');
const getUserController = async (req, res) => {
const appConfig = await getAppConfig({ role: req.user?.role });
@ -64,7 +35,7 @@ const getUserController = async (req, res) => {
const originalAvatar = userData.avatar;
try {
userData.avatar = await getNewS3URL(userData.avatar);
await updateUser(userData.id, { avatar: userData.avatar });
await db.updateUser(userData.id, { avatar: userData.avatar });
} catch (error) {
userData.avatar = originalAvatar;
logger.error('Error getting new S3 URL for avatar:', error);
@ -75,7 +46,7 @@ const getUserController = async (req, res) => {
const getTermsStatusController = async (req, res) => {
try {
const user = await User.findById(req.user.id);
const user = await db.getUserById(req.user.id, 'termsAccepted');
if (!user) {
return res.status(404).json({ message: 'User not found' });
}
@ -88,7 +59,7 @@ const getTermsStatusController = async (req, res) => {
const acceptTermsController = async (req, res) => {
try {
const user = await User.findByIdAndUpdate(req.user.id, { termsAccepted: true }, { new: true });
const user = await db.updateUser(req.user.id, { termsAccepted: true });
if (!user) {
return res.status(404).json({ message: 'User not found' });
}
@ -101,7 +72,7 @@ const acceptTermsController = async (req, res) => {
const deleteUserFiles = async (req) => {
try {
const userFiles = await getFiles({ user: req.user.id });
const userFiles = await db.getFiles({ user: req.user.id });
await processDeleteRequest({
req,
files: userFiles,
@ -117,7 +88,7 @@ const updateUserPluginsController = async (req, res) => {
const { pluginKey, action, auth, isEntityTool } = req.body;
try {
if (!isEntityTool) {
await updateUserPlugins(user._id, user.plugins, pluginKey, action);
await db.updateUserPlugins(user._id, user.plugins, pluginKey, action);
}
if (auth == null) {
@ -241,33 +212,33 @@ const deleteUserController = async (req, res) => {
const { user } = req;
try {
await deleteMessages({ user: user.id }); // delete user messages
await deleteAllUserSessions({ userId: user.id }); // delete user sessions
await Transaction.deleteMany({ user: user.id }); // delete user transactions
await deleteUserKey({ userId: user.id, all: true }); // delete user keys
await Balance.deleteMany({ user: user._id }); // delete user balances
await deletePresets(user.id); // delete user presets
await db.deleteMessages({ user: user.id });
await db.deleteAllUserSessions({ userId: user.id });
await db.deleteTransactions({ user: user.id });
await db.deleteUserKey({ userId: user.id, all: true });
await db.deleteBalances({ user: user._id });
await db.deletePresets(user.id);
try {
await deleteConvos(user.id); // delete user convos
await db.deleteConvos(user.id);
} catch (error) {
logger.error('[deleteUserController] Error deleting user convos, likely no convos', error);
}
await deleteUserPluginAuth(user.id, null, true); // delete user plugin auth
await deleteUserById(user.id); // delete user
await deleteAllSharedLinks(user.id); // delete user shared links
await deleteUserFiles(req); // delete user files
await deleteFiles(null, user.id); // delete database files in case of orphaned files from previous steps
await deleteToolCalls(user.id); // delete user tool calls
await deleteUserAgents(user.id); // delete user agents
await AgentApiKey.deleteMany({ user: user._id }); // delete user agent API keys
await Assistant.deleteMany({ user: user.id }); // delete user assistants
await ConversationTag.deleteMany({ user: user.id }); // delete user conversation tags
await MemoryEntry.deleteMany({ userId: user.id }); // delete user memory entries
await deleteUserPrompts(req, user.id); // delete user prompts
await Action.deleteMany({ user: user.id }); // delete user actions
await Token.deleteMany({ userId: user.id }); // delete user OAuth tokens
await Group.updateMany({ memberIds: user.id }, { $pullAll: { memberIds: [user.id] } });
await AclEntry.deleteMany({ principalId: user._id }); // delete user ACL entries
await deleteUserPluginAuth(user.id, null, true);
await db.deleteUserById(user.id);
await db.deleteAllSharedLinks(user.id);
await deleteUserFiles(req);
await db.deleteFiles(null, user.id);
await db.deleteToolCalls(user.id);
await db.deleteUserAgents(user.id);
await db.deleteAllAgentApiKeys(user._id);
await db.deleteAssistants({ user: user.id });
await db.deleteConversationTags({ user: user.id });
await db.deleteAllUserMemories(user.id);
await db.deleteUserPrompts(user.id);
await db.deleteActions({ user: user.id });
await db.deleteTokens({ userId: user.id });
await db.removeUserFromAllGroups(user.id);
await db.deleteAclEntries({ principalId: user._id });
logger.info(`User deleted account. Email: ${user.email} ID: ${user.id}`);
res.status(200).send({ message: 'User deleted' });
} catch (err) {
@ -327,7 +298,7 @@ const maybeUninstallOAuthMCP = async (userId, pluginKey, appConfig) => {
const clientTokenData = await MCPTokenStorage.getClientInfoAndMetadata({
userId,
serverName,
findToken,
findToken: db.findToken,
});
if (clientTokenData == null) {
return;
@ -338,7 +309,7 @@ const maybeUninstallOAuthMCP = async (userId, pluginKey, appConfig) => {
const tokens = await MCPTokenStorage.getTokens({
userId,
serverName,
findToken,
findToken: db.findToken,
});
// 3. revoke OAuth tokens at the provider
@ -394,7 +365,7 @@ const maybeUninstallOAuthMCP = async (userId, pluginKey, appConfig) => {
userId,
serverName,
deleteToken: async (filter) => {
await Token.deleteOne(filter);
await db.deleteTokens(filter);
},
});

View file

@ -14,20 +14,40 @@ jest.mock('@librechat/data-schemas', () => {
};
});
jest.mock('~/models', () => ({
deleteAllUserSessions: jest.fn().mockResolvedValue(undefined),
deleteAllSharedLinks: jest.fn().mockResolvedValue(undefined),
updateUserPlugins: jest.fn(),
deleteUserById: jest.fn().mockResolvedValue(undefined),
deleteMessages: jest.fn().mockResolvedValue(undefined),
deletePresets: jest.fn().mockResolvedValue(undefined),
deleteUserKey: jest.fn().mockResolvedValue(undefined),
deleteConvos: jest.fn().mockResolvedValue(undefined),
deleteFiles: jest.fn().mockResolvedValue(undefined),
updateUser: jest.fn(),
findToken: jest.fn(),
getFiles: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models', () => {
const _mongoose = require('mongoose');
return {
deleteAllUserSessions: jest.fn().mockResolvedValue(undefined),
deleteAllSharedLinks: jest.fn().mockResolvedValue(undefined),
deleteAllAgentApiKeys: jest.fn().mockResolvedValue(undefined),
deleteConversationTags: jest.fn().mockResolvedValue(undefined),
deleteAllUserMemories: jest.fn().mockResolvedValue(undefined),
deleteTransactions: jest.fn().mockResolvedValue(undefined),
deleteAclEntries: jest.fn().mockResolvedValue(undefined),
updateUserPlugins: jest.fn(),
deleteAssistants: jest.fn().mockResolvedValue(undefined),
deleteUserById: jest.fn().mockResolvedValue(undefined),
deleteUserPrompts: jest.fn().mockResolvedValue(undefined),
deleteMessages: jest.fn().mockResolvedValue(undefined),
deleteBalances: jest.fn().mockResolvedValue(undefined),
deleteActions: jest.fn().mockResolvedValue(undefined),
deletePresets: jest.fn().mockResolvedValue(undefined),
deleteUserKey: jest.fn().mockResolvedValue(undefined),
deleteToolCalls: jest.fn().mockResolvedValue(undefined),
deleteUserAgents: jest.fn().mockResolvedValue(undefined),
deleteTokens: jest.fn().mockResolvedValue(undefined),
deleteConvos: jest.fn().mockResolvedValue(undefined),
deleteFiles: jest.fn().mockResolvedValue(undefined),
updateUser: jest.fn(),
getUserById: jest.fn().mockResolvedValue(null),
findToken: jest.fn(),
getFiles: jest.fn().mockResolvedValue([]),
removeUserFromAllGroups: jest.fn().mockImplementation(async (userId) => {
const Group = _mongoose.models.Group;
await Group.updateMany({ memberIds: userId }, { $pullAll: { memberIds: [userId] } });
}),
};
});
jest.mock('~/server/services/PluginService', () => ({
updateUserPluginAuth: jest.fn(),
@ -55,18 +75,6 @@ jest.mock('~/server/services/Config', () => ({
getMCPServersRegistry: jest.fn(),
}));
jest.mock('~/models/ToolCall', () => ({
deleteToolCalls: jest.fn().mockResolvedValue(undefined),
}));
jest.mock('~/models/Prompt', () => ({
deleteUserPrompts: jest.fn().mockResolvedValue(undefined),
}));
jest.mock('~/models/Agent', () => ({
deleteUserAgents: jest.fn().mockResolvedValue(undefined),
}));
jest.mock('~/cache', () => ({
getLogStores: jest.fn(),
}));

View file

@ -77,11 +77,6 @@ jest.mock('~/server/services/ToolService', () => ({
loadToolsForExecution: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/spendTokens', () => ({
spendTokens: mockSpendTokens,
spendStructuredTokens: mockSpendStructuredTokens,
}));
jest.mock('~/server/controllers/agents/callbacks', () => ({
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
}));
@ -90,20 +85,8 @@ jest.mock('~/server/services/PermissionService', () => ({
findAccessibleResources: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/Conversation', () => ({
getConvoFiles: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/Agent', () => ({
getAgent: jest.fn().mockResolvedValue({
id: 'agent-123',
provider: 'openAI',
model_parameters: { model: 'gpt-4' },
}),
getAgents: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models', () => ({
getAgent: jest.fn().mockResolvedValue({ id: 'agent-123', name: 'Test Agent' }),
getFiles: jest.fn(),
getUserKey: jest.fn(),
getMessages: jest.fn(),
@ -112,6 +95,9 @@ jest.mock('~/models', () => ({
getUserCodeFiles: jest.fn(),
getToolFilesByIds: jest.fn(),
getCodeGeneratedFiles: jest.fn(),
spendTokens: mockSpendTokens,
spendStructuredTokens: mockSpendStructuredTokens,
getConvoFiles: jest.fn().mockResolvedValue([]),
}));
describe('OpenAIChatCompletionController', () => {

View file

@ -101,11 +101,6 @@ jest.mock('~/server/services/ToolService', () => ({
loadToolsForExecution: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/spendTokens', () => ({
spendTokens: mockSpendTokens,
spendStructuredTokens: mockSpendStructuredTokens,
}));
jest.mock('~/server/controllers/agents/callbacks', () => ({
createToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
createResponsesToolEndCallback: jest.fn().mockReturnValue(jest.fn()),
@ -115,23 +110,8 @@ jest.mock('~/server/services/PermissionService', () => ({
findAccessibleResources: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/Conversation', () => ({
getConvoFiles: jest.fn().mockResolvedValue([]),
saveConvo: jest.fn().mockResolvedValue({}),
getConvo: jest.fn().mockResolvedValue(null),
}));
jest.mock('~/models/Agent', () => ({
getAgent: jest.fn().mockResolvedValue({
id: 'agent-123',
name: 'Test Agent',
provider: 'anthropic',
model_parameters: { model: 'claude-3' },
}),
getAgents: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models', () => ({
getAgent: jest.fn().mockResolvedValue({ id: 'agent-123', name: 'Test Agent' }),
getFiles: jest.fn(),
getUserKey: jest.fn(),
getMessages: jest.fn().mockResolvedValue([]),
@ -141,6 +121,11 @@ jest.mock('~/models', () => ({
getUserCodeFiles: jest.fn(),
getToolFilesByIds: jest.fn(),
getCodeGeneratedFiles: jest.fn(),
spendTokens: mockSpendTokens,
spendStructuredTokens: mockSpendStructuredTokens,
getConvoFiles: jest.fn().mockResolvedValue([]),
saveConvo: jest.fn().mockResolvedValue({}),
getConvo: jest.fn().mockResolvedValue(null),
}));
describe('createResponse controller', () => {

View file

@ -1,10 +1,8 @@
const { duplicateAgent } = require('../v1');
const { getAgent, createAgent } = require('~/models/Agent');
const { getActions } = require('~/models/Action');
const { getAgent, createAgent, getActions } = require('~/models');
const { nanoid } = require('nanoid');
jest.mock('~/models/Agent');
jest.mock('~/models/Action');
jest.mock('~/models');
jest.mock('nanoid');
describe('duplicateAgent', () => {

View file

@ -21,6 +21,7 @@ const {
GenerationJobManager,
getTransactionsConfig,
createMemoryProcessor,
loadAgent: loadAgentFn,
createMultiAgentMapper,
filterMalformedContentParts,
} = require('@librechat/api');
@ -43,16 +44,15 @@ const {
isEphemeralAgentId,
removeNullishValues,
} = require('librechat-data-provider');
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
const { encodeAndFormat } = require('~/server/services/Files/images/encode');
const { createContextHandlers } = require('~/app/clients/prompts');
const { getConvoFiles } = require('~/models/Conversation');
const { getMCPServerTools } = require('~/server/services/Config');
const BaseClient = require('~/app/clients/BaseClient');
const { getRoleByName } = require('~/models/Role');
const { loadAgent } = require('~/models/Agent');
const { getMCPManager } = require('~/config');
const db = require('~/models');
const loadAgent = (params) => loadAgentFn(params, { getAgent: db.getAgent, getMCPServerTools });
class AgentClient extends BaseClient {
constructor(options = {}) {
super(null, options);
@ -409,7 +409,7 @@ class AgentClient extends BaseClient {
user,
permissionType: PermissionTypes.MEMORIES,
permissions: [Permissions.USE],
getRoleByName,
getRoleByName: db.getRoleByName,
});
if (!hasAccess) {
@ -469,9 +469,9 @@ class AgentClient extends BaseClient {
},
},
{
getConvoFiles,
getFiles: db.getFiles,
getUserKey: db.getUserKey,
getConvoFiles: db.getConvoFiles,
updateFilesUsage: db.updateFilesUsage,
getUserKeyValues: db.getUserKeyValues,
getToolFilesByIds: db.getToolFilesByIds,
@ -670,7 +670,7 @@ class AgentClient extends BaseClient {
};
if (cache_creation > 0 || cache_read > 0) {
spendStructuredTokens(txMetadata, {
db.spendStructuredTokens(txMetadata, {
promptTokens: {
input: usage.input_tokens,
write: cache_creation,
@ -685,7 +685,7 @@ class AgentClient extends BaseClient {
});
continue;
}
spendTokens(txMetadata, {
db.spendTokens(txMetadata, {
promptTokens: usage.input_tokens,
completionTokens: usage.output_tokens,
}).catch((err) => {
@ -1180,7 +1180,7 @@ class AgentClient extends BaseClient {
context = 'message',
}) {
try {
await spendTokens(
await db.spendTokens(
{
model,
context,
@ -1198,7 +1198,7 @@ class AgentClient extends BaseClient {
'reasoning_tokens' in usage &&
typeof usage.reasoning_tokens === 'number'
) {
await spendTokens(
await db.spendTokens(
{
model,
balance,

View file

@ -15,13 +15,15 @@ jest.mock('@librechat/api', () => ({
checkAccess: jest.fn(),
initializeAgent: jest.fn(),
createMemoryProcessor: jest.fn(),
}));
jest.mock('~/models/Agent', () => ({
loadAgent: jest.fn(),
}));
jest.mock('~/models/Role', () => ({
jest.mock('~/server/services/Config', () => ({
getMCPServerTools: jest.fn(),
}));
jest.mock('~/models', () => ({
getAgent: jest.fn(),
getRoleByName: jest.fn(),
}));
@ -2137,7 +2139,7 @@ describe('AgentClient - titleConvo', () => {
};
mockCheckAccess = require('@librechat/api').checkAccess;
mockLoadAgent = require('~/models/Agent').loadAgent;
mockLoadAgent = require('@librechat/api').loadAgent;
mockInitializeAgent = require('@librechat/api').initializeAgent;
mockCreateMemoryProcessor = require('@librechat/api').createMemoryProcessor;
});
@ -2194,6 +2196,7 @@ describe('AgentClient - titleConvo', () => {
expect.objectContaining({
agent_id: differentAgentId,
}),
expect.any(Object),
);
expect(mockInitializeAgent).toHaveBeenCalledWith(
expect.objectContaining({

View file

@ -3,8 +3,8 @@ const { logger } = require('@librechat/data-schemas');
const { CacheKeys, ViolationTypes } = require('librechat-data-provider');
const { sendResponse } = require('~/server/middleware/error');
const { recordUsage } = require('~/server/services/Threads');
const { getConvo } = require('~/models/Conversation');
const getLogStores = require('~/cache/getLogStores');
const { getConvo } = require('~/models');
/**
* @typedef {Object} ErrorHandlerContext

View file

@ -24,9 +24,6 @@ const {
const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
const { createToolEndCallback } = require('~/server/controllers/agents/callbacks');
const { findAccessibleResources } = require('~/server/services/PermissionService');
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
const { getConvoFiles } = require('~/models/Conversation');
const { getAgent, getAgents } = require('~/models/Agent');
const db = require('~/models');
/**
@ -139,7 +136,7 @@ const OpenAIChatCompletionController = async (req, res) => {
const agentId = request.model;
// Look up the agent
const agent = await getAgent({ id: agentId });
const agent = await db.getAgent({ id: agentId });
if (!agent) {
return sendErrorResponse(
res,
@ -206,7 +203,7 @@ const OpenAIChatCompletionController = async (req, res) => {
isInitialAgent: true,
},
{
getConvoFiles,
getConvoFiles: db.getConvoFiles,
getFiles: db.getFiles,
getUserKey: db.getUserKey,
getMessages: db.getMessages,
@ -490,7 +487,7 @@ const OpenAIChatCompletionController = async (req, res) => {
const balanceConfig = getBalanceConfig(appConfig);
const transactionsConfig = getTransactionsConfig(appConfig);
recordCollectedUsage(
{ spendTokens, spendStructuredTokens },
{ spendTokens: db.spendTokens, spendStructuredTokens: db.spendStructuredTokens },
{
user: userId,
conversationId,
@ -599,7 +596,7 @@ const ListModelsController = async (req, res) => {
// Get the accessible agents
let agents = [];
if (accessibleAgentIds.length > 0) {
agents = await getAgents({ _id: { $in: accessibleAgentIds } });
agents = await db.getAgents({ _id: { $in: accessibleAgentIds } });
}
const models = agents.map((agent) => ({
@ -642,7 +639,7 @@ const GetModelController = async (req, res) => {
return sendErrorResponse(res, 401, 'Authentication required', 'auth_error');
}
const agent = await getAgent({ id: model });
const agent = await db.getAgent({ id: model });
if (!agent) {
return sendErrorResponse(

View file

@ -14,7 +14,7 @@ const { EModelEndpoint } = require('librechat-data-provider');
const mockSpendTokens = jest.fn().mockResolvedValue();
const mockSpendStructuredTokens = jest.fn().mockResolvedValue();
jest.mock('~/models/spendTokens', () => ({
jest.mock('~/models', () => ({
spendTokens: (...args) => mockSpendTokens(...args),
spendStructuredTokens: (...args) => mockSpendStructuredTokens(...args),
}));

View file

@ -131,9 +131,15 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
partialMessage.agent_id = req.body.agent_id;
}
await saveMessage(req, partialMessage, {
context: 'api/server/controllers/agents/request.js - partial response on disconnect',
});
await saveMessage(
{
userId: req?.user?.id,
isTemporary: req?.body?.isTemporary,
interfaceConfig: req?.config?.interfaceConfig,
},
partialMessage,
{ context: 'api/server/controllers/agents/request.js - partial response on disconnect' },
);
logger.debug(
`[ResumableAgentController] Saved partial response for ${streamId}, content parts: ${aggregatedContent.length}`,
@ -274,8 +280,14 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
// Save user message BEFORE sending final event to avoid race condition
// where client refetch happens before database is updated
const reqCtx = {
userId: req?.user?.id,
isTemporary: req?.body?.isTemporary,
interfaceConfig: req?.config?.interfaceConfig,
};
if (!client.skipSaveUserMessage && userMessage) {
await saveMessage(req, userMessage, {
await saveMessage(reqCtx, userMessage, {
context: 'api/server/controllers/agents/request.js - resumable user message',
});
}
@ -285,7 +297,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
// before the response is saved to the database, causing orphaned parentMessageIds.
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
await saveMessage(
req,
reqCtx,
{ ...response, user: userId, unfinished: wasAbortedBeforeComplete },
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
);
@ -668,7 +680,11 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle
// Save the message if needed
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
await saveMessage(
req,
{
userId: req?.user?.id,
isTemporary: req?.body?.isTemporary,
interfaceConfig: req?.config?.interfaceConfig,
},
{ ...finalResponse, user: userId },
{ context: 'api/server/controllers/agents/request.js - response end' },
);
@ -697,9 +713,15 @@ const _LegacyAgentController = async (req, res, next, initializeClient, addTitle
// Save user message if needed
if (!client.skipSaveUserMessage) {
await saveMessage(req, userMessage, {
context: "api/server/controllers/agents/request.js - don't skip saving user message",
});
await saveMessage(
{
userId: req?.user?.id,
isTemporary: req?.body?.isTemporary,
interfaceConfig: req?.config?.interfaceConfig,
},
userMessage,
{ context: "api/server/controllers/agents/request.js - don't skip saving user message" },
);
}
// Add title if needed - extract minimal data

View file

@ -36,9 +36,6 @@ const {
} = require('~/server/controllers/agents/callbacks');
const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService');
const { findAccessibleResources } = require('~/server/services/PermissionService');
const { getConvoFiles, saveConvo, getConvo } = require('~/models/Conversation');
const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
const { getAgent, getAgents } = require('~/models/Agent');
const db = require('~/models');
/** @type {import('@librechat/api').AppConfig | null} */
@ -213,8 +210,12 @@ async function saveResponseOutput(req, conversationId, responseId, response, age
* @returns {Promise<void>}
*/
async function saveConversation(req, conversationId, agentId, agent) {
await saveConvo(
req,
await db.saveConvo(
{
userId: req?.user?.id,
isTemporary: req?.body?.isTemporary,
interfaceConfig: req?.config?.interfaceConfig,
},
{
conversationId,
endpoint: EModelEndpoint.agents,
@ -278,7 +279,7 @@ const createResponse = async (req, res) => {
const isStreaming = request.stream === true;
// Look up the agent
const agent = await getAgent({ id: agentId });
const agent = await db.getAgent({ id: agentId });
if (!agent) {
return sendResponsesErrorResponse(
res,
@ -341,7 +342,7 @@ const createResponse = async (req, res) => {
isInitialAgent: true,
},
{
getConvoFiles,
getConvoFiles: db.getConvoFiles,
getFiles: db.getFiles,
getUserKey: db.getUserKey,
getMessages: db.getMessages,
@ -505,7 +506,7 @@ const createResponse = async (req, res) => {
const balanceConfig = getBalanceConfig(req.config);
const transactionsConfig = getTransactionsConfig(req.config);
recordCollectedUsage(
{ spendTokens, spendStructuredTokens },
{ spendTokens: db.spendTokens, spendStructuredTokens: db.spendStructuredTokens },
{
user: userId,
conversationId,
@ -649,7 +650,7 @@ const createResponse = async (req, res) => {
const balanceConfig = getBalanceConfig(req.config);
const transactionsConfig = getTransactionsConfig(req.config);
recordCollectedUsage(
{ spendTokens, spendStructuredTokens },
{ spendTokens: db.spendTokens, spendStructuredTokens: db.spendStructuredTokens },
{
user: userId,
conversationId,
@ -746,7 +747,7 @@ const listModels = async (req, res) => {
// Get the accessible agents
let agents = [];
if (accessibleAgentIds.length > 0) {
agents = await getAgents({ _id: { $in: accessibleAgentIds } });
agents = await db.getAgents({ _id: { $in: accessibleAgentIds } });
}
// Convert to models format
@ -796,7 +797,7 @@ const getResponse = async (req, res) => {
// The responseId could be either the response ID or the conversation ID
// Try to find a conversation with this ID
const conversation = await getConvo(userId, responseId);
const conversation = await db.getConvo(userId, responseId);
if (!conversation) {
return sendResponsesErrorResponse(

View file

@ -24,15 +24,6 @@ const {
actionDelimiter,
removeNullishValues,
} = require('librechat-data-provider');
const {
getListAgentsByAccess,
countPromotedAgents,
revertAgentVersion,
createAgent,
updateAgent,
deleteAgent,
getAgent,
} = require('~/models/Agent');
const {
findPubliclyAccessibleResources,
findAccessibleResources,
@ -40,14 +31,13 @@ const {
grantPermission,
} = require('~/server/services/PermissionService');
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
const { getCategoriesWithCounts, deleteFileByFilter } = require('~/models');
const { resizeAvatar } = require('~/server/services/Files/images/avatar');
const { getFileStrategy } = require('~/server/utils/getFileStrategy');
const { refreshS3Url } = require('~/server/services/Files/S3/crud');
const { filterFile } = require('~/server/services/Files/process');
const { updateAction, getActions } = require('~/models/Action');
const { getCachedTools } = require('~/server/services/Config');
const { getLogStores } = require('~/cache');
const db = require('~/models');
const systemTools = {
[Tools.execute_code]: true,
@ -92,7 +82,7 @@ const createAgentHandler = async (req, res) => {
}
}
const agent = await createAgent(agentData);
const agent = await db.createAgent(agentData);
try {
await Promise.all([
@ -152,7 +142,7 @@ const getAgentHandler = async (req, res, expandProperties = false) => {
// Permissions are validated by middleware before calling this function
// Simply load the agent by ID
const agent = await getAgent({ id });
const agent = await db.getAgent({ id });
if (!agent) {
return res.status(404).json({ error: 'Agent not found' });
@ -240,7 +230,7 @@ const updateAgentHandler = async (req, res) => {
// Convert OCR to context in incoming updateData
convertOcrToContextInPlace(updateData);
const existingAgent = await getAgent({ id });
const existingAgent = await db.getAgent({ id });
if (!existingAgent) {
return res.status(404).json({ error: 'Agent not found' });
@ -257,7 +247,7 @@ const updateAgentHandler = async (req, res) => {
let updatedAgent =
Object.keys(updateData).length > 0
? await updateAgent({ id }, updateData, {
? await db.updateAgent({ id }, updateData, {
updatingUserId: req.user.id,
})
: existingAgent;
@ -307,7 +297,7 @@ const duplicateAgentHandler = async (req, res) => {
const sensitiveFields = ['api_key', 'oauth_client_id', 'oauth_client_secret'];
try {
const agent = await getAgent({ id });
const agent = await db.getAgent({ id });
if (!agent) {
return res.status(404).json({
error: 'Agent not found',
@ -355,7 +345,7 @@ const duplicateAgentHandler = async (req, res) => {
});
const newActionsList = [];
const originalActions = (await getActions({ agent_id: id }, true)) ?? [];
const originalActions = (await db.getActions({ agent_id: id }, true)) ?? [];
const promises = [];
/**
@ -374,7 +364,7 @@ const duplicateAgentHandler = async (req, res) => {
delete filteredMetadata[field];
}
const newAction = await updateAction(
const newAction = await db.updateAction(
{ action_id: newActionId },
{
metadata: filteredMetadata,
@ -397,7 +387,7 @@ const duplicateAgentHandler = async (req, res) => {
const agentActions = await Promise.all(promises);
newAgentData.actions = agentActions;
const newAgent = await createAgent(newAgentData);
const newAgent = await db.createAgent(newAgentData);
try {
await Promise.all([
@ -450,11 +440,11 @@ const duplicateAgentHandler = async (req, res) => {
const deleteAgentHandler = async (req, res) => {
try {
const id = req.params.id;
const agent = await getAgent({ id });
const agent = await db.getAgent({ id });
if (!agent) {
return res.status(404).json({ error: 'Agent not found' });
}
await deleteAgent({ id });
await db.deleteAgent({ id });
return res.json({ message: 'Agent deleted' });
} catch (error) {
logger.error('[/Agents/:id] Error deleting Agent', error);
@ -529,7 +519,7 @@ const getListAgentsHandler = async (req, res) => {
logger.debug('[/Agents] S3 avatar refresh already checked, skipping');
} else {
try {
const fullList = await getListAgentsByAccess({
const fullList = await db.getListAgentsByAccess({
accessibleIds,
otherParams: {},
limit: MAX_AVATAR_REFRESH_AGENTS,
@ -539,7 +529,7 @@ const getListAgentsHandler = async (req, res) => {
agents: fullList?.data ?? [],
userId,
refreshS3Url,
updateAgent,
updateAgent: db.updateAgent,
});
await cache.set(refreshKey, true, Time.THIRTY_MINUTES);
} catch (err) {
@ -548,7 +538,7 @@ const getListAgentsHandler = async (req, res) => {
}
// Use the new ACL-aware function
const data = await getListAgentsByAccess({
const data = await db.getListAgentsByAccess({
accessibleIds,
otherParams: filter,
limit,
@ -604,7 +594,7 @@ const uploadAgentAvatarHandler = async (req, res) => {
return res.status(400).json({ message: 'Agent ID is required' });
}
const existingAgent = await getAgent({ id: agent_id });
const existingAgent = await db.getAgent({ id: agent_id });
if (!existingAgent) {
return res.status(404).json({ error: 'Agent not found' });
@ -636,7 +626,7 @@ const uploadAgentAvatarHandler = async (req, res) => {
const { deleteFile } = getStrategyFunctions(_avatar.source);
try {
await deleteFile(req, { filepath: _avatar.filepath });
await deleteFileByFilter({ user: req.user.id, filepath: _avatar.filepath });
await db.deleteFileByFilter({ user: req.user.id, filepath: _avatar.filepath });
} catch (error) {
logger.error('[/:agent_id/avatar] Error deleting old avatar', error);
}
@ -649,7 +639,7 @@ const uploadAgentAvatarHandler = async (req, res) => {
},
};
const updatedAgent = await updateAgent({ id: agent_id }, data, {
const updatedAgent = await db.updateAgent({ id: agent_id }, data, {
updatingUserId: req.user.id,
});
res.status(201).json(updatedAgent);
@ -697,7 +687,7 @@ const revertAgentVersionHandler = async (req, res) => {
return res.status(400).json({ error: 'version_index is required' });
}
const existingAgent = await getAgent({ id });
const existingAgent = await db.getAgent({ id });
if (!existingAgent) {
return res.status(404).json({ error: 'Agent not found' });
@ -705,7 +695,7 @@ const revertAgentVersionHandler = async (req, res) => {
// Permissions are enforced via route middleware (ACL EDIT)
const updatedAgent = await revertAgentVersion({ id }, version_index);
const updatedAgent = await db.revertAgentVersion({ id }, version_index);
if (updatedAgent.author) {
updatedAgent.author = updatedAgent.author.toString();
@ -729,8 +719,8 @@ const revertAgentVersionHandler = async (req, res) => {
*/
const getAgentCategories = async (_req, res) => {
try {
const categories = await getCategoriesWithCounts();
const promotedCount = await countPromotedAgents();
const categories = await db.getCategoriesWithCounts();
const promotedCount = await db.countPromotedAgents();
const formattedCategories = categories.map((category) => ({
value: category.value,
label: category.label,

View file

@ -30,15 +30,6 @@ jest.mock('~/server/services/Files/process', () => ({
filterFile: jest.fn(),
}));
jest.mock('~/models/Action', () => ({
updateAction: jest.fn(),
getActions: jest.fn().mockResolvedValue([]),
}));
jest.mock('~/models/File', () => ({
deleteFileByFilter: jest.fn(),
}));
jest.mock('~/server/services/PermissionService', () => ({
findAccessibleResources: jest.fn().mockResolvedValue([]),
findPubliclyAccessibleResources: jest.fn().mockResolvedValue([]),
@ -47,9 +38,18 @@ jest.mock('~/server/services/PermissionService', () => ({
checkPermission: jest.fn().mockResolvedValue(true),
}));
jest.mock('~/models', () => ({
getCategoriesWithCounts: jest.fn(),
}));
jest.mock('~/models', () => {
const mongoose = require('mongoose');
const { createMethods } = require('@librechat/data-schemas');
const methods = createMethods(mongoose, {
removeAllPermissions: jest.fn().mockResolvedValue(undefined),
});
return {
...methods,
getCategoriesWithCounts: jest.fn(),
deleteFileByFilter: jest.fn(),
};
});
// Mock cache for S3 avatar refresh tests
const mockCache = {

View file

@ -1,7 +1,13 @@
const { v4 } = require('uuid');
const { sleep } = require('@librechat/agents');
const { logger } = require('@librechat/data-schemas');
const { sendEvent, getBalanceConfig, getModelMaxTokens, countTokens } = require('@librechat/api');
const {
sendEvent,
countTokens,
checkBalance,
getBalanceConfig,
getModelMaxTokens,
} = require('@librechat/api');
const {
Time,
Constants,
@ -31,10 +37,14 @@ const { createRun, StreamRunManager } = require('~/server/services/Runs');
const { addTitle } = require('~/server/services/Endpoints/assistants');
const { createRunBody } = require('~/server/services/createRunBody');
const { sendResponse } = require('~/server/middleware/error');
const { getTransactions } = require('~/models/Transaction');
const { checkBalance } = require('~/models/balanceMethods');
const { getConvo } = require('~/models/Conversation');
const getLogStores = require('~/cache/getLogStores');
const {
createAutoRefillTransaction,
findBalanceByUser,
getTransactions,
getMultiplier,
getConvo,
} = require('~/models');
const { logViolation, getLogStores } = require('~/cache');
const { getOpenAIClient } = require('./helpers');
/**
@ -275,16 +285,19 @@ const chatV1 = async (req, res) => {
// Count tokens up to the current context window
promptTokens = Math.min(promptTokens, getModelMaxTokens(model));
await checkBalance({
req,
res,
txData: {
model,
user: req.user.id,
tokenType: 'prompt',
amount: promptTokens,
await checkBalance(
{
req,
res,
txData: {
model,
user: req.user.id,
tokenType: 'prompt',
amount: promptTokens,
},
},
});
{ findBalanceByUser, getMultiplier, createAutoRefillTransaction, logViolation },
);
};
const { openai: _openai } = await getOpenAIClient({

View file

@ -1,7 +1,13 @@
const { v4 } = require('uuid');
const { sleep } = require('@librechat/agents');
const { logger } = require('@librechat/data-schemas');
const { sendEvent, getBalanceConfig, getModelMaxTokens, countTokens } = require('@librechat/api');
const {
sendEvent,
countTokens,
checkBalance,
getBalanceConfig,
getModelMaxTokens,
} = require('@librechat/api');
const {
Time,
Constants,
@ -26,10 +32,14 @@ const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
const { createRun, StreamRunManager } = require('~/server/services/Runs');
const { addTitle } = require('~/server/services/Endpoints/assistants');
const { createRunBody } = require('~/server/services/createRunBody');
const { getTransactions } = require('~/models/Transaction');
const { checkBalance } = require('~/models/balanceMethods');
const { getConvo } = require('~/models/Conversation');
const getLogStores = require('~/cache/getLogStores');
const {
getConvo,
getMultiplier,
getTransactions,
findBalanceByUser,
createAutoRefillTransaction,
} = require('~/models');
const { logViolation, getLogStores } = require('~/cache');
const { getOpenAIClient } = require('./helpers');
/**
@ -148,16 +158,19 @@ const chatV2 = async (req, res) => {
// Count tokens up to the current context window
promptTokens = Math.min(promptTokens, getModelMaxTokens(model));
await checkBalance({
req,
res,
txData: {
model,
user: req.user.id,
tokenType: 'prompt',
amount: promptTokens,
await checkBalance(
{
req,
res,
txData: {
model,
user: req.user.id,
tokenType: 'prompt',
amount: promptTokens,
},
},
});
{ findBalanceByUser, getMultiplier, createAutoRefillTransaction, logViolation },
);
};
const { openai: _openai } = await getOpenAIClient({

View file

@ -3,8 +3,8 @@ const { logger } = require('@librechat/data-schemas');
const { CacheKeys, ViolationTypes, ContentTypes } = require('librechat-data-provider');
const { recordUsage, checkMessageGaps } = require('~/server/services/Threads');
const { sendResponse } = require('~/server/middleware/error');
const { getConvo } = require('~/models/Conversation');
const getLogStores = require('~/cache/getLogStores');
const { getConvo } = require('~/models');
/**
* @typedef {Object} ErrorHandlerContext

View file

@ -1,15 +1,14 @@
const fs = require('fs').promises;
const { logger } = require('@librechat/data-schemas');
const { FileContext } = require('librechat-data-provider');
const { deleteFileByFilter, updateAssistantDoc, getAssistants } = require('~/models');
const { uploadImageBuffer, filterFile } = require('~/server/services/Files/process');
const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
const { deleteAssistantActions } = require('~/server/services/ActionService');
const { updateAssistantDoc, getAssistants } = require('~/models/Assistant');
const { getOpenAIClient, fetchAssistants } = require('./helpers');
const { getCachedTools } = require('~/server/services/Config');
const { manifestToolMap } = require('~/app/clients/tools');
const { deleteFileByFilter } = require('~/models');
/**
* Create an assistant.

View file

@ -3,8 +3,8 @@ const { ToolCallTypes } = require('librechat-data-provider');
const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
const { validateAndUpdateTool } = require('~/server/services/ActionService');
const { getCachedTools } = require('~/server/services/Config');
const { updateAssistantDoc } = require('~/models/Assistant');
const { manifestToolMap } = require('~/app/clients/tools');
const { updateAssistantDoc } = require('~/models');
const { getOpenAIClient } = require('./helpers');
/**

View file

@ -9,13 +9,11 @@ const {
ToolCallTypes,
PermissionTypes,
} = require('librechat-data-provider');
const { getRoleByName, createToolCall, getToolCallsByConvo, getMessage } = require('~/models');
const { processFileURL, uploadImageBuffer } = require('~/server/services/Files/process');
const { processCodeOutput } = require('~/server/services/Files/Code/process');
const { createToolCall, getToolCallsByConvo } = require('~/models/ToolCall');
const { loadAuthValues } = require('~/server/services/Tools/credentials');
const { loadTools } = require('~/app/clients/tools/util');
const { getRoleByName } = require('~/models/Role');
const { getMessage } = require('~/models/Message');
const fieldsMap = {
[Tools.execute_code]: [EnvVar.CODE_API_KEY],