fix: initialize a scheduled fire's balance with setOnInsert

The balance READ and the upsert that follows it are separate statements, and the
upsert applied every field with $set. A concurrent charge that created the record
in that window was therefore overwritten with a fresh starting balance, handing
back credits the user had already spent.

buildBalanceUpdateFields mixes two kinds of field, so a blanket $setOnInsert
would have been wrong too: the credit fields are initialization values, while the
refill-config fields are a genuine sync that must keep updating an existing
record. The two are now split, and only when the record was ABSENT — an existing
record with a null tokenCredits has no charge to clobber, and $setOnInsert would
never fire for it.

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,

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<ReturnType<SchedulesServiceDeps['getAppConfig']>>;
+
+  function serviceWithBalance(existing: Record<string, unknown> | 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<string, unknown>; setOnInsert: Record<string, unknown> },
+      ]
+    )[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

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<IBalance | null>;
-  upsertBalance: (userId: string, fields: BalanceUpdateFields) => Promise<IBalance | null>;
+  /**
+   * 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<BalanceUpdateFields>; setOnInsert: Partial<BalanceUpdateFields> },
+  ) => Promise<IBalance | null>;
   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;