File size: 3,802 Bytes
f0743f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | import { logger } from '@librechat/data-schemas';
import type { NextFunction, Request as ServerRequest, Response as ServerResponse } from 'express';
import type { IBalance, IUser, BalanceConfig, ObjectId, AppConfig } from '@librechat/data-schemas';
import type { Model } from 'mongoose';
import type { BalanceUpdateFields } from '~/types';
import { getBalanceConfig } from '~/app/config';
export interface BalanceMiddlewareOptions {
getAppConfig: (options?: { role?: string; refresh?: boolean }) => Promise<AppConfig>;
Balance: Model<IBalance>;
}
/**
* Build an object containing fields that need updating
* @param config - The balance configuration
* @param userRecord - The user's current balance record, if any
* @param userId - The user's ID
* @returns Fields that need updating
*/
function buildUpdateFields(
config: BalanceConfig,
userRecord: IBalance | null,
userId: string,
): BalanceUpdateFields {
const updateFields: BalanceUpdateFields = {};
// Ensure user record has the required fields
if (!userRecord) {
updateFields.user = userId;
updateFields.tokenCredits = config.startBalance;
}
if (userRecord?.tokenCredits == null && config.startBalance != null) {
updateFields.tokenCredits = config.startBalance;
}
const isAutoRefillConfigValid =
config.autoRefillEnabled &&
config.refillIntervalValue != null &&
config.refillIntervalUnit != null &&
config.refillAmount != null;
if (!isAutoRefillConfigValid) {
return updateFields;
}
if (userRecord?.autoRefillEnabled !== config.autoRefillEnabled) {
updateFields.autoRefillEnabled = config.autoRefillEnabled;
}
if (userRecord?.refillIntervalValue !== config.refillIntervalValue) {
updateFields.refillIntervalValue = config.refillIntervalValue;
}
if (userRecord?.refillIntervalUnit !== config.refillIntervalUnit) {
updateFields.refillIntervalUnit = config.refillIntervalUnit;
}
if (userRecord?.refillAmount !== config.refillAmount) {
updateFields.refillAmount = config.refillAmount;
}
// Initialize lastRefill if it's missing when auto-refill is enabled
if (config.autoRefillEnabled && !userRecord?.lastRefill) {
updateFields.lastRefill = new Date();
}
return updateFields;
}
/**
* Factory function to create middleware that synchronizes user balance settings with current balance configuration.
* @param options - Options containing getBalanceConfig function and Balance model
* @returns Express middleware function
*/
export function createSetBalanceConfig({
getAppConfig,
Balance,
}: BalanceMiddlewareOptions): (
req: ServerRequest,
res: ServerResponse,
next: NextFunction,
) => Promise<void> {
return async (req: ServerRequest, res: ServerResponse, next: NextFunction): Promise<void> => {
try {
const user = req.user as IUser & { _id: string | ObjectId };
const appConfig = await getAppConfig({ role: user?.role });
const balanceConfig = getBalanceConfig(appConfig);
if (!balanceConfig?.enabled) {
return next();
}
if (balanceConfig.startBalance == null) {
return next();
}
if (!user || !user._id) {
return next();
}
const userId = typeof user._id === 'string' ? user._id : user._id.toString();
const userBalanceRecord = await Balance.findOne({ user: userId }).lean();
const updateFields = buildUpdateFields(balanceConfig, userBalanceRecord, userId);
if (Object.keys(updateFields).length === 0) {
return next();
}
await Balance.findOneAndUpdate(
{ user: userId },
{ $set: updateFields },
{ upsert: true, new: true },
);
next();
} catch (error) {
logger.error('Error setting user balance:', error);
next(error);
}
};
}
|