⚛️ fix: Make Ledger, Plugin, Memory and Preset Writes Atomic

packages/api/src/agents/transactions.spec.ts

@@ -398,6 +398,33 @@ describe('bulkWriteTransactions (real DB)', () => {
     const bal = await Balance.findOne({ user: testUserId }).lean();
     expect(bal).toBeNull();
   });
+
+  it('should not debit the balance when inserting the ledger entries fails', async () => {
+    await Balance.create({ user: testUserId, tokenCredits: 1000 });
+    const docs: PreparedEntry[] = [
+      {
+        doc: { user: testUserId, conversationId: 'c1', tokenType: 'prompt', tokenValue: -100 },
+        tokenValue: -100,
+        balance: { enabled: true },
+      },
+    ];
+    const updateBalance = jest.fn(dbMethods.updateBalance);
+    const dbOps = {
+      insertMany: jest.fn().mockRejectedValue(new Error('insert failed')),
+      updateBalance,
+    };
+
+    await expect(bulkWriteTransactions({ user: testUserId, docs }, dbOps)).rejects.toThrow(
+      'insert failed',
+    );
+
+    expect(updateBalance).not.toHaveBeenCalled();
+    const bal = (await Balance.findOne({ user: testUserId }).lean()) as Record<
+      string,
+      unknown
+    > | null;
+    expect(bal!.tokenCredits).toBe(1000);
+  });
 });
 
 describe('end-to-end: prepare → bulk write → verify', () => {

packages/api/src/agents/transactions.ts

@@ -371,9 +371,13 @@ export async function bulkWriteTransactions(
     return doc;
   });
 
+  /**
+   * Ledger first, then balance (same order as `createTransaction`): a failed insert
+   * must never leave a debit without its transaction records.
+   */
+  await dbOps.insertMany(plainDocs);
+
   if (balanceEnabled) {
     await dbOps.updateBalance({ user, incrementValue: totalTokenValue });
   }
-
-  await dbOps.insertMany(plainDocs);
 }

packages/data-schemas/src/methods/memory.spec.ts

@@ -239,6 +239,20 @@ describe('memory partitions', () => {
     expect(created.ok).toBe(true);
   });
 
+  it('createMemory never overwrites an existing entry and stores one document per key', async () => {
+    await methods.createMemory({ userId, key: 'context', value: 'original', tokenCount: 1 });
+    const results = await Promise.allSettled([
+      methods.createMemory({ userId, key: 'context', value: 'dupe 1', tokenCount: 9 }),
+      methods.createMemory({ userId, key: 'context', value: 'dupe 2', tokenCount: 9 }),
+    ]);
+
+    expect(results.every((result) => result.status === 'rejected')).toBe(true);
+    const stored = await MemoryEntry.find({ userId, key: 'context' }).lean();
+    expect(stored).toHaveLength(1);
+    expect(stored[0].value).toBe('original');
+    expect(stored[0].tokenCount).toBe(1);
+  });
+
   it('formats memories per partition with partition-scoped totals', async () => {
     await methods.setMemory({ userId, key: 'one', value: 'personal one', tokenCount: 5 });
     await methods.setMemory({ userId, key: 'two', value: 'agent two', tokenCount: 7, agentId });

packages/data-schemas/src/methods/memory.ts

@@ -57,24 +57,23 @@ export function createMemoryMethods(mongoose: typeof import('mongoose')): {
       }
 
       const MemoryEntry = mongoose.models.MemoryEntry;
-      const existingMemory = await MemoryEntry.findOne({
-        userId,
-        key,
-        ...partitionFilter(agentId),
-      });
-      if (existingMemory) {
+      /** Insert-only upsert: one round trip decides existence, unlike a find-then-create. */
+      const result = await MemoryEntry.updateOne(
+        { userId, key, ...partitionFilter(agentId) },
+        {
+          $setOnInsert: {
+            value,
+            tokenCount,
+            ...(agentId ? { agentId } : {}),
+            updated_at: new Date(),
+          },
+        },
+        { upsert: true },
+      );
+      if (!result.upsertedCount) {
         throw new Error('Memory with this key already exists');
       }
 
-      await MemoryEntry.create({
-        userId,
-        key,
-        value,
-        tokenCount,
-        ...(agentId ? { agentId } : {}),
-        updated_at: new Date(),
-      });
-
       return { ok: true };
     } catch (error) {
       throw new Error(

packages/data-schemas/src/methods/pluginAuth.spec.ts

@@ -0,0 +1,52 @@
+import mongoose from 'mongoose';
+import { MongoMemoryServer } from 'mongodb-memory-server';
+import type { PluginAuthMethods } from './pluginAuth';
+import { createPluginAuthMethods } from './pluginAuth';
+import { createModels } from '~/models';
+
+jest.mock('~/config/winston', () => ({
+  error: jest.fn(),
+  warn: jest.fn(),
+  info: jest.fn(),
+  debug: jest.fn(),
+}));
+
+let mongoServer: InstanceType<typeof MongoMemoryServer>;
+let methods: PluginAuthMethods;
+
+const userId = new mongoose.Types.ObjectId().toString();
+const key = { userId, pluginKey: 'web_search', authField: 'API_KEY' };
+
+beforeAll(async () => {
+  mongoServer = await MongoMemoryServer.create();
+  createModels(mongoose);
+  methods = createPluginAuthMethods(mongoose);
+  await mongoose.connect(mongoServer.getUri());
+});
+
+afterAll(async () => {
+  await mongoose.disconnect();
+  await mongoServer.stop();
+});
+
+beforeEach(async () => {
+  await mongoose.models.PluginAuth.deleteMany({});
+});
+
+describe('updatePluginAuth', () => {
+  it('creates the credential when missing and returns it', async () => {
+    const created = await methods.updatePluginAuth({ ...key, value: 'first' });
+
+    expect(created).toMatchObject({ ...key, value: 'first' });
+    expect(await mongoose.models.PluginAuth.countDocuments(key)).toBe(1);
+  });
+
+  it('updates the existing credential in place instead of inserting another', async () => {
+    const created = await methods.updatePluginAuth({ ...key, value: 'first' });
+    const updated = await methods.updatePluginAuth({ ...key, value: 'second' });
+
+    expect(updated?._id?.toString()).toBe(created?._id?.toString());
+    expect(updated?.value).toBe('second');
+    expect(await mongoose.models.PluginAuth.countDocuments(key)).toBe(1);
+  });
+});

packages/data-schemas/src/methods/pluginAuth.ts

@@ -89,28 +89,12 @@ export function createPluginAuthMethods(mongoose: typeof import('mongoose')): {
   }: UpdatePluginAuthParams): Promise<IPluginAuth> {
     try {
       const PluginAuth: Model<IPluginAuth> = mongoose.models.PluginAuth;
-      const existingAuth = await PluginAuth.findOne({
-        userId,
-        pluginKey,
-        authField,
-      }).lean<IPluginAuth>();
-
-      if (existingAuth) {
-        return await PluginAuth.findOneAndUpdate(
-          { userId, pluginKey, authField },
-          { $set: { value } },
-          { new: true, upsert: true },
-        ).lean<IPluginAuth>();
-      } else {
-        const newPluginAuth = await new PluginAuth({
-          userId,
-          authField,
-          value,
-          pluginKey,
-        });
-        await newPluginAuth.save();
-        return newPluginAuth.toObject();
-      }
+      /** Single atomic upsert: a find-then-save let concurrent writers insert duplicates. */
+      return await PluginAuth.findOneAndUpdate(
+        { userId, pluginKey, authField },
+        { $set: { value } },
+        { new: true, upsert: true, setDefaultsOnInsert: true },
+      ).lean<IPluginAuth>();
     } catch (error) {
       throw new Error(
         `Failed to update plugin auth: ${error instanceof Error ? error.message : 'Unknown error'}`,

packages/data-schemas/src/methods/preset.spec.ts

@@ -0,0 +1,82 @@
+import mongoose from 'mongoose';
+import { MongoMemoryServer } from 'mongodb-memory-server';
+import { createPresetMethods } from './preset';
+import { createModels } from '~/models';
+
+jest.mock('~/config/winston', () => ({
+  error: jest.fn(),
+  warn: jest.fn(),
+  info: jest.fn(),
+  debug: jest.fn(),
+}));
+
+let mongoServer: InstanceType<typeof MongoMemoryServer>;
+let methods: ReturnType<typeof createPresetMethods>;
+
+const user = new mongoose.Types.ObjectId().toString();
+const otherUser = new mongoose.Types.ObjectId().toString();
+
+beforeAll(async () => {
+  mongoServer = await MongoMemoryServer.create();
+  createModels(mongoose);
+  methods = createPresetMethods(mongoose);
+  await mongoose.connect(mongoServer.getUri());
+});
+
+afterAll(async () => {
+  await mongoose.disconnect();
+  await mongoServer.stop();
+});
+
+beforeEach(async () => {
+  await mongoose.models.Preset.deleteMany({});
+});
+
+async function defaultsFor(owner: string): Promise<string[]> {
+  const docs = await mongoose.models.Preset.find({ user: owner, defaultPreset: true }).lean();
+  return docs.map((doc) => doc.presetId as string).sort();
+}
+
+describe('savePreset default handling', () => {
+  it('demotes every other default of the same user, leaving other users untouched', async () => {
+    await mongoose.models.Preset.create([
+      { user, presetId: 'a', title: 'a', endpoint: 'openAI', defaultPreset: true, order: 0 },
+      { user, presetId: 'b', title: 'b', endpoint: 'openAI', defaultPreset: true, order: 0 },
+      {
+        user: otherUser,
+        presetId: 'x',
+        title: 'x',
+        endpoint: 'openAI',
+        defaultPreset: true,
+        order: 0,
+      },
+    ]);
+
+    await methods.savePreset(user, {
+      presetId: 'c',
+      title: 'c',
+      endpoint: 'openAI',
+      defaultPreset: true,
+    });
+
+    expect(await defaultsFor(user)).toEqual(['c']);
+    expect(await defaultsFor(otherUser)).toEqual(['x']);
+  });
+
+  it('keeps the default flag when re-saving the current default', async () => {
+    await methods.savePreset(user, {
+      presetId: 'a',
+      title: 'a',
+      endpoint: 'openAI',
+      defaultPreset: true,
+    });
+    await methods.savePreset(user, {
+      presetId: 'a',
+      title: 'renamed',
+      endpoint: 'openAI',
+      defaultPreset: true,
+    });
+
+    expect(await defaultsFor(user)).toEqual(['a']);
+  });
+});

packages/data-schemas/src/methods/preset.ts

@@ -159,13 +159,11 @@ export function createPresetMethods(mongoose: typeof import('mongoose')): {
         update.defaultPreset = defaultPreset;
         update.order = 0;
 
-        const currentDefault = await Preset.findOne({ defaultPreset: true, user });
-
-        if (currentDefault && currentDefault.presetId !== presetId) {
-          await Preset.findByIdAndUpdate(currentDefault._id, {
-            $unset: { defaultPreset: '', order: '' },
-          });
-        }
+        /** One write demotes every other default, instead of a read then a by-id update. */
+        await Preset.updateMany(
+          { user, defaultPreset: true, ...(presetId != null && { presetId: { $ne: presetId } }) },
+          { $unset: { defaultPreset: '', order: '' } },
+        );
       } else if (defaultPreset === false) {
         update.defaultPreset = undefined;
         update.order = undefined;

packages/data-schemas/src/methods/user.methods.spec.ts

@@ -1939,6 +1939,48 @@ describe('User Methods - Database Tests', () => {
     });
   });
 
+  describe('updateUserPlugins', () => {
+    test('keeps concurrent installs made from the same stale snapshot', async () => {
+      const user = await User.create({
+        name: 'Test User',
+        email: 'test@example.com',
+        provider: 'local',
+        plugins: ['existing'],
+      });
+      const userId = user._id?.toString() ?? '';
+      const staleSnapshot = ['existing'];
+
+      await Promise.all([
+        methods.updateUserPlugins(userId, staleSnapshot, 'alpha', 'install'),
+        methods.updateUserPlugins(userId, staleSnapshot, 'beta', 'install'),
+        methods.updateUserPlugins(userId, staleSnapshot, 'alpha', 'install'),
+      ]);
+
+      const stored = await User.findById(userId).lean();
+      expect([...(stored?.plugins ?? [])].sort()).toEqual(['alpha', 'beta', 'existing']);
+    });
+
+    test('uninstall removes only the target plugin, ignoring the stale snapshot', async () => {
+      const user = await User.create({
+        name: 'Test User',
+        email: 'test@example.com',
+        provider: 'local',
+        plugins: ['alpha', 'beta', 'gamma'],
+      });
+      const userId = user._id?.toString() ?? '';
+
+      const updated = await methods.updateUserPlugins(userId, ['alpha'], 'alpha', 'uninstall');
+
+      expect(updated?.plugins).toEqual(['beta', 'gamma']);
+    });
+
+    test('returns null for an unknown user', async () => {
+      const fakeId = new mongoose.Types.ObjectId().toString();
+      const result = await methods.updateUserPlugins(fakeId, [], 'alpha', 'install');
+      expect(result).toBeNull();
+    });
+  });
+
   describe('toggleUserMemories', () => {
     test('should enable memories for user', async () => {
       const user = await User.create({