diff --git a/api/server/services/Schedules/index.js b/api/server/services/Schedules/index.js index 40f592ebdb0..cd96bc2561b 100644 --- a/api/server/services/Schedules/index.js +++ b/api/server/services/Schedules/index.js @@ -27,7 +27,8 @@ function getService() { getAppConfig, findUserById: (userId) => mongoose.models.User.findById(userId).select('_id tenantId role').lean(), - findBalance: (userId) => methods.findBalanceByUser(userId, { includeReservedCredits: true }), + findBalance: (userId, options) => + methods.findBalanceByUser(userId, { includeReservedCredits: true, ...options }), upsertBalance: (userId, { set, setOnInsert }) => methods.upsertBalanceFields(userId, set ?? {}, setOnInsert ?? {}), // Compare-and-set: only initialize an existing record while its credit is still null. diff --git a/client/src/components/Balance/Summary.tsx b/client/src/components/Balance/Summary.tsx index 62dcd270738..fc7cca81c6a 100644 --- a/client/src/components/Balance/Summary.tsx +++ b/client/src/components/Balance/Summary.tsx @@ -3,9 +3,10 @@ import { useTranslation } from 'react-i18next'; import { Progress, Skeleton } from '@librechat/client'; import type { BalanceSummary, BalanceTone, CurrencyConfig } from '~/utils'; import type { BalanceState } from '~/hooks/useBalanceSummary'; -import { formatBalanceAmount, formatTimeUntil, cn } from '~/utils'; +import { formatBalanceAmount, formatTimeUntil } from '~/utils/balance'; import useTimeTick from '~/hooks/useTimeTick'; -import { useLocalize } from '~/hooks'; +import useLocalize from '~/hooks/useLocalize'; +import cn from '~/utils/cn'; const TONE_VARIANT = { normal: 'default', @@ -31,12 +32,23 @@ function useRefillText(summary: BalanceSummary, currency?: CurrencyConfig): stri /** Re-render once a minute so "in 5 minutes" counts down while the view stays open */ useTimeTick(); - const { display, refillAmount, nextRefill } = summary; + const { display, refillAmount, refillMode, nextRefill } = summary; if (refillAmount == null) { return localize('com_ui_balance_no_refill'); } const when = nextRefill != null ? formatTimeUntil(nextRefill, Date.now(), i18n.resolvedLanguage) : null; + if (refillMode === 'reset') { + if (display === 'percent') { + return when != null + ? localize('com_ui_balance_reset_in', { 0: when }) + : localize('com_ui_balance_reset_due'); + } + const amount = formatBalanceAmount(refillAmount, display, currency); + return when != null + ? localize('com_ui_balance_reset_amount_in', { 0: amount, 1: when }) + : localize('com_ui_balance_reset_amount_due', { 0: amount }); + } if (display === 'percent') { return when != null ? localize('com_ui_balance_refill_in', { 0: when }) diff --git a/client/src/components/Balance/__tests__/Summary.spec.tsx b/client/src/components/Balance/__tests__/Summary.spec.tsx index 0afaddb794a..9e8b6a0e008 100644 --- a/client/src/components/Balance/__tests__/Summary.spec.tsx +++ b/client/src/components/Balance/__tests__/Summary.spec.tsx @@ -1,7 +1,7 @@ import { render, screen } from '@testing-library/react'; import type { TBalanceResponse } from 'librechat-data-provider'; import type { BalanceState } from '~/hooks/useBalanceSummary'; -import { summarizeBalance } from '~/utils'; +import { summarizeBalance } from '~/utils/balance'; import Summary from '../Summary'; const NOW = Date.parse('2026-07-05T00:00:00.000Z'); @@ -64,6 +64,31 @@ describe('balance Summary', () => { expect(screen.getByTestId('balance-refill')).toHaveTextContent('Refills when you run out'); }); + it('shows reset timing without exposing the allowance in percent mode', () => { + const { container } = render(); + expect(screen.getByTestId('balance-refill')).toHaveTextContent('Usage resets in 3 days'); + expect(container).not.toHaveTextContent(/credits|5,000,000|3,100,000|\$/); + }); + + it('shows a due reset on the next request rather than waiting for exhaustion', () => { + render( + , + ); + expect(screen.getByTestId('balance-refill')).toHaveTextContent( + 'Usage resets on your next request', + ); + }); + + it('shows the reset target in currency mode', () => { + render(); + expect(screen.getByTestId('balance-refill')).toHaveTextContent('Resets to $5.00 in 3 days'); + }); + it('flags an empty balance', () => { render(); expect(screen.getByTestId('balance-status')).toHaveTextContent('Out of credits'); diff --git a/client/src/components/Nav/Settings/BillingControls.tsx b/client/src/components/Nav/Settings/BillingControls.tsx index be35792984b..652a12e9583 100644 --- a/client/src/components/Nav/Settings/BillingControls.tsx +++ b/client/src/components/Nav/Settings/BillingControls.tsx @@ -62,6 +62,7 @@ export function AutoRefill() { return ( ; /** When the next refill becomes eligible — the same date the balance summary counts down to */ nextRefill: Date | null; @@ -21,6 +26,7 @@ interface AutoRefillSettingsProps { } const AutoRefillSettings: React.FC = ({ + refillMode = 'add', lastRefill, nextRefill, refillAmount, @@ -29,6 +35,7 @@ const AutoRefillSettings: React.FC = ({ }) => { const localize = useLocalize(); const hour12 = useClockFormat(); + const reset = refillMode === 'reset'; const lastRefillDate = lastRefill ? new Date(lastRefill) : null; @@ -63,14 +70,22 @@ const AutoRefillSettings: React.FC = ({ return (
-

{localize('com_nav_balance_auto_refill_settings')}

+

+ {localize( + reset ? 'com_nav_balance_auto_reset_settings' : 'com_nav_balance_auto_refill_settings', + )} +

- {localize('com_nav_balance_last_refill')} + + {localize(reset ? 'com_nav_balance_last_reset' : 'com_nav_balance_last_refill')} + {lastRefillDate ? lastRefillDate.toLocaleString(undefined, { hour12 }) : '-'}
{refillAmount != null && (
- {localize('com_nav_balance_refill_amount')} + + {localize(reset ? 'com_nav_balance_reset_amount' : 'com_nav_balance_refill_amount')} + {refillAmount}
)} @@ -83,8 +98,15 @@ const AutoRefillSettings: React.FC = ({
- - + +
diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 1b468eb66e3..d58e3780d17 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -65,6 +65,15 @@ "com_ui_balance_no_refill": "Doesn't refill automatically", "com_ui_balance_refill_amount_in": "+{{0}} available {{1}}", "com_ui_balance_refill_amount_when_empty": "+{{0}} when you run out", + "com_ui_balance_reset_in": "Usage resets {{0}}", + "com_ui_balance_reset_due": "Usage resets on your next request", + "com_ui_balance_reset_amount_in": "Resets to {{0}} {{1}}", + "com_ui_balance_reset_amount_due": "Resets to {{0}} on your next request", + "com_nav_balance_auto_reset_settings": "Automatic reset settings", + "com_nav_balance_last_reset": "Last reset", + "com_nav_balance_reset_amount": "Reset allowance", + "com_nav_balance_next_reset": "Next reset", + "com_nav_balance_next_reset_info": "Once the interval has passed, reading your balance or sending a request replaces it with the configured allowance. Unused credits do not carry over. In-flight requests still count against the allowance.", "com_ui_balance_refill_in": "Refill available {{0}}", "com_ui_balance_refill_when_empty": "Refills when you run out", "com_ui_balance_usage_limit": "Usage limit", diff --git a/client/src/utils/balance.ts b/client/src/utils/balance.ts index ce36fd913b2..ff7cc94e083 100644 --- a/client/src/utils/balance.ts +++ b/client/src/utils/balance.ts @@ -1,5 +1,5 @@ import { CREDITS_PER_USD, getRefillEligibilityDate } from 'librechat-data-provider'; -import type { BalanceDisplay, TBalanceResponse } from 'librechat-data-provider'; +import type { BalanceDisplay, BalanceRefillMode, TBalanceResponse } from 'librechat-data-provider'; import type { CurrencyConfig } from './tokens'; import { formatCost } from './tokens'; @@ -23,6 +23,7 @@ export interface BalanceSummary { tone: BalanceTone; /** Credits the next auto-refill adds; null without auto-refill. */ refillAmount: number | null; + refillMode: BalanceRefillMode; /** When the next auto-refill becomes eligible; null without auto-refill. */ nextRefill: Date | null; } @@ -86,6 +87,7 @@ export function summarizeBalance( usedPercent, tone, refillAmount, + refillMode: balance.refillMode ?? 'add', nextRefill: balance.autoRefillEnabled ? getNextRefill(balance) : null, }; } diff --git a/librechat.example.yaml b/librechat.example.yaml index 0b69903cf79..a40971434ad 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -449,9 +449,11 @@ registration: # enabled: false # startBalance: 20000 # autoRefillEnabled: false -# refillIntervalValue: 30 +# refillIntervalValue: 30 # Reset mode requires a positive integer # refillIntervalUnit: 'days' # refillAmount: 10000 +# # Deploy reset-capable servers to every replica before selecting 'reset'. +# refillMode: 'add' # 'reset' replaces the allowance each period without rollover # reservationTtlMs: 1800000 # How long an in-flight request's reserved credits count if never released # display: 'credits' # credits | currency (uses interface.currency) | percent (share of the allotment used) diff --git a/packages/api/src/middleware/balance.spec.ts b/packages/api/src/middleware/balance.spec.ts index 20fd333c66f..f65db49426a 100644 --- a/packages/api/src/middleware/balance.spec.ts +++ b/packages/api/src/middleware/balance.spec.ts @@ -3,7 +3,7 @@ import { MongoMemoryServer } from 'mongodb-memory-server'; import { logger, balanceSchema } from '@librechat/data-schemas'; import type { NextFunction, Request as ServerRequest, Response as ServerResponse } from 'express'; import type { IBalance, IBalanceUpdate } from '@librechat/data-schemas'; -import { createSetBalanceConfig } from './balance'; +import { buildBalanceUpdateFields, createSetBalanceConfig } from './balance'; jest.mock('@librechat/data-schemas', () => ({ ...jest.requireActual('@librechat/data-schemas'), @@ -381,6 +381,41 @@ describe('createSetBalanceConfig', () => { // This should have fixed the issue - user should no longer get the error }); + test('disables a due stored reset before balance reads and admission', async () => { + const userId = new mongoose.Types.ObjectId(); + const lastRefill = new Date('2020-01-01'); + await Balance.create({ + user: userId, + tokenCredits: 500, + autoRefillEnabled: true, + refillMode: 'reset', + refillAmount: 1000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + lastRefill, + }); + const read = jest.fn(findBalanceByUser); + const middleware = createSetBalanceConfig({ + getAppConfig: jest.fn().mockResolvedValue({ + balance: { enabled: true, startBalance: 1000, autoRefillEnabled: false }, + }), + findBalanceByUser: read, + upsertBalanceFields, + }); + await middleware( + createMockRequest(userId) as ServerRequest, + createMockResponse() as ServerResponse, + mockNext, + ); + expect(read).toHaveBeenCalledTimes(1); + expect(read).toHaveBeenCalledWith(userId.toString(), { applyReset: false }); + expect(await findBalanceByUser(userId.toString())).toMatchObject({ + tokenCredits: 500, + autoRefillEnabled: false, + lastRefill, + }); + }); + test('should not set lastRefill when auto-refill is disabled', async () => { const userId = new mongoose.Types.ObjectId(); @@ -731,3 +766,21 @@ describe('createSetBalanceConfig', () => { ); }); }); + +describe('balance refill mode synchronization', () => { + const config = { + autoRefillEnabled: true, + refillAmount: 1000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks' as const, + }; + test('persists reset mode for existing balances', () => { + expect( + buildBalanceUpdateFields({ ...config, refillMode: 'reset' }, null, 'user-1'), + ).toMatchObject({ refillMode: 'reset' }); + }); + test('restores additive mode when the reset option is removed', () => { + const record = { ...config, refillMode: 'reset', lastRefill: new Date() } as IBalance; + expect(buildBalanceUpdateFields(config, record, 'user-1')).toEqual({ refillMode: 'add' }); + }); +}); diff --git a/packages/api/src/middleware/balance.ts b/packages/api/src/middleware/balance.ts index 44cb34d7e93..845403a07b0 100644 --- a/packages/api/src/middleware/balance.ts +++ b/packages/api/src/middleware/balance.ts @@ -1,4 +1,5 @@ import { logger } from '@librechat/data-schemas'; +import { isBalanceRefillDue } from 'librechat-data-provider'; import type { IBalanceUpdate, BalanceConfig, @@ -18,7 +19,10 @@ export interface BalanceMiddlewareOptions { tenantId?: string; refresh?: boolean; }) => Promise; - findBalanceByUser: (userId: string) => Promise; + findBalanceByUser: ( + userId: string, + options?: { applyReset?: boolean }, + ) => Promise; upsertBalanceFields: ( userId: string, fields: IBalanceUpdate, @@ -80,9 +84,14 @@ export function buildBalanceUpdateFields( config.autoRefillEnabled && config.refillIntervalValue != null && config.refillIntervalUnit != null && - config.refillAmount != null; + config.refillAmount != null && + (config.refillMode !== 'reset' || + (Number.isInteger(config.refillIntervalValue) && config.refillIntervalValue > 0)); if (!isAutoRefillConfigValid) { + if (userRecord?.autoRefillEnabled === true) { + updateFields.autoRefillEnabled = false; + } return updateFields; } @@ -98,6 +107,11 @@ export function buildBalanceUpdateFields( updateFields.refillIntervalUnit = config.refillIntervalUnit; } + const refillMode = config.refillMode ?? 'add'; + if ((userRecord?.refillMode ?? 'add') !== refillMode) { + updateFields.refillMode = refillMode; + } + if (userRecord?.refillAmount !== config.refillAmount) { updateFields.refillAmount = config.refillAmount; } @@ -147,11 +161,16 @@ export function createSetBalanceConfig({ } const userId = typeof user._id === 'string' ? user._id : user._id.toString(); await runBalanceUpdate(userId, async () => { - const userBalanceRecord = await findBalanceByUser(userId); + const userBalanceRecord = await findBalanceByUser(userId, { applyReset: false }); const updateFields = buildBalanceUpdateFields(balanceConfig, userBalanceRecord, userId); if (Object.keys(updateFields).length === 0) { - balanceLocals.balanceData = userBalanceRecord; + balanceLocals.balanceData = + balanceConfig.autoRefillEnabled === true && + userBalanceRecord?.refillMode === 'reset' && + isBalanceRefillDue(userBalanceRecord, new Date()) + ? await findBalanceByUser(userId) + : userBalanceRecord; return; } diff --git a/packages/api/src/middleware/checkBalance.spec.ts b/packages/api/src/middleware/checkBalance.spec.ts index d976f51cb05..4796ab36ba4 100644 --- a/packages/api/src/middleware/checkBalance.spec.ts +++ b/packages/api/src/middleware/checkBalance.spec.ts @@ -533,6 +533,80 @@ describe('checkBalance', () => { expect(refills).toHaveLength(1); }); + it.each([ + ['increased', 2000, true, 'reset' as const, 2000], + ['decreased', 400, true, 'reset' as const, 400], + ['disabled', 1000, false, 'reset' as const, 500], + ['changed to additive', 1000, true, 'add' as const, 500], + ])( + 'uses the %s current policy for direct admission without a balance read', + async (_case, refillAmount, autoRefillEnabled, refillMode, expectedCredits) => { + const user = new mongoose.Types.ObjectId().toString(); + const lastRefill = new Date('2020-01-01'); + await Balance.create({ + user, + tokenCredits: 500, + autoRefillEnabled: true, + refillMode: 'reset', + refillAmount: 1000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + lastRefill, + }); + const reservation = await checkBalance( + { req, res, txData: { ...baseTxData, user, amount: 100 } }, + realDeps({ + startBalance: 9999, + refillAmount, + autoRefillEnabled, + refillMode, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + }), + ); + await reservation.release(); + expect((await Balance.findOne({ user }).lean())?.tokenCredits).toBe(expectedCredits); + const ledger = await methods.getTransactions({ user, context: 'balanceReset' }); + expect(ledger).toHaveLength(autoRefillEnabled && refillMode === 'reset' ? 1 : 0); + }, + ); + + it('preserves the period and held credits when policy is synchronized before concurrent admissions', async () => { + const user = new mongoose.Types.ObjectId().toString(); + const lastRefill = new Date(); + await Balance.create({ + user, + tokenCredits: 500, + autoRefillEnabled: true, + refillMode: 'reset', + refillAmount: 1000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + lastRefill, + reservations: [{ id: 'held', amount: 400, expiresAt: new Date(Date.now() + 60_000) }], + reservedCredits: 400, + }); + const deps = realDeps({ + startBalance: 9999, + autoRefillEnabled: true, + refillMode: 'reset', + refillAmount: 2000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + }); + const admitted = admittedOf(await admitConcurrently(user, 10, deps, 100)); + expect(admitted).toHaveLength(1); + await Promise.all(admitted.map((reservation) => reservation.release())); + const record = await Balance.findOne({ user }).select('+reservedCredits').lean(); + expect(record).toMatchObject({ + tokenCredits: 500, + reservedCredits: 400, + refillAmount: 2000, + lastRefill, + }); + expect(await methods.getTransactions({ user, context: 'balanceReset' })).toHaveLength(0); + }); + it('lazily initializes a missing record and reserves against it', async () => { const user = new mongoose.Types.ObjectId().toString(); const deps = realDeps({ startBalance: 1000 }); diff --git a/packages/api/src/middleware/checkBalance.ts b/packages/api/src/middleware/checkBalance.ts index 070abca38fb..aaa849e7ede 100644 --- a/packages/api/src/middleware/checkBalance.ts +++ b/packages/api/src/middleware/checkBalance.ts @@ -15,6 +15,7 @@ import type { } from '@librechat/data-schemas'; import type { Response } from 'express'; import type { ServerRequest } from '~/types/http'; +import { buildBalanceUpdateFields } from './balance'; interface TxData { user: string; @@ -186,20 +187,7 @@ function buildInitialBalance(user: string, config?: BalanceConfig): IBalanceUpda if (config?.startBalance == null) { return undefined; } - const fields: IBalanceUpdate = { user, tokenCredits: config.startBalance }; - if ( - config.autoRefillEnabled && - config.refillIntervalValue != null && - config.refillIntervalUnit != null && - config.refillAmount != null - ) { - fields.autoRefillEnabled = config.autoRefillEnabled; - fields.refillIntervalValue = config.refillIntervalValue; - fields.refillIntervalUnit = config.refillIntervalUnit; - fields.refillAmount = config.refillAmount; - fields.lastRefill = new Date(); - } - return fields; + return buildBalanceUpdateFields(config, null, user); } /** @@ -237,12 +225,28 @@ export async function checkBalance( }); const expiresAt = new Date(Date.now() + ttlMs); + let refillPolicy: BalanceReservationRequest['refillPolicy']; + if (deps.balanceConfig != null) { + // Initialization fields must never restore credits or restart an existing period. + const { + user: _user, + tokenCredits: _credits, + lastRefill: _lastRefill, + ...policy + } = buildBalanceUpdateFields(deps.balanceConfig, null, user); + refillPolicy = { + autoRefillEnabled: false, + refillMode: deps.balanceConfig.refillMode ?? 'add', + ...policy, + }; + } const result = await deps.reserveBalance({ user, reservationId, amount: tokenCost, expiresAt, initialBalance: buildInitialBalance(user, deps.balanceConfig), + ...(refillPolicy != null ? { refillPolicy } : {}), }); if (result?.reserved) { diff --git a/packages/api/src/schedules/service.spec.ts b/packages/api/src/schedules/service.spec.ts index 5eb442581ac..2f16fc87283 100644 --- a/packages/api/src/schedules/service.spec.ts +++ b/packages/api/src/schedules/service.spec.ts @@ -1,5 +1,7 @@ -import { logger } from '@librechat/data-schemas'; +import mongoose from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; import { McpError, ErrorCode } from '@modelcontextprotocol/sdk/types.js'; +import { logger, createModels, createMethods } from '@librechat/data-schemas'; import type { SchedulesServiceDeps } from './service'; import { createSchedulesService, @@ -2743,3 +2745,87 @@ it('interrupts an A3 admission wait at shutdown without a model-visible policy r signal.mockRestore(); } }); + +describe('scheduled reset policy synchronization', () => { + let mongo: MongoMemoryServer; + beforeAll(async () => { + mongo = await MongoMemoryServer.create(); + createModels(mongoose); + await mongoose.connect(mongo.getUri()); + }); + afterAll(async () => { + await mongoose.disconnect(); + await mongo.stop(); + }); + beforeEach(async () => { + await mongoose.connection.dropDatabase(); + }); + + it.each([ + ['increased', 2000, true], + ['decreased', 400, true], + ['unchanged', 1000, true], + ['disabled', 1000, false], + ])( + 'applies the %s policy before resetting a scheduled balance', + async (_case, refillAmount, autoRefillEnabled) => { + const user = new mongoose.Types.ObjectId(); + const lastRefill = new Date('2020-01-01'); + await mongoose.models.Balance.create({ + user, + tokenCredits: 500, + autoRefillEnabled: true, + refillMode: 'reset', + refillAmount: 1000, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + lastRefill, + reservations: [{ id: 'held', amount: 450, expiresAt: new Date(Date.now() + 60_000) }], + reservedCredits: 450, + }); + const db = createMethods(mongoose); + const findBalance: SchedulesServiceDeps['findBalance'] = (userId, options) => + db.findBalanceByUser(userId, { includeReservedCredits: true, ...options }); + const service = createSchedulesService({ + methods: {} as SchedulesServiceDeps['methods'], + getAppConfig: async () => + ({ + balance: { + enabled: true, + startBalance: 1000, + autoRefillEnabled, + refillMode: 'reset', + refillAmount, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + }, + }) as Awaited>, + findBalance, + upsertBalance: (userId, { set, setOnInsert }) => + db.upsertBalanceFields(userId, set, setOnInsert), + initializeNullBalance: jest.fn(async () => null), + findUserById: jest.fn(async () => null), + preflightMCP: jest.fn().mockResolvedValue([]), + resolveAgentFireAccess: jest.fn(async () => 'ok' as const), + getChatProject: jest.fn(async () => null), + isUserDeleting: jest.fn(async () => false), + enqueueAgentTrigger: jest.fn(async () => undefined), + getTriggerDelivery: jest.fn(async () => null), + }); + await expect(service.engineDeps.isOutOfBalance({ id: user.toString() })).resolves.toBe( + autoRefillEnabled && refillAmount <= 450, + ); + const record = await db.findBalanceByUser(user.toString()); + expect(record?.tokenCredits).toBe(autoRefillEnabled ? refillAmount : 500); + expect(record?.lastRefill).toEqual(autoRefillEnabled ? expect.any(Date) : lastRefill); + const ledger = await mongoose.models.Transaction.find({ + user, + context: 'balanceReset', + }).lean(); + expect(ledger).toHaveLength(autoRefillEnabled ? 1 : 0); + if (autoRefillEnabled) { + expect(ledger[0].rawAmount).toBe(refillAmount - 500); + } + }, + ); +}); diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index d048d1c4175..5592aed9a69 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -1,5 +1,5 @@ +import { isBalanceRefillDue, Permissions, PermissionTypes } from 'librechat-data-provider'; import { logger, runAsSystem, tenantStorage, isRuntimeDisabled } from '@librechat/data-schemas'; -import { getRefillEligibilityDate, Permissions, PermissionTypes } from 'librechat-data-provider'; import { DEFAULT_SCHEDULE_MCP_CONSENT_LIFETIME_HOURS, isScheduleMCPAuthorizationFailure, @@ -169,7 +169,7 @@ export interface SchedulesServiceDeps { userId: string | Types.ObjectId, ) => Promise<{ _id: Types.ObjectId; tenantId?: string; role?: string } | null>; /** Reads the balance record together with the credits unexpired in-flight reservations hold. */ - findBalance: (userId: string) => Promise; + findBalance: (userId: string, options?: { applyReset?: boolean }) => Promise; /** * Upserts a balance record. `setOnInsert` carries fields that must ONLY apply to a * document this call creates — chiefly the starting credit — so a record created by a @@ -676,27 +676,7 @@ export function createSchedulesService( * mirroring the chat balance check's auto-refill eligibility (record-based). */ function isRefillEligible(record: IBalance | null | undefined): boolean { - if (record?.autoRefillEnabled !== true) { - return false; - } - if (!(typeof record.refillAmount === 'number' && record.refillAmount > 0)) { - return false; - } - const lastRefillDate = new Date(record.lastRefill ?? 0); - if (Number.isNaN(lastRefillDate.getTime())) { - return true; - } - // Mirror checkBalanceRecord's fallbacks exactly (interval 0 / 'days' when a - // partially-synced record is missing them) so we never pre-skip a record the - // interactive chat balance check would have refilled. - return ( - new Date() >= - getRefillEligibilityDate( - lastRefillDate, - record.refillIntervalValue ?? 0, - record.refillIntervalUnit ?? 'days', - ) - ); + return record != null && isBalanceRefillDue(record, new Date()); } const engineDeps: ScheduleEngineDeps = { @@ -726,11 +706,11 @@ export function createSchedulesService( if (balanceConfig?.enabled !== true) { return false; } - let record = await deps.findBalance(user.id); + let record = await deps.findBalance(user.id, { applyReset: false }); // Credits in-flight requests hold are unavailable to this fire as well: the chat // balance check admits against the unreserved amount. Taken from this read because // the initialization/sync writes below return the record without the total. - const reservedCredits = record?.reservedCredits ?? 0; + let reservedCredits = record?.reservedCredits ?? 0; // Initialize/sync the record exactly as the chat's balance middleware would, // so a new user's startBalance is applied before we read it (avoids skipping // a schedule that an interactive chat would have allowed). @@ -763,7 +743,8 @@ export function createSchedulesService( tokenCredits, sync: syncFields, }); - record = initialized ?? (await deps.findBalance(user.id)) ?? record; + record = + initialized ?? (await deps.findBalance(user.id, { applyReset: false })) ?? record; } else if (Object.keys(syncFields).length > 0) { // EXISTING record with a real credit: only refill-config sync remains, which // never touches `tokenCredits` and is safe to `$set` directly. @@ -771,6 +752,14 @@ export function createSchedulesService( } } } + if ( + balanceConfig.autoRefillEnabled === true && + record?.refillMode === 'reset' && + isBalanceRefillDue(record, new Date()) + ) { + record = await deps.findBalance(user.id, { applyReset: true }); + reservedCredits = record?.reservedCredits ?? 0; + } const credits = (record?.tokenCredits ?? 0) - reservedCredits; if (credits > 0) { return false; diff --git a/packages/api/src/types/balance.ts b/packages/api/src/types/balance.ts index cdf72feb07a..a47ceed7327 100644 --- a/packages/api/src/types/balance.ts +++ b/packages/api/src/types/balance.ts @@ -1,4 +1,4 @@ -import type { RefillIntervalUnit } from 'librechat-data-provider'; +import type { BalanceRefillMode, RefillIntervalUnit } from 'librechat-data-provider'; export interface BalanceUpdateFields { user?: string; @@ -7,5 +7,6 @@ export interface BalanceUpdateFields { refillIntervalValue?: number; refillIntervalUnit?: RefillIntervalUnit; refillAmount?: number; + refillMode?: BalanceRefillMode; lastRefill?: Date; } diff --git a/packages/data-provider/specs/config-schemas.spec.ts b/packages/data-provider/specs/config-schemas.spec.ts index f934fe79534..ea04d70335f 100644 --- a/packages/data-provider/specs/config-schemas.spec.ts +++ b/packages/data-provider/specs/config-schemas.spec.ts @@ -7,6 +7,7 @@ import { isAllDataRetention, isForcedTemporaryRetention, configSchema, + balanceSchema, interfaceSchema, fileStorageSchema, fileStrategiesSchema, @@ -1993,3 +1994,40 @@ describe('interfaceSchema agentSelectorLimit', () => { expect(normalizeAgentSelectorLimit('10')).toBe(10); }); }); + +describe('balance refill mode', () => { + test('keeps additive refill as the default for existing configurations', () => { + expect(balanceSchema.parse({ autoRefillEnabled: true }).refillMode).toBe('add'); + }); + test('accepts weekly non-accumulating reset with percentage display', () => { + expect( + balanceSchema.parse({ + enabled: true, + startBalance: 300000000, + autoRefillEnabled: true, + refillIntervalValue: 1, + refillIntervalUnit: 'weeks', + refillAmount: 300000000, + refillMode: 'reset', + display: 'percent', + }), + ).toMatchObject({ refillMode: 'reset', display: 'percent' }); + }); + test('rejects unsupported refill modes', () => { + expect(balanceSchema.safeParse({ refillMode: 'rollover' }).success).toBe(false); + }); +}); + +describe('reset interval validation', () => { + test.each([0, -1, 0.5, 1.5])( + 'rejects reset interval %s while preserving additive settings', + (refillIntervalValue) => { + expect(balanceSchema.safeParse({ refillMode: 'reset', refillIntervalValue }).success).toBe( + false, + ); + expect(balanceSchema.safeParse({ refillMode: 'add', refillIntervalValue }).success).toBe( + true, + ); + }, + ); +}); diff --git a/packages/data-provider/src/balance.ts b/packages/data-provider/src/balance.ts index 17e2f07db26..ce28af9957c 100644 --- a/packages/data-provider/src/balance.ts +++ b/packages/data-provider/src/balance.ts @@ -1,3 +1,5 @@ +import type { TBalanceResponse } from './types'; + export const REFILL_INTERVAL_UNITS = [ 'seconds', 'minutes', @@ -7,6 +9,10 @@ export const REFILL_INTERVAL_UNITS = [ 'months', ] as const; +export const BALANCE_REFILL_MODES = ['add', 'reset'] as const; + +export type BalanceRefillMode = (typeof BALANCE_REFILL_MODES)[number]; + export type RefillIntervalUnit = (typeof REFILL_INTERVAL_UNITS)[number]; /** How the UI presents a balance: raw credits, their currency value, or the share of the @@ -58,3 +64,42 @@ export function getRefillEligibilityDate( } } } + +/** Whether the configured refill/reset period has elapsed. */ +export function isBalanceRefillDue( + record: Pick< + TBalanceResponse, + | 'refillMode' + | 'autoRefillEnabled' + | 'refillAmount' + | 'lastRefill' + | 'refillIntervalValue' + | 'refillIntervalUnit' + >, + now: Date, +): boolean { + if (!record.autoRefillEnabled || !(record.refillAmount != null && record.refillAmount > 0)) { + return false; + } + if ( + record.refillMode === 'reset' && + !( + record.refillIntervalValue != null && + Number.isInteger(record.refillIntervalValue) && + record.refillIntervalValue > 0 + ) + ) { + return false; + } + const lastRefill = new Date(record.lastRefill ?? 0); + if (isNaN(lastRefill.getTime())) { + return true; + } + const eligibleAt = getRefillEligibilityDate( + lastRefill, + record.refillIntervalValue ?? 0, + record.refillIntervalUnit ?? 'days', + ); + // Date setters can round a fractional interval down to the same instant. + return (record.refillMode !== 'reset' || eligibleAt > lastRefill) && now >= eligibleAt; +} diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 5048aaee942..2471aa8ff00 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -21,6 +21,13 @@ import { MAX_PII_PATTERNS_PER_SOURCE, MAX_PII_PATTERN_LENGTH, } from './filters'; +import { + REFILL_INTERVAL_UNITS, + BALANCE_DISPLAY_MODES, + BALANCE_REFILL_MODES, + MIN_BALANCE_RESERVATION_TTL_MS, + DEFAULT_BALANCE_RESERVATION_TTL_MS, +} from './balance'; import { EModelEndpoint, eModelEndpointSchema, @@ -28,12 +35,6 @@ import { eReasoningParameterFormatSchema, eReasoningResponseKeySchema, } from './schemas'; -import { - REFILL_INTERVAL_UNITS, - BALANCE_DISPLAY_MODES, - MIN_BALANCE_RESERVATION_TTL_MS, - DEFAULT_BALANCE_RESERVATION_TTL_MS, -} from './balance'; import { scheduledMCPResourceBindingSchema, scheduledMCPReadOnlyPolicySchema, @@ -3383,22 +3384,37 @@ export const ocrSchema = z.object({ strategy: z.nativeEnum(OCRStrategy).default(OCRStrategy.MISTRAL_OCR), }); -export const balanceSchema = z.object({ - enabled: z.boolean().optional().default(false), - startBalance: z.number().optional().default(20000), - autoRefillEnabled: z.boolean().optional().default(false), - refillIntervalValue: z.number().optional().default(30), - refillIntervalUnit: z.enum(REFILL_INTERVAL_UNITS).optional().default('days'), - refillAmount: z.number().optional().default(10000), - reservationTtlMs: z - .number() - .int() - .min(MIN_BALANCE_RESERVATION_TTL_MS) - .optional() - .default(DEFAULT_BALANCE_RESERVATION_TTL_MS), - /** How the UI presents the balance; `credits` keeps the raw figure. */ - display: z.enum(BALANCE_DISPLAY_MODES).optional().default('credits'), -}); +export const balanceSchema = z + .object({ + enabled: z.boolean().optional().default(false), + startBalance: z.number().optional().default(20000), + autoRefillEnabled: z.boolean().optional().default(false), + refillIntervalValue: z.number().optional().default(30), + refillIntervalUnit: z.enum(REFILL_INTERVAL_UNITS).optional().default('days'), + refillAmount: z.number().optional().default(10000), + /** Add credits on exhaustion, or replace the allowance on the next read/request once due. */ + refillMode: z.enum(BALANCE_REFILL_MODES).optional().default('add'), + reservationTtlMs: z + .number() + .int() + .min(MIN_BALANCE_RESERVATION_TTL_MS) + .optional() + .default(DEFAULT_BALANCE_RESERVATION_TTL_MS), + /** How the UI presents the balance; `credits` keeps the raw figure. */ + display: z.enum(BALANCE_DISPLAY_MODES).optional().default('credits'), + }) + .superRefine((balance, ctx) => { + if ( + balance.refillMode === 'reset' && + !(Number.isInteger(balance.refillIntervalValue) && balance.refillIntervalValue > 0) + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + path: ['refillIntervalValue'], + message: 'Reset intervals must be positive integers', + }); + } + }); export const transactionsSchema = z.object({ enabled: z.boolean().optional().default(true), diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index dc33077a1c6..2b51072af56 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -21,10 +21,10 @@ import type { CodeEnvironmentUserSettings, TAgentsEndpoint, } from './config'; +import type { BalanceRefillMode, RefillIntervalUnit } from './balance'; import type { StatefulCodeEnvironment } from './stateful-code'; import type { CodeApprovalMode } from './code/approval'; import type { EToolResources } from './types/tools'; -import type { RefillIntervalUnit } from './balance'; import type { SettingDefinition } from './generate'; import type { TMinimalFeedback } from './feedback'; import type { ContentTypes } from './types/runs'; @@ -1257,6 +1257,7 @@ export type TBalanceResponse = { refillIntervalUnit?: RefillIntervalUnit; lastRefill?: Date | string; refillAmount?: number; + refillMode?: BalanceRefillMode; }; /* -------------------------------------------------------------------------- */ diff --git a/packages/data-schemas/src/methods/transaction.spec.ts b/packages/data-schemas/src/methods/transaction.spec.ts index 181ae960ca4..9418da2a301 100644 --- a/packages/data-schemas/src/methods/transaction.spec.ts +++ b/packages/data-schemas/src/methods/transaction.spec.ts @@ -1568,6 +1568,129 @@ describe('Balance Reservations', () => { lastRefill: new Date('2020-01-01T00:00:00.000Z'), }; + describe('reset mode', () => { + const resettable = { ...refillable, refillMode: 'reset' as const }; + + test.each([500, 1500, -100])( + 'replaces %s credits with one allowance on read', + async (credits) => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: credits }); + + const balance = await findBalanceByUser(user.toString()); + expect(balance?.tokenCredits).toBe(1000); + expect(balance?.refillMode).toBe('reset'); + expect(balance).not.toHaveProperty('pendingRefill'); + expect(balance).not.toHaveProperty('reservations'); + expect(balance).not.toHaveProperty('reservedCredits'); + const resets = await Transaction.find({ user, context: 'balanceReset' }).lean(); + expect(resets).toHaveLength(1); + expect(resets[0].rawAmount).toBe(1000 - credits); + expect(resets[0].tokenValue).toBe(1000 - credits); + }, + ); + + test.each([0, -1, 0.0001, 0.5, 1.5])( + 'does not grant repeated allowances for invalid interval %s', + async (refillIntervalValue) => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, refillIntervalValue, tokenCredits: 500 }); + await findBalanceByUser(user.toString()); + await reserve(user.toString(), 100); + await findBalanceByUser(user.toString()); + expect((await readState(user))?.tokenCredits).toBe(500); + expect(await Transaction.countDocuments({ user })).toBe(0); + }, + ); + + test('disabled resets remain disabled on reads and admission', async () => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: 500 }); + await upsertBalanceFields(user.toString(), { autoRefillEnabled: false }); + expect((await findBalanceByUser(user.toString()))?.tokenCredits).toBe(500); + expect(await reserve(user.toString(), 100)).toEqual({ reserved: true, balance: 500 }); + expect(await Transaction.countDocuments({ user })).toBe(0); + }); + + test('resets before admission even when the remaining balance covers the request', async () => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: 500 }); + expect(await reserve(user.toString(), 100)).toEqual({ reserved: true, balance: 1000 }); + expect((await readState(user))?.tokenCredits).toBe(1000); + }); + + test('preserves reservations and applies one reset across concurrent reads and admissions', async () => { + const user = new mongoose.Types.ObjectId(); + const heldId = newId(); + await Balance.create({ + user, + ...resettable, + tokenCredits: 1500, + reservedCredits: 700, + reservations: [{ id: heldId, amount: 700, expiresAt: new Date(Date.now() + 60_000) }], + }); + await Promise.all([ + findBalanceByUser(user.toString()), + findBalanceByUser(user.toString()), + reserve(user.toString(), 200), + ]); + const stored = await readState(user); + expect(stored?.tokenCredits).toBe(1000); + expect(stored?.reservedCredits).toBe(900); + expect(stored?.reservations?.some((held) => held.id === heldId)).toBe(true); + expect(await Transaction.countDocuments({ user, context: 'balanceReset' })).toBe(1); + }); + + test('does not reset again before the next period or accumulate missed periods', async () => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: 400 }); + await findBalanceByUser(user.toString()); + await Balance.updateOne({ user }, { $set: { tokenCredits: 600 } }); + expect((await findBalanceByUser(user.toString()))?.tokenCredits).toBe(600); + expect(await Transaction.countDocuments({ user, context: 'balanceReset' })).toBe(1); + }); + + test('replays a failed reset ledger write without resetting or duplicating the allowance', async () => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: 1500 }); + const save = jest + .spyOn(Transaction.prototype, 'save') + .mockRejectedValue(new Error('ledger unavailable')); + expect((await findBalanceByUser(user.toString()))?.tokenCredits).toBe(1000); + save.mockRestore(); + expect((await findBalanceByUser(user.toString()))?.tokenCredits).toBe(1000); + expect((await readState(user))?.pendingRefill).toBeUndefined(); + const resets = await Transaction.find({ user, context: 'balanceReset' }).lean(); + expect(resets).toHaveLength(1); + expect(resets[0].rawAmount).toBe(-500); + }); + + test('synchronizes a changed allowance before applying a due reset', async () => { + const user = new mongoose.Types.ObjectId(); + await Balance.create({ user, ...resettable, tokenCredits: 500 }); + expect( + (await findBalanceByUser(user.toString(), { applyReset: false }))?.tokenCredits, + ).toBe(500); + expect(await Transaction.countDocuments({ user })).toBe(0); + expect( + (await upsertBalanceFields(user.toString(), { refillAmount: 2000 }))?.tokenCredits, + ).toBe(2000); + const resets = await Transaction.find({ user, context: 'balanceReset' }).lean(); + expect(resets).toHaveLength(1); + expect(resets[0].rawAmount).toBe(1500); + }); + + 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 }); + interleaveBeforeNextWrite(() => + Balance.collection.updateOne({ user }, { $set: { refillMode: 'add' } }), + ); + expect(await reserve(user.toString(), 100)).toEqual({ reserved: true, balance: 500 }); + expect(await Transaction.countDocuments({ user })).toBe(0); + }); + }); + test('applies one refill per eligibility window across concurrent admissions', async () => { const user = new mongoose.Types.ObjectId(); await Balance.create({ user, ...refillable }); diff --git a/packages/data-schemas/src/methods/transaction.ts b/packages/data-schemas/src/methods/transaction.ts index 736d6720d74..b0ea64483ab 100644 --- a/packages/data-schemas/src/methods/transaction.ts +++ b/packages/data-schemas/src/methods/transaction.ts @@ -1,4 +1,4 @@ -import { getRefillEligibilityDate } from 'librechat-data-provider'; +import { isBalanceRefillDue } from 'librechat-data-provider'; import type { AnyBulkWriteOperation, FilterQuery, Model, Types } from 'mongoose'; import type { BalanceReservationRequest, @@ -99,7 +99,7 @@ export function createTransactionMethods( bulkInsertTransactions: (docs: TransactionData[]) => Promise; findBalanceByUser: ( user: string, - options?: { includeReservedCredits?: boolean }, + options?: { includeReservedCredits?: boolean; applyReset?: boolean }, ) => Promise; upsertBalanceFields: ( user: string, @@ -292,24 +292,6 @@ export function createTransactionMethods( ); } - function isAutoRefillDue(record: IBalance, now: Date): boolean { - if (!record.autoRefillEnabled || !(record.refillAmount > 0)) { - return false; - } - const lastRefill = new Date(record.lastRefill ?? 0); - if (isNaN(lastRefill.getTime())) { - return true; - } - return ( - now >= - getRefillEligibilityDate( - lastRefill, - record.refillIntervalValue ?? 0, - record.refillIntervalUnit ?? 'days', - ) - ); - } - function isDuplicateKeyError(error: unknown): boolean { return error instanceof Error && 'code' in error && (error as { code: number }).code === 11000; } @@ -322,7 +304,7 @@ export function createTransactionMethods( async function settleAutoRefill( balanceId: unknown, user: Types.ObjectId, - { transactionId, rawAmount }: IBalancePendingRefill, + { transactionId, rawAmount, context }: IBalancePendingRefill, ): Promise { try { const Transaction = mongoose.models.Transaction; @@ -330,7 +312,7 @@ export function createTransactionMethods( _id: transactionId, user, tokenType: 'credits', - context: 'autoRefill', + context: context ?? 'autoRefill', rawAmount, }); calculateTokenValue(transaction); @@ -361,9 +343,14 @@ export function createTransactionMethods( */ async function applyAutoRefill(record: IBalance, now: Date): Promise { const Balance = mongoose.models.Balance as Model; + const reset = record.refillMode === 'reset'; + const credits = reset + ? record.refillAmount + : Math.max(0, (record.tokenCredits ?? 0) + record.refillAmount); const pendingRefill: IBalancePendingRefill = { transactionId: new mongoose.Types.ObjectId(), - rawAmount: record.refillAmount, + rawAmount: reset ? credits - (record.tokenCredits ?? 0) : record.refillAmount, + ...(reset ? { context: 'balanceReset' as const } : {}), }; const result = await Balance.updateOne( { @@ -374,12 +361,13 @@ export function createTransactionMethods( pendingRefill: null, autoRefillEnabled: record.autoRefillEnabled, refillAmount: record.refillAmount, + refillMode: record.refillMode ?? null, refillIntervalValue: record.refillIntervalValue ?? null, refillIntervalUnit: record.refillIntervalUnit ?? null, }, { $set: { - tokenCredits: Math.max(0, (record.tokenCredits ?? 0) + record.refillAmount), + tokenCredits: credits, lastRefill: now, pendingRefill, }, @@ -471,7 +459,9 @@ export function createTransactionMethods( * the value read and `reservedCredits` still leaves room for the amount, so concurrent * admissions of a funded balance all commit, and an admission that lost the credits re-reads. * A due auto-refill is applied first by its own fenced write, at most once per admission, so a - * refill interval that is due again the instant it is applied still refills once. Returns null + * refill interval that is due again the instant it is applied still refills once. Reset mode + * replaces the allowance whenever due, even before exhaustion, while retaining reservations. + * Returns null * when the user has no balance record and no `initialBalance` was given. */ async function reserveBalance({ @@ -480,6 +470,7 @@ export function createTransactionMethods( amount, expiresAt, initialBalance, + refillPolicy, }: BalanceReservationRequest): Promise { const Balance = mongoose.models.Balance as Model; let delay = 10; @@ -498,6 +489,19 @@ export function createTransactionMethods( continue; } + if ( + refillPolicy != null && + (record.autoRefillEnabled !== refillPolicy.autoRefillEnabled || + (refillPolicy.autoRefillEnabled === true && + ((record.refillMode ?? 'add') !== (refillPolicy.refillMode ?? 'add') || + record.refillAmount !== refillPolicy.refillAmount || + record.refillIntervalValue !== refillPolicy.refillIntervalValue || + record.refillIntervalUnit !== refillPolicy.refillIntervalUnit))) + ) { + await upsertBalanceRecord(user, refillPolicy); + continue; + } + const now = new Date(); const refillSettled = record.pendingRefill == null || @@ -513,7 +517,12 @@ export function createTransactionMethods( const credits = record.tokenCredits ?? 0; const balance = credits - (record.reservedCredits ?? 0); - if (refillSettled && !refilled && balance - amount <= 0 && isAutoRefillDue(record, now)) { + if ( + refillSettled && + !refilled && + (record.refillMode === 'reset' || balance - amount <= 0) && + isBalanceRefillDue(record, now) + ) { refilled = await applyAutoRefill(record, now); continue; } @@ -685,23 +694,40 @@ export function createTransactionMethods( /** * Retrieves a user's balance record. With `includeReservedCredits`, `reservedCredits` is the * total of the reservations that have not expired, so a reservation left by a crashed request - * stops counting at its expiry even before a reservation write prunes it. + * stops counting at its expiry even before a reservation write prunes it. A due reset is + * applied before returning the reading; inactive periods do not accumulate allowances. */ async function findBalanceByUser( user: string, - options?: { includeReservedCredits?: boolean }, + options?: { includeReservedCredits?: boolean; applyReset?: boolean }, ): Promise { const Balance = mongoose.models.Balance as Model; - const query = Balance.findOne({ user }).sort(oldestFirst); - if (!options?.includeReservedCredits) { - return query.lean(); - } - const record = await query.select('+reservations').lean(); + const read = () => + Balance.findOne({ user }) + .sort(oldestFirst) + .select(options?.includeReservedCredits ? '+reservations +pendingRefill' : '+pendingRefill') + .lean(); + let record = await read(); if (!record) { return null; } + if ( + options?.applyReset !== false && + record.refillMode === 'reset' && + (record.pendingRefill != null || isBalanceRefillDue(record, new Date())) + ) { + // Reuse the fenced admission path without holding any credits for a balance read. + await reserveBalance({ user, reservationId: '', amount: 0, expiresAt: new Date() }); + record = await read(); + if (!record) { + return null; + } + } const now = new Date(); - const { reservations, ...balance } = record; + const { reservations, reservedCredits: _held, pendingRefill: _pending, ...balance } = record; + if (!options?.includeReservedCredits) { + return balance as IBalance; + } const reservedCredits = (reservations ?? []).reduce( (sum, reservation) => (reservation.expiresAt > now ? sum + reservation.amount : sum), 0, @@ -715,7 +741,11 @@ export function createTransactionMethods( fields: IBalanceUpdate, insertOnly?: IBalanceUpdate, ): Promise { - return upsertBalanceRecord(user, fields, insertOnly); + const record = await upsertBalanceRecord(user, fields, insertOnly); + if (record?.refillMode === 'reset' && isBalanceRefillDue(record, new Date())) { + return findBalanceByUser(user); + } + return record; } /** Deletes transactions matching a filter. */ diff --git a/packages/data-schemas/src/methods/user.ts b/packages/data-schemas/src/methods/user.ts index 7c31842fdb8..5237d415ec4 100644 --- a/packages/data-schemas/src/methods/user.ts +++ b/packages/data-schemas/src/methods/user.ts @@ -1,9 +1,9 @@ import mongoose, { FilterQuery } from 'mongoose'; -import { - AUTH_USER_DOC_CACHE_TTL_MS, - CacheKeys, - type RefillIntervalUnit, - type StatefulCodeEnvironment, +import { AUTH_USER_DOC_CACHE_TTL_MS, CacheKeys } from 'librechat-data-provider'; +import type { + BalanceRefillMode, + RefillIntervalUnit, + StatefulCodeEnvironment, } from 'librechat-data-provider'; import type { IUser, BalanceConfig, CreateUserRequest, UserDeleteResult } from '~/types'; import type { TwoFactorEnrollmentGuard, TwoFactorEnrollmentUpdate } from '~/types'; @@ -298,6 +298,7 @@ export function createUserMethods( refillIntervalValue: number; refillIntervalUnit: RefillIntervalUnit; refillAmount: number; + refillMode?: BalanceRefillMode; }; } = { $inc: { tokenCredits: balanceConfig.startBalance }, @@ -314,6 +315,7 @@ export function createUserMethods( refillIntervalValue: balanceConfig.refillIntervalValue, refillIntervalUnit: balanceConfig.refillIntervalUnit, refillAmount: balanceConfig.refillAmount, + refillMode: balanceConfig.refillMode ?? 'add', }; } diff --git a/packages/data-schemas/src/schema/balance.ts b/packages/data-schemas/src/schema/balance.ts index 1191098f48c..48f672679ea 100644 --- a/packages/data-schemas/src/schema/balance.ts +++ b/packages/data-schemas/src/schema/balance.ts @@ -1,5 +1,5 @@ import { Schema } from 'mongoose'; -import { REFILL_INTERVAL_UNITS } from 'librechat-data-provider'; +import { BALANCE_REFILL_MODES, REFILL_INTERVAL_UNITS } from 'librechat-data-provider'; import type * as t from '~/types'; const balanceSchema: Schema = new Schema({ @@ -37,6 +37,11 @@ const balanceSchema: Schema = new Schema({ type: Number, default: 0, }, + refillMode: { + type: String, + enum: BALANCE_REFILL_MODES, + default: 'add', + }, tenantId: { type: String, index: true, @@ -62,6 +67,7 @@ const balanceSchema: Schema = new Schema({ type: { transactionId: { type: Schema.Types.ObjectId, required: true }, rawAmount: { type: Number, required: true }, + context: { type: String, enum: ['autoRefill', 'balanceReset'] }, }, _id: false, default: undefined, diff --git a/packages/data-schemas/src/types/balance.ts b/packages/data-schemas/src/types/balance.ts index e52b4b3f6b4..48c75f461c1 100644 --- a/packages/data-schemas/src/types/balance.ts +++ b/packages/data-schemas/src/types/balance.ts @@ -1,4 +1,4 @@ -import type { RefillIntervalUnit } from 'librechat-data-provider'; +import type { BalanceRefillMode, RefillIntervalUnit } from 'librechat-data-provider'; import type { Document, Types } from 'mongoose'; /** Whole credits held against a balance while the request that reserved them is in flight */ @@ -12,6 +12,7 @@ export interface IBalanceReservation { export interface IBalancePendingRefill { transactionId: Types.ObjectId; rawAmount: number; + context?: 'autoRefill' | 'balanceReset'; } export interface IBalance extends Document { @@ -23,6 +24,7 @@ export interface IBalance extends Document { refillIntervalUnit: RefillIntervalUnit; lastRefill: Date; refillAmount: number; + refillMode?: BalanceRefillMode; tenantId?: string; /** Reservation state is excluded from reads unless explicitly selected */ reservations?: IBalanceReservation[]; @@ -39,6 +41,7 @@ export interface IBalanceUpdate { refillIntervalValue?: number; refillIntervalUnit?: RefillIntervalUnit; refillAmount?: number; + refillMode?: BalanceRefillMode; lastRefill?: Date; } @@ -53,6 +56,14 @@ export interface BalanceReservationRequest { expiresAt: Date; /** Creates the balance record with these fields when the user has none */ initialBalance?: IBalanceUpdate; + /** Current request policy; synchronized before existing-record admission without changing credits. */ + refillPolicy?: Pick< + IBalanceUpdate, + 'refillAmount' | 'refillIntervalValue' | 'refillIntervalUnit' + > & { + autoRefillEnabled: boolean; + refillMode: BalanceRefillMode; + }; } export interface BalanceReservationRenewal { diff --git a/packages/data-schemas/src/types/user.ts b/packages/data-schemas/src/types/user.ts index c589705e5e1..138de99ebc7 100644 --- a/packages/data-schemas/src/types/user.ts +++ b/packages/data-schemas/src/types/user.ts @@ -1,6 +1,7 @@ import type { TUserFavorite, RefillIntervalUnit, + BalanceRefillMode, StatefulCodeEnvironment, } from 'librechat-data-provider'; import type { Document, Types } from 'mongoose'; @@ -129,6 +130,7 @@ export interface BalanceConfig { refillIntervalValue?: number; refillIntervalUnit?: RefillIntervalUnit; refillAmount?: number; + refillMode?: BalanceRefillMode; reservationTtlMs?: number; }