🧮 fix: Scale Code Environment Reconciliation (#15625)

* fix: Scale code environment reconciliation

* fix: Bound reservation cleanup index scan

* fix: Bound reservation expiry index scans

* fix: Make reservation cleanup compatible with rolling upgrades

* fix: Fence cleanup sweeps against continuous inserts

packages/api/src/code/environments.integration.spec.ts

@@ -1159,6 +1159,122 @@ describe('code environment registry', () => {
     expect(claimed?.deletionLeaseId).not.toBe('abandoned-removal');
   });
 
+  test('bounds expired agent reservation cleanup per reconciliation tick', async () => {
+    const methods = createMethods(mongoose);
+    const ownerId = new Types.ObjectId();
+    const expiredAt = new Date(Date.now() - 1_000);
+    await mongoose.models.User.create({
+      _id: ownerId,
+      email: 'reservation-owner@example.com',
+      provider: 'local',
+    });
+    const environments = await Promise.all(
+      ['first', 'second'].map((suffix) =>
+        methods.createCodeEnvironment({
+          environmentId: `expired-reservation-${suffix}`,
+          name: `Expired reservation ${suffix}`,
+          type: 'attached',
+          baseURL: 'https://code.example.com',
+          controlPlaneId: 'shared-code-api',
+          createdBy: ownerId,
+        }),
+      ),
+    );
+    await Promise.all(
+      environments.map((environment, index) =>
+        mongoose.models.CodeEnvironment.updateOne(
+          { _id: environment._id },
+          {
+            $set: {
+              pendingAgentReferences: Array.from(
+                { length: index === 0 ? 100 : 1 },
+                (_, offset) => ({
+                  reservationId: `reservation-${index}-${offset}`,
+                  expiresAt: expiredAt,
+                }),
+              ),
+            },
+          },
+        ),
+      ),
+    );
+
+    const cleanupPlan = await mongoose.models.CodeEnvironment.collection
+      .find({})
+      .sort({ _id: 1 })
+      .hint('_id_')
+      .limit(1)
+      .explain('executionStats');
+    expect(JSON.stringify(cleanupPlan.queryPlanner.winningPlan)).not.toContain('"stage":"SORT"');
+    expect(cleanupPlan.executionStats.totalDocsExamined).toBeLessThanOrEqual(1);
+    expect(cleanupPlan.executionStats.totalKeysExamined).toBeLessThanOrEqual(1);
+
+    await reconcileCodeEnvironmentLifecycle({ mongoose, limit: 1 });
+
+    const afterFirstTick = await mongoose.models.CodeEnvironment.countDocuments({
+      _id: { $in: environments.map(({ _id }) => _id) },
+      'pendingAgentReferences.0': { $exists: true },
+    });
+    expect(afterFirstTick).toBe(1);
+
+    await reconcileCodeEnvironmentLifecycle({ mongoose, limit: 1 });
+
+    await expect(
+      mongoose.models.CodeEnvironment.countDocuments({
+        _id: { $in: environments.map(({ _id }) => _id) },
+        'pendingAgentReferences.0': { $exists: true },
+      }),
+    ).resolves.toBe(0);
+  });
+
+  test('revisits legacy writes behind the durable cleanup cursor and preserves live leases', async () => {
+    const model = mongoose.models.CodeEnvironment;
+    const owner = await mongoose.models.User.create({
+      email: 'legacy-reference-owner@example.com',
+      provider: 'local',
+    });
+    const environment = await createMethods(mongoose).createCodeEnvironment({
+      environmentId: 'legacy-reference',
+      name: 'Legacy reference',
+      type: 'attached',
+      baseURL: 'https://code.example.com',
+      controlPlaneId: 'shared-code-api',
+      createdBy: owner._id,
+    });
+    await reconcileCodeEnvironmentLifecycle({ mongoose, limit: 1 });
+    const checkpoints = mongoose.connection.db!.collection<{
+      _id: string;
+      lastId: Types.ObjectId | null;
+    }>('code_environment_reconciliation');
+    expect((await checkpoints.findOne({ _id: 'agent-reference-cleanup' }))?.lastId).toBeNull();
+    const live = { reservationId: 'live', expiresAt: new Date(Date.now() + 60_000) };
+    await model.updateOne(
+      { _id: environment._id },
+      {
+        $push: {
+          pendingAgentReferences: {
+            $each: [{ reservationId: 'abandoned', expiresAt: new Date(0) }, live],
+          },
+        },
+      },
+    );
+    for (let index = 0; index < 3; index++) {
+      await createMethods(mongoose).createCodeEnvironment({
+        environmentId: `new-tail-${index}`,
+        name: 'New tail',
+        type: 'attached',
+        baseURL: 'https://code.example.com',
+        controlPlaneId: 'shared-code-api',
+        createdBy: owner._id,
+      });
+      await reconcileCodeEnvironmentLifecycle({ mongoose, limit: 1 });
+    }
+    const result = await model
+      .findById(environment._id)
+      .lean<{ pendingAgentReferences: (typeof live)[] }>();
+    expect(result?.pendingAgentReferences).toEqual([expect.objectContaining(live)]);
+  });
+
   test('preserves persisted agent references during interrupted-removal recovery', async () => {
     const ownerId = new Types.ObjectId();
     const registry = createCodeEnvironmentRegistry(mongoose);

packages/api/src/code/lifecycle.spec.ts

@@ -0,0 +1,27 @@
+import mongoose from 'mongoose';
+import { startCodeEnvironmentLifecycleReconciler } from './lifecycle';
+import { isLeader } from '~/cluster';
+
+jest.mock('~/cluster', () => ({ isLeader: jest.fn() }));
+
+describe('code environment lifecycle scheduler', () => {
+  afterEach(() => {
+    jest.clearAllTimers();
+    jest.useRealTimers();
+  });
+
+  test('runs reconciliation only on the elected API replica', async () => {
+    jest.useFakeTimers();
+    const leader = jest.mocked(isLeader);
+    leader.mockResolvedValueOnce(false).mockResolvedValueOnce(true);
+
+    startCodeEnvironmentLifecycleReconciler({ mongoose });
+    await jest.advanceTimersByTimeAsync(0);
+
+    expect(leader).toHaveBeenCalledTimes(1);
+
+    await jest.advanceTimersByTimeAsync(60_000);
+
+    expect(leader).toHaveBeenCalledTimes(2);
+  });
+});

packages/api/src/code/lifecycle.ts

@@ -1,8 +1,10 @@
 import { EModelEndpoint, ResourceType } from 'librechat-data-provider';
 import { createMethods, logger, runAsSystem } from '@librechat/data-schemas';
 import type { AppConfig, CodeEnvironmentDocument } from '@librechat/data-schemas';
+import type { Types } from 'mongoose';
 import type { CodeBridgeFetch } from './bridge';
 import { readCodeBridgeSecret, revokeCodeBridgeWorker } from './bridge';
+import { isLeader } from '~/cluster';
 
 const RECONCILE_INTERVAL_MS = 60_000;
 const RECONCILE_LEASE_MS = 2 * 60_000;
@@ -36,9 +38,59 @@ export async function reconcileCodeEnvironmentLifecycle({
     if (CodeEnvironment == null || AclEntry == null) return;
     const methods = createMethods(mongoose);
     const now = new Date();
-    await CodeEnvironment.updateMany(
-      {},
-      { $pull: { pendingAgentReferences: { expiresAt: { $lte: now } } } },
+    const checkpoints = mongoose.connection.db!.collection<{
+      _id: string;
+      lastId: Types.ObjectId | null;
+      upperId?: Types.ObjectId | null;
+    }>('code_environment_reconciliation');
+    const checkpointId = 'agent-reference-cleanup';
+    await checkpoints.updateOne(
+      { _id: checkpointId },
+      { $setOnInsert: { lastId: null } },
+      { upsert: true },
+    );
+    const checkpoint = await checkpoints.findOne({ _id: checkpointId });
+    const lastId = checkpoint?.lastId ?? null;
+    let upperId = checkpoint?.upperId;
+    if (upperId == null) {
+      const tail = await CodeEnvironment.findOne({})
+        .hint('_id_')
+        .sort({ _id: -1 })
+        .select('_id')
+        .lean<Pick<CodeEnvironmentDocument, '_id'>>();
+      upperId = tail?._id ?? null;
+      const captured = await checkpoints.updateOne(
+        { _id: checkpointId, lastId, upperId: null },
+        { $set: { upperId } },
+      );
+      if (captured.matchedCount !== 1) return;
+    }
+    // Persist progress across leader changes and include old-replica writes on each sweep.
+    const expiredReferenceCandidates = await CodeEnvironment.find({
+      _id: { $lte: upperId, ...(lastId == null ? {} : { $gt: lastId }) },
+    })
+      .hint('_id_')
+      .sort({ _id: 1 })
+      .limit(limit)
+      .select('_id')
+      .lean<Array<Pick<CodeEnvironmentDocument, '_id'>>>();
+    if (expiredReferenceCandidates.length > 0) {
+      await CodeEnvironment.updateMany(
+        { _id: { $in: expiredReferenceCandidates.map(({ _id }) => _id) } },
+        { $pull: { pendingAgentReferences: { expiresAt: { $lte: now } } } },
+      );
+    }
+    // Advance only after cleanup succeeds; a competing sweep must not rewind progress.
+    const nextId = expiredReferenceCandidates[expiredReferenceCandidates.length - 1]?._id;
+    const sweepComplete = nextId == null || nextId.equals(upperId);
+    await checkpoints.updateOne(
+      { _id: checkpointId, lastId, upperId },
+      {
+        $set: {
+          lastId: sweepComplete ? null : nextId,
+          upperId: sweepComplete ? null : upperId,
+        },
+      },
     );
     const expiredRemovals = await CodeEnvironment.find({
       deletionCommittedAt: { $exists: false },
@@ -329,7 +381,11 @@ export function startCodeEnvironmentLifecycleReconciler(
   if (reconcileTimer != null) return;
   const run = (): void => {
     if (reconcileInFlight != null) return;
-    const current = reconcileCodeEnvironmentLifecycle(options)
+    const current = isLeader()
+      .then(async (leader) => {
+        if (!leader) return;
+        await reconcileCodeEnvironmentLifecycle(options);
+      })
       .catch((error) => {
         logger.error('[code-environments] lifecycle reconciliation failed:', error);
       })