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 */