From 57ab67cbacbc92ccd886754b2af5a3bc075fce40 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 9 Nov 2026 08:27:13 -0500 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=8F=20fix:=20Measure=20Usage=20Percent?= =?UTF-8?q?=20Against=20the=20Allowance=20Each=20Period=20Granted=20(#1785?= =?UTF-8?q?0)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 📏 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. --- api/server/controllers/Balance.js | 1 - api/server/controllers/Balance.spec.js | 25 ++++++++++ client/src/utils/__tests__/balance.spec.ts | 32 +++++++++++++ client/src/utils/balance.ts | 11 +++-- packages/api/src/middleware/balance.spec.ts | 47 ++++++++++++++++++- packages/api/src/middleware/balance.ts | 17 ++++++- .../api/src/middleware/checkBalance.spec.ts | 28 +++++++++-- packages/api/src/middleware/checkBalance.ts | 2 + packages/api/src/schedules/service.ts | 16 +++++-- packages/api/src/types/balance.ts | 3 ++ packages/data-provider/src/balance.ts | 20 ++++++++ packages/data-provider/src/types.ts | 4 ++ .../src/methods/transaction.spec.ts | 15 ++++++ .../data-schemas/src/methods/transaction.ts | 2 + packages/data-schemas/src/methods/user.ts | 11 ++++- packages/data-schemas/src/schema/balance.ts | 8 ++++ packages/data-schemas/src/types/balance.ts | 6 +++ 17 files changed, 233 insertions(+), 15 deletions(-) diff --git a/api/server/controllers/Balance.js b/api/server/controllers/Balance.js index 8df579e5c6..ebf963da6c 100644 --- a/api/server/controllers/Balance.js +++ b/api/server/controllers/Balance.js @@ -18,7 +18,6 @@ async function balanceController(req, res) { if (!result.autoRefillEnabled) { delete result.refillIntervalValue; delete result.refillIntervalUnit; - delete result.lastRefill; delete result.refillAmount; } diff --git a/api/server/controllers/Balance.spec.js b/api/server/controllers/Balance.spec.js index 833f2d8c54..86ae35bbec 100644 --- a/api/server/controllers/Balance.spec.js +++ b/api/server/controllers/Balance.spec.js @@ -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 = { diff --git a/client/src/utils/__tests__/balance.spec.ts b/client/src/utils/__tests__/balance.spec.ts index 5b1b4aa5c0..d81bd34148 100644 --- a/client/src/utils/__tests__/balance.spec.ts +++ b/client/src/utils/__tests__/balance.spec.ts @@ -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); diff --git a/client/src/utils/balance.ts b/client/src/utils/balance.ts index ff7cc94e08..0930611c78 100644 --- a/client/src/utils/balance.ts +++ b/client/src/utils/balance.ts @@ -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'; diff --git a/packages/api/src/middleware/balance.spec.ts b/packages/api/src/middleware/balance.spec.ts index ead52b46dc..d2cbc51f57 100644 --- a/packages/api/src/middleware/balance.spec.ts +++ b/packages/api/src/middleware/balance.spec.ts @@ -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({ diff --git a/packages/api/src/middleware/balance.ts b/packages/api/src/middleware/balance.ts index 1c551c874e..5fe33824f3 100644 --- a/packages/api/src/middleware/balance.ts +++ b/packages/api/src/middleware/balance.ts @@ -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; } diff --git a/packages/api/src/middleware/checkBalance.spec.ts b/packages/api/src/middleware/checkBalance.spec.ts index 4796ab36ba..910a186ad1 100644 --- a/packages/api/src/middleware/checkBalance.spec.ts +++ b/packages/api/src/middleware/checkBalance.spec.ts @@ -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); diff --git a/packages/api/src/middleware/checkBalance.ts b/packages/api/src/middleware/checkBalance.ts index aaa849e7ed..975616855d 100644 --- a/packages/api/src/middleware/checkBalance.ts +++ b/packages/api/src/middleware/checkBalance.ts @@ -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); diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index 5592aed9a6..0514b21c11 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -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; diff --git a/packages/api/src/types/balance.ts b/packages/api/src/types/balance.ts index a47ceed732..e6794c85af 100644 --- a/packages/api/src/types/balance.ts +++ b/packages/api/src/types/balance.ts @@ -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; } diff --git a/packages/data-provider/src/balance.ts b/packages/data-provider/src/balance.ts index ce28af9957..e657a5876f 100644 --- a/packages/data-provider/src/balance.ts +++ b/packages/data-provider/src/balance.ts @@ -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, +): 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< diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index c0a439cff5..4285d01a98 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -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; }; /* -------------------------------------------------------------------------- */ diff --git a/packages/data-schemas/src/methods/transaction.spec.ts b/packages/data-schemas/src/methods/transaction.spec.ts index 9418da2a30..2e1d93def8 100644 --- a/packages/data-schemas/src/methods/transaction.spec.ts +++ b/packages/data-schemas/src/methods/transaction.spec.ts @@ -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 }); diff --git a/packages/data-schemas/src/methods/transaction.ts b/packages/data-schemas/src/methods/transaction.ts index 01874c090a..50fca1e819 100644 --- a/packages/data-schemas/src/methods/transaction.ts +++ b/packages/data-schemas/src/methods/transaction.ts @@ -384,6 +384,8 @@ export function createTransactionMethods( $set: { tokenCredits: credits, lastRefill: now, + allotment: record.refillAmount, + allottedAt: now, pendingRefill, }, }, diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index 22e153bb8b..4820a8827f 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -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 && diff --git a/packages/data-schemas/src/schema/balance.ts b/packages/data-schemas/src/schema/balance.ts index 48f672679e..815809b042 100644 --- a/packages/data-schemas/src/schema/balance.ts +++ b/packages/data-schemas/src/schema/balance.ts @@ -42,6 +42,14 @@ const balanceSchema: Schema = new Schema({ 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, diff --git a/packages/data-schemas/src/types/balance.ts b/packages/data-schemas/src/types/balance.ts index 48c75f461c..8bae6f79c7 100644 --- a/packages/data-schemas/src/types/balance.ts +++ b/packages/data-schemas/src/types/balance.ts @@ -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 */