diff --git a/api/server/services/Schedules/index.js b/api/server/services/Schedules/index.js index 0460265683..54726a2efa 100644 --- a/api/server/services/Schedules/index.js +++ b/api/server/services/Schedules/index.js @@ -10,10 +10,15 @@ const service = createSchedulesService({ findUserById: (userId) => mongoose.models.User.findById(userId).select('_id tenantId role').lean(), findBalance: (userId) => mongoose.models.Balance.findOne({ user: userId }).lean(), - upsertBalance: (userId, fields) => + upsertBalance: (userId, { set, setOnInsert }) => mongoose.models.Balance.findOneAndUpdate( { user: userId }, - { $set: fields }, + { + ...(set && Object.keys(set).length > 0 ? { $set: set } : {}), + ...(setOnInsert && Object.keys(setOnInsert).length > 0 + ? { $setOnInsert: setOnInsert } + : {}), + }, { upsert: true, new: true }, ).lean(), resolveAgentFireAccess, diff --git a/packages/api/src/schedules/service.spec.ts b/packages/api/src/schedules/service.spec.ts index 8fd077e876..db308989f9 100644 --- a/packages/api/src/schedules/service.spec.ts +++ b/packages/api/src/schedules/service.spec.ts @@ -54,6 +54,61 @@ const run = (): ActiveRun => ({ conversationId: 'c1', }); +describe('balance initialization', () => { + const balanceConfig = { + interfaceConfig: {}, + balance: { enabled: true, startBalance: 20000 }, + } as unknown as Awaited>; + + function serviceWithBalance(existing: Record | null) { + const upsertBalance = jest.fn(async () => ({ tokenCredits: 20000 })); + const service = createSchedulesService({ + methods: {} as unknown as SchedulesServiceDeps['methods'], + getAppConfig: (async () => balanceConfig) as SchedulesServiceDeps['getAppConfig'], + findUserById: jest.fn(async () => null), + findBalance: jest.fn(async () => existing), + upsertBalance, + resolveAgentFireAccess: jest.fn(async () => 'ok' as const), + isUserDeleting: jest.fn(async () => false), + } as unknown as SchedulesServiceDeps); + return { service, upsertBalance }; + } + + const updateFrom = (spy: jest.Mock) => + ( + spy.mock.calls[0] as unknown as [ + string, + { set: Record; setOnInsert: Record }, + ] + )[1]; + + /** + * The balance READ and this write are separate statements. A concurrent charge that + * creates the record in between would be overwritten by a blind `$set`, handing back + * credits the user had already spent. + */ + it('initializes the starting credit via setOnInsert, never $set', async () => { + const { service, upsertBalance } = serviceWithBalance(null); + + await service.engineDeps.isOutOfBalance({ id: 'user-1' } as never); + + expect(upsertBalance).toHaveBeenCalledTimes(1); + const update = updateFrom(upsertBalance); + expect(update.setOnInsert).toMatchObject({ tokenCredits: 20000 }); + expect(update.set).not.toHaveProperty('tokenCredits'); + }); + + it('still $sets a null credit on an EXISTING record, which has no charge to clobber', async () => { + const { service, upsertBalance } = serviceWithBalance({ autoRefillEnabled: false }); + + await service.engineDeps.isOutOfBalance({ id: 'user-1' } as never); + + const update = updateFrom(upsertBalance); + expect(update.set).toMatchObject({ tokenCredits: 20000 }); + expect(update.setOnInsert).toEqual({}); + }); +}); + describe('deleteScheduleForOwner', () => { /** * markScheduleDeleting runs FIRST and is one-shot: it matches only a not-yet-deleting diff --git a/packages/api/src/schedules/service.ts b/packages/api/src/schedules/service.ts index 0abf45e43c..eb1b8856b5 100644 --- a/packages/api/src/schedules/service.ts +++ b/packages/api/src/schedules/service.ts @@ -112,7 +112,15 @@ export interface SchedulesServiceDeps { userId: string | Types.ObjectId, ) => Promise<{ _id: Types.ObjectId; tenantId?: string; role?: string } | null>; findBalance: (userId: string) => Promise; - upsertBalance: (userId: string, fields: BalanceUpdateFields) => 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 + * concurrent charge is never overwritten with a fresh balance. + */ + upsertBalance: ( + userId: string, + update: { set: Partial; setOnInsert: Partial }, + ) => Promise; resolveAgentFireAccess: ( agentId: string, user: ScheduleUserContext, @@ -308,7 +316,26 @@ export function createSchedulesService(deps: SchedulesServiceDeps): SchedulesSer if (balanceConfig.startBalance != null) { const updateFields = buildBalanceUpdateFields(balanceConfig, record, user.id); if (Object.keys(updateFields).length > 0) { - record = await deps.upsertBalance(user.id, updateFields); + // The read above and this write are separate statements, and the credit + // fields are INITIALIZATION values: a concurrent charge that created the + // record in between would be overwritten by a blind `$set`, restoring credits + // the user had already spent. Route those through `$setOnInsert` so they only + // apply to a document this call actually creates. The refill-config fields are + // a genuine sync and stay on `$set`. + // + // Only when the record was ABSENT: an existing record with a null + // tokenCredits has no charge to clobber, so initializing it via `$set` is + // both safe and necessary (`$setOnInsert` would never fire for it). + const { user: initUser, tokenCredits, ...syncFields } = updateFields; + const setOnInsert = + record == null + ? { + ...(initUser != null ? { user: initUser } : {}), + ...(tokenCredits != null ? { tokenCredits } : {}), + } + : {}; + const set = record == null ? syncFields : updateFields; + record = await deps.upsertBalance(user.id, { set, setOnInsert }); } } const credits = record?.tokenCredits ?? 0;