mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-11-12 09:45:47 +00:00
📏 fix: Measure Usage Percent Against the Allowance Each Period Granted (#17850)
* 📏 fix: Measure Usage Percent Against the Allowance Each Period Granted Store the credits granted at the start of each balance period on the record and measure the usage percent against it, so changing refillAmount or startBalance no longer moves a user's percent until their next grant. Legacy records fall back to the configured amount. * 🕰️ fix: Ignore a Stored Allotment a Later Refill Has Replaced Stamp each stored allotment with the time of its grant and measure against it only while no refill is more recent, so a refill by a replica that predates the field falls back to the configured amount instead of reporting usage against the prior period's allowance. * 🔒 fix: Write the Allotment Grant Atomically and Return Its Period Start Insert lastRefill with the credit, allotment and grant time when a balance record is created, so a racing creator never moves another's period start past its grant. Return lastRefill even without auto-refill so the client can place a stored allotment in its period, and fall back when it cannot.
This commit is contained in:
parent
23aa0f822d
commit
57ab67cbac
17 changed files with 233 additions and 15 deletions
|
|
@ -18,7 +18,6 @@ async function balanceController(req, res) {
|
|||
if (!result.autoRefillEnabled) {
|
||||
delete result.refillIntervalValue;
|
||||
delete result.refillIntervalUnit;
|
||||
delete result.lastRefill;
|
||||
delete result.refillAmount;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -55,6 +55,31 @@ describe('balanceController', () => {
|
|||
});
|
||||
});
|
||||
|
||||
it('keeps the period start that validates a stored allotment without auto-refill', async () => {
|
||||
const lastRefill = new Date('2026-07-01T00:00:00.000Z');
|
||||
const res = createResponse();
|
||||
res.locals.balanceConfigEnabled = true;
|
||||
res.locals.balanceData = {
|
||||
_id: 'balance-1',
|
||||
tokenCredits: 100,
|
||||
autoRefillEnabled: false,
|
||||
refillAmount: 500,
|
||||
lastRefill,
|
||||
allotment: 1000,
|
||||
allottedAt: lastRefill,
|
||||
};
|
||||
|
||||
await balanceController({ user: { id: 'user-1' } }, res);
|
||||
|
||||
expect(res.json).toHaveBeenCalledWith({
|
||||
tokenCredits: 100,
|
||||
autoRefillEnabled: false,
|
||||
lastRefill,
|
||||
allotment: 1000,
|
||||
allottedAt: lastRefill,
|
||||
});
|
||||
});
|
||||
|
||||
it('returns not found when balance is enabled and no record exists', async () => {
|
||||
findBalanceByUser.mockResolvedValue(null);
|
||||
const req = {
|
||||
|
|
|
|||
|
|
@ -39,6 +39,38 @@ describe('summarizeBalance', () => {
|
|||
expect(summary.tone).toBe('normal');
|
||||
});
|
||||
|
||||
it('measures against the stored allotment after the configured amount changes', () => {
|
||||
const summary = summarizeBalance(
|
||||
refilling({
|
||||
allotment: 40_000,
|
||||
allottedAt: '2026-07-01T00:00:00.000Z',
|
||||
refillAmount: 20_000,
|
||||
}),
|
||||
);
|
||||
expect(summary.allotment).toBe(40_000);
|
||||
expect(summary.usedPercent).toBe(69);
|
||||
expect(summary.refillAmount).toBe(20_000);
|
||||
});
|
||||
|
||||
it('falls back to the configured amount when no allotment is stored', () => {
|
||||
expect(summarizeBalance(refilling({ allotment: 0 })).allotment).toBe(20_000);
|
||||
expect(summarizeBalance(refilling({ allotment: 40_000 })).allotment).toBe(20_000);
|
||||
});
|
||||
|
||||
it('ignores an allotment from before the last refill', () => {
|
||||
const stale = refilling({ allotment: 40_000, allottedAt: '2026-06-24T00:00:00.000Z' });
|
||||
expect(summarizeBalance(stale).allotment).toBe(20_000);
|
||||
});
|
||||
|
||||
it('ignores an allotment it cannot place in a period', () => {
|
||||
const unplaced = refilling({
|
||||
allotment: 40_000,
|
||||
allottedAt: '2026-07-01T00:00:00.000Z',
|
||||
lastRefill: undefined,
|
||||
});
|
||||
expect(summarizeBalance(unplaced).allotment).toBe(20_000);
|
||||
});
|
||||
|
||||
it('reads 0% used when saved-up credits exceed the allotment', () => {
|
||||
const summary = summarizeBalance(refilling({ tokenCredits: 45_000 }));
|
||||
expect(summary.usedPercent).toBe(0);
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
import { CREDITS_PER_USD, getRefillEligibilityDate } from 'librechat-data-provider';
|
||||
import {
|
||||
CREDITS_PER_USD,
|
||||
getGrantedAllotment,
|
||||
getRefillEligibilityDate,
|
||||
} from 'librechat-data-provider';
|
||||
import type { BalanceDisplay, BalanceRefillMode, TBalanceResponse } from 'librechat-data-provider';
|
||||
import type { CurrencyConfig } from './tokens';
|
||||
import { formatCost } from './tokens';
|
||||
|
|
@ -15,7 +19,8 @@ export type BalanceTone = 'normal' | 'warning' | 'danger';
|
|||
export interface BalanceSummary {
|
||||
display: BalanceDisplay;
|
||||
credits: number;
|
||||
/** Credits one period grants: the refill amount under auto-refill, else the starting balance. */
|
||||
/** Credits the current period started with: the stored grant, else the configured refill
|
||||
* amount under auto-refill, else the starting balance. */
|
||||
allotment: number | null;
|
||||
/** Whole percent of the allotment spent, 0–100; 100 once credits run out, else null without
|
||||
* an allotment to measure against. */
|
||||
|
|
@ -70,7 +75,7 @@ export function summarizeBalance(
|
|||
): BalanceSummary {
|
||||
const credits = Number.isFinite(balance.tokenCredits) ? balance.tokenCredits : 0;
|
||||
const refillAmount = balance.autoRefillEnabled ? positive(balance.refillAmount) : null;
|
||||
const allotment = refillAmount ?? positive(startBalance);
|
||||
const allotment = getGrantedAllotment(balance) ?? refillAmount ?? positive(startBalance);
|
||||
const usedPercent = getUsedPercent(credits, allotment);
|
||||
|
||||
let tone: BalanceTone = 'normal';
|
||||
|
|
|
|||
|
|
@ -68,9 +68,16 @@ describe('createSetBalanceConfig', () => {
|
|||
const getAppConfig = jest.fn().mockResolvedValue({
|
||||
balance: { enabled: true, startBalance: 1000 },
|
||||
});
|
||||
const winnerGrant = new Date('2026-07-01T00:00:00.000Z');
|
||||
const findThenCharge = async (id: string) => {
|
||||
const record = await findBalanceByUser(id);
|
||||
await Balance.create({ user: id, tokenCredits: 50 });
|
||||
await Balance.create({
|
||||
user: id,
|
||||
tokenCredits: 50,
|
||||
allotment: 1000,
|
||||
allottedAt: winnerGrant,
|
||||
lastRefill: winnerGrant,
|
||||
});
|
||||
return record;
|
||||
};
|
||||
|
||||
|
|
@ -83,10 +90,46 @@ describe('createSetBalanceConfig', () => {
|
|||
|
||||
await middleware(createMockRequest(userId) as ServerRequest, res as ServerResponse, mockNext);
|
||||
|
||||
expect((await Balance.findOne({ user: userId }).lean())?.tokenCredits).toBe(50);
|
||||
const stored = await Balance.findOne({ user: userId }).lean();
|
||||
expect(stored).toMatchObject({
|
||||
tokenCredits: 50,
|
||||
allotment: 1000,
|
||||
allottedAt: winnerGrant,
|
||||
lastRefill: winnerGrant,
|
||||
});
|
||||
expect((res.locals as { balanceData?: IBalance }).balanceData?.tokenCredits).toBe(50);
|
||||
});
|
||||
|
||||
test('stores the start balance as the allotment only when it grants credits', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const config = { enabled: true, startBalance: 1000 };
|
||||
const getAppConfig = jest.fn().mockResolvedValue({ balance: config });
|
||||
const middleware = createSetBalanceConfig({
|
||||
getAppConfig,
|
||||
findBalanceByUser,
|
||||
upsertBalanceFields,
|
||||
});
|
||||
|
||||
await middleware(
|
||||
createMockRequest(userId) as ServerRequest,
|
||||
createMockResponse() as ServerResponse,
|
||||
mockNext,
|
||||
);
|
||||
const created = await Balance.findOne({ user: userId }).lean();
|
||||
expect(created?.allotment).toBe(1000);
|
||||
expect(created?.allottedAt).toEqual(created?.lastRefill);
|
||||
|
||||
getAppConfig.mockResolvedValue({ balance: { ...config, startBalance: 5000 } });
|
||||
await middleware(
|
||||
createMockRequest(userId) as ServerRequest,
|
||||
createMockResponse() as ServerResponse,
|
||||
mockNext,
|
||||
);
|
||||
const record = await Balance.findOne({ user: userId }).lean();
|
||||
expect(record?.tokenCredits).toBe(1000);
|
||||
expect(record?.allotment).toBe(1000);
|
||||
});
|
||||
|
||||
test('should create balance record for new user with start balance', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const getAppConfig = jest.fn().mockResolvedValue({
|
||||
|
|
|
|||
|
|
@ -80,6 +80,15 @@ export function buildBalanceUpdateFields(
|
|||
updateFields.tokenCredits = config.startBalance;
|
||||
}
|
||||
|
||||
if (updateFields.tokenCredits != null) {
|
||||
const grantedAt = userRecord?.lastRefill ?? new Date();
|
||||
updateFields.allotment = updateFields.tokenCredits;
|
||||
updateFields.allottedAt = grantedAt;
|
||||
if (!userRecord?.lastRefill) {
|
||||
updateFields.lastRefill = grantedAt;
|
||||
}
|
||||
}
|
||||
|
||||
const isAutoRefillConfigValid =
|
||||
config.autoRefillEnabled &&
|
||||
config.refillIntervalValue != null &&
|
||||
|
|
@ -118,7 +127,7 @@ export function buildBalanceUpdateFields(
|
|||
|
||||
// Initialize lastRefill if it's missing when auto-refill is enabled
|
||||
if (config.autoRefillEnabled && !userRecord?.lastRefill) {
|
||||
updateFields.lastRefill = new Date();
|
||||
updateFields.lastRefill ??= new Date();
|
||||
}
|
||||
|
||||
return updateFields;
|
||||
|
|
@ -177,9 +186,13 @@ export function createSetBalanceConfig({
|
|||
}
|
||||
|
||||
if (userBalanceRecord == null) {
|
||||
const { tokenCredits, ...syncFields } = updateFields;
|
||||
// The grant tuple is insert-only so a racing creator keeps its own period and allotment.
|
||||
const { tokenCredits, allotment, allottedAt, lastRefill, ...syncFields } = updateFields;
|
||||
balanceLocals.balanceData = await upsertBalanceFields(userId, syncFields, {
|
||||
tokenCredits,
|
||||
allotment,
|
||||
allottedAt,
|
||||
lastRefill,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -298,7 +298,13 @@ describe('checkBalance', () => {
|
|||
|
||||
await checkBalance({ req, res, txData: baseTxData }, deps);
|
||||
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({ user: 'user-1', tokenCredits: 5000 });
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({
|
||||
user: 'user-1',
|
||||
tokenCredits: 5000,
|
||||
allotment: 5000,
|
||||
allottedAt: expect.any(Date),
|
||||
lastRefill: expect.any(Date),
|
||||
});
|
||||
});
|
||||
|
||||
it('includes auto-refill fields when configured', async () => {
|
||||
|
|
@ -317,6 +323,8 @@ describe('checkBalance', () => {
|
|||
expect(reserveRequest(deps).initialBalance).toEqual({
|
||||
user: 'user-1',
|
||||
tokenCredits: 5000,
|
||||
allotment: 5000,
|
||||
allottedAt: expect.any(Date),
|
||||
autoRefillEnabled: true,
|
||||
refillIntervalValue: 1,
|
||||
refillIntervalUnit: 'days',
|
||||
|
|
@ -332,7 +340,13 @@ describe('checkBalance', () => {
|
|||
|
||||
await checkBalance({ req, res, txData: baseTxData }, deps);
|
||||
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({ user: 'user-1', tokenCredits: 5000 });
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({
|
||||
user: 'user-1',
|
||||
tokenCredits: 5000,
|
||||
allotment: 5000,
|
||||
allottedAt: expect.any(Date),
|
||||
lastRefill: expect.any(Date),
|
||||
});
|
||||
});
|
||||
|
||||
it('creates a record with a startBalance of 0', async () => {
|
||||
|
|
@ -340,7 +354,13 @@ describe('checkBalance', () => {
|
|||
|
||||
await checkBalance({ req, res, txData: baseTxData }, deps);
|
||||
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({ user: 'user-1', tokenCredits: 0 });
|
||||
expect(reserveRequest(deps).initialBalance).toEqual({
|
||||
user: 'user-1',
|
||||
tokenCredits: 0,
|
||||
allotment: 0,
|
||||
allottedAt: expect.any(Date),
|
||||
lastRefill: expect.any(Date),
|
||||
});
|
||||
});
|
||||
|
||||
it.each([
|
||||
|
|
@ -580,6 +600,7 @@ describe('checkBalance', () => {
|
|||
autoRefillEnabled: true,
|
||||
refillMode: 'reset',
|
||||
refillAmount: 1000,
|
||||
allotment: 1000,
|
||||
refillIntervalValue: 1,
|
||||
refillIntervalUnit: 'weeks',
|
||||
lastRefill,
|
||||
|
|
@ -602,6 +623,7 @@ describe('checkBalance', () => {
|
|||
tokenCredits: 500,
|
||||
reservedCredits: 400,
|
||||
refillAmount: 2000,
|
||||
allotment: 1000,
|
||||
lastRefill,
|
||||
});
|
||||
expect(await methods.getTransactions({ user, context: 'balanceReset' })).toHaveLength(0);
|
||||
|
|
|
|||
|
|
@ -231,6 +231,8 @@ export async function checkBalance(
|
|||
const {
|
||||
user: _user,
|
||||
tokenCredits: _credits,
|
||||
allotment: _allotment,
|
||||
allottedAt: _allottedAt,
|
||||
lastRefill: _lastRefill,
|
||||
...policy
|
||||
} = buildBalanceUpdateFields(deps.balanceConfig, null, user);
|
||||
|
|
|
|||
|
|
@ -722,7 +722,16 @@ export function createSchedulesService(
|
|||
// record in between must never be handed back its starting balance. Each case
|
||||
// fences the credit write so it can only create, never restore, credits. The
|
||||
// refill-config fields are a genuine sync and never carry a credit.
|
||||
const { user: initUser, tokenCredits, ...syncFields } = updateFields;
|
||||
const {
|
||||
user: initUser,
|
||||
tokenCredits,
|
||||
allotment,
|
||||
allottedAt,
|
||||
lastRefill,
|
||||
...syncFields
|
||||
} = updateFields;
|
||||
// The grant tuple is written only with the credit, so a racing creator keeps its period.
|
||||
const grant = { allotment, allottedAt, lastRefill };
|
||||
if (record == null) {
|
||||
// ABSENT record: upsert. The credit rides `$setOnInsert` so a document created
|
||||
// by a concurrent charge keeps its charged balance; the refill-config sync is
|
||||
|
|
@ -731,7 +740,7 @@ export function createSchedulesService(
|
|||
set: syncFields,
|
||||
setOnInsert: {
|
||||
...(initUser != null ? { user: initUser } : {}),
|
||||
...(tokenCredits != null ? { tokenCredits } : {}),
|
||||
...(tokenCredits != null ? { tokenCredits, ...grant } : {}),
|
||||
},
|
||||
});
|
||||
} else if (tokenCredits != null) {
|
||||
|
|
@ -739,9 +748,10 @@ export function createSchedulesService(
|
|||
// concurrent initializer/charge already set and spent between the read and
|
||||
// here. Initialize under a `{ tokenCredits: null }` CAS; on a miss re-read the
|
||||
// winner's balance instead of overwriting it.
|
||||
// The allotment rides the CAS so it is written only with the credit it describes.
|
||||
const initialized = await deps.initializeNullBalance(user.id, {
|
||||
tokenCredits,
|
||||
sync: syncFields,
|
||||
sync: { ...syncFields, ...grant },
|
||||
});
|
||||
record =
|
||||
initialized ?? (await deps.findBalance(user.id, { applyReset: false })) ?? record;
|
||||
|
|
|
|||
|
|
@ -9,4 +9,7 @@ export interface BalanceUpdateFields {
|
|||
refillAmount?: number;
|
||||
refillMode?: BalanceRefillMode;
|
||||
lastRefill?: Date;
|
||||
/** Written only together with `tokenCredits`: the credits that grant starts the period with */
|
||||
allotment?: number;
|
||||
allottedAt?: Date;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -65,6 +65,26 @@ export function getRefillEligibilityDate(
|
|||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The credits the current period was granted, or null when none is stored, the period start is
|
||||
* unknown, or a refill since (by a server that predates the field) left it describing an earlier
|
||||
* period.
|
||||
*/
|
||||
export function getGrantedAllotment(
|
||||
record: Pick<TBalanceResponse, 'allotment' | 'allottedAt' | 'lastRefill'>,
|
||||
): number | null {
|
||||
const { allotment, allottedAt, lastRefill } = record;
|
||||
if (allotment == null || !(allotment > 0) || allottedAt == null) {
|
||||
return null;
|
||||
}
|
||||
const granted = new Date(allottedAt).getTime();
|
||||
if (Number.isNaN(granted)) {
|
||||
return null;
|
||||
}
|
||||
const refilled = lastRefill == null ? NaN : new Date(lastRefill).getTime();
|
||||
return Number.isNaN(refilled) || granted < refilled ? null : allotment;
|
||||
}
|
||||
|
||||
/** Whether the configured refill/reset period has elapsed. */
|
||||
export function isBalanceRefillDue(
|
||||
record: Pick<
|
||||
|
|
|
|||
|
|
@ -1354,6 +1354,10 @@ export type TBalanceResponse = {
|
|||
lastRefill?: Date | string;
|
||||
refillAmount?: number;
|
||||
refillMode?: BalanceRefillMode;
|
||||
/** Credits granted at the start of the current period; absent on records not yet granted one. */
|
||||
allotment?: number;
|
||||
/** When `allotment` was granted; a later `lastRefill` means a writer replaced the period without it. */
|
||||
allottedAt?: Date | string;
|
||||
};
|
||||
|
||||
/* -------------------------------------------------------------------------- */
|
||||
|
|
|
|||
|
|
@ -1680,6 +1680,21 @@ describe('Balance Reservations', () => {
|
|||
expect(resets[0].rawAmount).toBe(1500);
|
||||
});
|
||||
|
||||
test('stores the granted allowance and keeps it until the next reset', async () => {
|
||||
const user = new mongoose.Types.ObjectId();
|
||||
await Balance.create({ user, ...resettable, tokenCredits: 500 });
|
||||
|
||||
const granted = await findBalanceByUser(user.toString());
|
||||
expect(granted?.allotment).toBe(1000);
|
||||
expect(granted?.allottedAt).toEqual(granted?.lastRefill);
|
||||
await reserve(user.toString(), 100);
|
||||
const changed = await upsertBalanceFields(user.toString(), { refillAmount: 3000 });
|
||||
|
||||
expect(changed?.tokenCredits).toBe(1000);
|
||||
expect(changed?.allotment).toBe(1000);
|
||||
expect((await readState(user))?.allotment).toBe(1000);
|
||||
});
|
||||
|
||||
test('uses current settings if reset mode changes during its fenced write', async () => {
|
||||
const user = new mongoose.Types.ObjectId();
|
||||
await Balance.create({ user, ...resettable, tokenCredits: 500 });
|
||||
|
|
|
|||
|
|
@ -384,6 +384,8 @@ export function createTransactionMethods(
|
|||
$set: {
|
||||
tokenCredits: credits,
|
||||
lastRefill: now,
|
||||
allotment: record.refillAmount,
|
||||
allottedAt: now,
|
||||
pendingRefill,
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -291,6 +291,7 @@ export function createUserMethods(
|
|||
}
|
||||
|
||||
const Balance = mongoose.models.Balance;
|
||||
const grantedAt = new Date();
|
||||
const initial: {
|
||||
tokenCredits: number;
|
||||
autoRefillEnabled?: boolean;
|
||||
|
|
@ -298,7 +299,15 @@ export function createUserMethods(
|
|||
refillIntervalUnit?: RefillIntervalUnit;
|
||||
refillAmount?: number;
|
||||
refillMode?: BalanceRefillMode;
|
||||
} = { tokenCredits: balanceConfig.startBalance };
|
||||
allotment: number;
|
||||
allottedAt: Date;
|
||||
lastRefill: Date;
|
||||
} = {
|
||||
tokenCredits: balanceConfig.startBalance,
|
||||
allotment: balanceConfig.startBalance,
|
||||
allottedAt: grantedAt,
|
||||
lastRefill: grantedAt,
|
||||
};
|
||||
|
||||
if (
|
||||
balanceConfig.autoRefillEnabled &&
|
||||
|
|
|
|||
|
|
@ -42,6 +42,14 @@ const balanceSchema: Schema<t.IBalance> = new Schema<t.IBalance>({
|
|||
enum: BALANCE_REFILL_MODES,
|
||||
default: 'add',
|
||||
},
|
||||
/** Credits granted at the start of the current period; usage percent is measured against it */
|
||||
allotment: {
|
||||
type: Number,
|
||||
},
|
||||
/** When `allotment` was granted; a later `lastRefill` marks it stale */
|
||||
allottedAt: {
|
||||
type: Date,
|
||||
},
|
||||
tenantId: {
|
||||
type: String,
|
||||
index: true,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,10 @@ export interface IBalance extends Document {
|
|||
lastRefill: Date;
|
||||
refillAmount: number;
|
||||
refillMode?: BalanceRefillMode;
|
||||
/** Credits granted at the start of the current period, kept when the configured amount changes */
|
||||
allotment?: number;
|
||||
/** When `allotment` was granted; stale once `lastRefill` is later */
|
||||
allottedAt?: Date;
|
||||
tenantId?: string;
|
||||
/** Reservation state is excluded from reads unless explicitly selected */
|
||||
reservations?: IBalanceReservation[];
|
||||
|
|
@ -43,6 +47,8 @@ export interface IBalanceUpdate {
|
|||
refillAmount?: number;
|
||||
refillMode?: BalanceRefillMode;
|
||||
lastRefill?: Date;
|
||||
allotment?: number;
|
||||
allottedAt?: Date;
|
||||
}
|
||||
|
||||
/** Holds credits against a user's balance for the lifetime of one in-flight request */
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue