diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 21cef79ed7..4672c1a30c 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/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); diff --git a/packages/api/src/code/lifecycle.spec.ts b/packages/api/src/code/lifecycle.spec.ts new file mode 100644 index 0000000000..e844dbe083 --- /dev/null +++ b/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); + }); +}); diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index f82de80e0b..fe394d95df 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/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>(); + 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>>(); + 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); })