diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 961ffbc4ef..cc5ba64662 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -370,6 +370,7 @@ const loadTools = async ({ codeApiBaseUrl: codeExecutionContext.baseUrl, executionProfile: codeExecutionContext.executionProfile, executionRouteKey: codeExecutionContext.executionRouteKey, + bridgeWorkerId: codeExecutionContext.bridgeWorkerId, }); if (toolContext) { dynamicToolContextMap[tool] = toolContext; diff --git a/api/app/clients/tools/util/handleTools.test.js b/api/app/clients/tools/util/handleTools.test.js index 6ceaa578cf..4b9f38e70b 100644 --- a/api/app/clients/tools/util/handleTools.test.js +++ b/api/app/clients/tools/util/handleTools.test.js @@ -11,8 +11,10 @@ const mockCreateMCPTool = jest.fn(); const mockCreateMCPTools = jest.fn(); const mockGetServerConfig = jest.fn(); const mockGetAccessibleMcpServerNames = jest.fn(async () => []); +const mockPrimeCodeFiles = jest.fn(async () => ({ files: [], toolContext: undefined })); const mockCreateSearchTool = jest.fn(() => ({ name: 'web_search' })); +const mockCreateCodeExecutionTool = jest.fn(() => ({ name: 'execute_code' })); const mockLoadWebSearchAuth = jest.fn(async () => ({ authenticated: true, authResult: { searchProvider: 'serper', searxngInstanceUrl: 'http://searxng.internal:8080' }, @@ -21,6 +23,11 @@ const mockLoadWebSearchAuth = jest.fn(async () => ({ jest.mock('@librechat/agents', () => ({ ...jest.requireActual('@librechat/agents'), createSearchTool: (...args) => mockCreateSearchTool(...args), + createCodeExecutionTool: (...args) => mockCreateCodeExecutionTool(...args), +})); + +jest.mock('~/server/services/Files/Code/process', () => ({ + primeFiles: (...args) => mockPrimeCodeFiles(...args), })); jest.mock('@librechat/api', () => ({ @@ -365,6 +372,32 @@ describe('Tool Handlers', () => { expect(loadedTools[0].name).toBe(ASK_USER_QUESTION_TOOL_NAME); }); + it('routes code file priming to the selected bridge worker', async () => { + const bridgeWorkerId = 'principal-worker'; + const toolMap = await loadTools({ + user: fakeUser._id.toString(), + tools: [Tools.execute_code], + returnMap: true, + agent: { id: 'agent-1' }, + options: { + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + bridgeWorkerId, + }, + }, + }); + + await toolMap[Tools.execute_code](); + + expect(mockPrimeCodeFiles).toHaveBeenCalledWith( + expect.objectContaining({ + agentId: 'agent-1', + codeApiBaseUrl: 'https://code.example.com/v1', + bridgeWorkerId, + }), + ); + }); + it('passes request body to chat MCP tool creation and skips stale cache for BODY-scoped servers', async () => { const serverName = 'body-scoped'; const toolKey = `search${Constants.mcp_delimiter}${serverName}`; diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index f4fe17d1e4..d90f7f0f5b 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -12,6 +12,7 @@ const { getWebSearchUninstallFields, deleteAgentCheckpoints, deleteAllSharedLinksWithCleanup, + revokeUserCodeEnvironmentWorkers, } = require('@librechat/api'); const { Tools, @@ -404,6 +405,7 @@ const deleteUserController = async (req, res) => { if (triggerDeletionFence != null) { await prepareAgentTriggerUserPurge(user.id, triggerDeletionFence, user.tenantId); } + const deletionAppConfig = await getAppConfig({ baseOnly: true }); await drainAgentTriggerDeliveriesForUser(user.id); await subagentThreadTaskStore.cancelAndDrainForOwner(user.id, user.tenantId); // Reversibly suspend the user's schedules under a per-attempt token BEFORE draining. @@ -460,10 +462,6 @@ const deleteUserController = async (req, res) => { await db.deleteAllUserMemories(user.id); await db.deleteUserPrompts(user.id); await db.deleteUserSkills(user.id); - await db.deleteUserCodeEnvironments(user.id); - await invalidateCodeEnvironmentConfigCache(user.tenantId).catch((error) => { - logger.error('[deleteUserController] code environment cache invalidation failed:', error); - }); await deleteUserMcpServers(user.id); await db.deleteActions({ user: user.id }); await db.deleteTokens({ userId: user.id }); @@ -475,6 +473,27 @@ const deleteUserController = async (req, res) => { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; + let codeEnvironmentCleanupSafe = true; + try { + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: user.id, + appConfig: deletionAppConfig, + }); + } catch (error) { + codeEnvironmentCleanupSafe = false; + logger.error('[deleteUserController] Failed to revoke code environment workers', error); + } + if (codeEnvironmentCleanupSafe) { + try { + await db.deleteUserCodeEnvironments(user.id); + } catch (error) { + logger.error('[deleteUserController] Failed to delete code environments', error); + } + } + await invalidateCodeEnvironmentConfigCache(user.tenantId).catch((error) => { + logger.error('[deleteUserController] code environment cache invalidation failed:', error); + }); await purgeAgentTriggerDeliveriesForUser(user.id); logger.info(`User deleted account. Email: ${user.email} ID: ${user.id}`); res.status(200).send({ message: 'User deleted' }); diff --git a/api/server/controllers/UserController.spec.js b/api/server/controllers/UserController.spec.js index 2c6b70cba9..7d3fd9ff8e 100644 --- a/api/server/controllers/UserController.spec.js +++ b/api/server/controllers/UserController.spec.js @@ -12,6 +12,7 @@ const mockQuiesceUserSchedules = jest.fn().mockResolvedValue(true); const mockRestoreUserSchedules = jest.fn().mockResolvedValue(undefined); const mockGetWebSearchInstallEntries = jest.fn(); const mockInvalidateCodeEnvironmentConfigCache = jest.fn().mockResolvedValue(undefined); +const mockRevokeUserCodeEnvironmentWorkers = jest.fn().mockResolvedValue(0); jest.mock('@librechat/data-schemas', () => { const actual = jest.requireActual('@librechat/data-schemas'); @@ -90,6 +91,7 @@ jest.mock('@librechat/api', () => ({ needsRefresh: jest.fn(), getNewS3URL: jest.fn(), getWebSearchInstallEntries: (...args) => mockGetWebSearchInstallEntries(...args), + revokeUserCodeEnvironmentWorkers: (...args) => mockRevokeUserCodeEnvironmentWorkers(...args), GenerationJobManager: { getCleanupBlockingJobIdsForUser: (...args) => mockGetActiveJobIdsForUser(...args), abortJob: (...args) => mockAbortJob(...args), @@ -161,6 +163,7 @@ const { deleteConvos, acceptTerms, deleteUserById, + deleteUserCodeEnvironments, deleteMessages, beginAgentTriggerUserDeletion, cancelAgentTriggerUserDeletion, @@ -491,6 +494,12 @@ describe('deleteUserController', () => { expect(deleteUserById.mock.invocationCallOrder[0]).toBeLessThan( mockPurgeAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0], ); + expect(deleteUserById.mock.invocationCallOrder[0]).toBeLessThan( + mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0], + ); + expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( + deleteUserCodeEnvironments.mock.invocationCallOrder[0], + ); expect(mockPurgeAgentTriggerDeliveriesForUser).toHaveBeenCalledWith(userId.toString()); expect(cancelAgentTriggerUserDeletion).not.toHaveBeenCalled(); expect(mockCancelAgentTriggerUserPurge).not.toHaveBeenCalled(); diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 23c152b30b..aabf4a63c0 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -32,6 +32,7 @@ const mockCancelAgentTriggerUserDeletion = jest.fn(); const mockCancelAndDrainSubagentThreads = jest.fn(); const mockQuiesceUserSchedules = jest.fn(); const mockDeleteSchedulesByUser = jest.fn(); +const mockRevokeUserCodeEnvironmentWorkers = jest.fn(); jest.mock('@librechat/data-schemas', () => ({ logger: { error: jest.fn(), info: jest.fn() }, @@ -53,6 +54,7 @@ jest.mock('@librechat/api', () => ({ needsRefresh: jest.fn(), getNewS3URL: jest.fn(), deleteAllSharedLinksWithCleanup: (...args) => mockDeleteAllSharedLinksWithCleanup(...args), + revokeUserCodeEnvironmentWorkers: (...args) => mockRevokeUserCodeEnvironmentWorkers(...args), GenerationJobManager: { getCleanupBlockingJobIdsForUser: (...args) => mockGetCleanupBlockingJobIdsForUser(...args), abortJob: (...args) => mockAbortJob(...args), @@ -138,7 +140,7 @@ jest.mock('~/server/services/Schedules', () => ({ })); jest.mock('~/server/services/Config', () => ({ - getAppConfig: jest.fn(), + getAppConfig: jest.fn().mockResolvedValue({}), invalidateCodeEnvironmentConfigCache: (...args) => mockInvalidateCodeEnvironmentConfigCache(...args), })); @@ -186,6 +188,7 @@ function stubDeletionMocks() { mockCancelAndDrainSubagentThreads.mockResolvedValue(); mockQuiesceUserSchedules.mockResolvedValue(true); mockDeleteSchedulesByUser.mockResolvedValue(); + mockRevokeUserCodeEnvironmentWorkers.mockResolvedValue(0); } beforeEach(() => { @@ -208,6 +211,9 @@ describe('deleteUserController - 2FA enforcement', () => { expect(mockDeleteUserPrompts).toHaveBeenCalledWith('user1'); expect(mockDeleteUserSkills).toHaveBeenCalledWith('user1'); expect(mockInvalidateCodeEnvironmentConfigCache).toHaveBeenCalledWith(undefined); + expect(mockRevokeUserCodeEnvironmentWorkers).toHaveBeenCalledWith( + expect.objectContaining({ userId: 'user1' }), + ); expect(mockVerifyOTPOrBackupCode).not.toHaveBeenCalled(); expect(mockBeginAgentTriggerUserDeletion.mock.invocationCallOrder[0]).toBeLessThan( mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0], @@ -221,6 +227,12 @@ describe('deleteUserController - 2FA enforcement', () => { expect(mockDeleteMessages.mock.invocationCallOrder[0]).toBeLessThan( mockDeleteUserById.mock.invocationCallOrder[0], ); + expect(mockDeleteUserById.mock.invocationCallOrder[0]).toBeLessThan( + mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0], + ); + expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( + mockDeleteUserCodeEnvironments.mock.invocationCallOrder[0], + ); expect(mockCancelAgentTriggerUserDeletion).not.toHaveBeenCalled(); expect(mockCancelAgentTriggerUserPurge).not.toHaveBeenCalled(); }); @@ -243,6 +255,23 @@ describe('deleteUserController - 2FA enforcement', () => { ); }); + it('preserves code environment records when revocation marking fails after deletion', async () => { + const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; + const res = createRes(); + mockGetUserById.mockResolvedValue({ _id: 'user1', twoFactorEnabled: false }); + mockRevokeUserCodeEnvironmentWorkers.mockRejectedValueOnce(new Error('worker unavailable')); + + await deleteUserController(req, res); + + expect(res.status).toHaveBeenCalledWith(200); + expect(mockBeginAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); + expect(mockDeleteMessages).toHaveBeenCalled(); + expect(mockDeleteUserById).toHaveBeenCalledWith('user1'); + expect(mockDeleteUserCodeEnvironments).not.toHaveBeenCalled(); + expect(mockCancelAgentTriggerUserPurge).not.toHaveBeenCalled(); + expect(mockCancelAgentTriggerUserDeletion).not.toHaveBeenCalled(); + }); + it('proceeds with deletion when user has no 2FA record', async () => { const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; const res = createRes(); diff --git a/api/server/controllers/agents/__tests__/v1.duplicate-actions.spec.js b/api/server/controllers/agents/__tests__/v1.duplicate-actions.spec.js index cc298bd03a..f0a4bf6095 100644 --- a/api/server/controllers/agents/__tests__/v1.duplicate-actions.spec.js +++ b/api/server/controllers/agents/__tests__/v1.duplicate-actions.spec.js @@ -102,6 +102,49 @@ describe('duplicateAgentHandler — action domain extraction', () => { expect(duplicatedAction.metadata.domain).toBe('api.example.com'); }); + it('removes cloned actions when Agent creation is rejected', async () => { + const userId = new mongoose.Types.ObjectId(); + const originalAgentId = 'agent_conflicted_clone_source'; + const agent = await mongoose.models.Agent.create({ + id: originalAgentId, + name: 'Conflicted clone source', + author: userId.toString(), + provider: 'openai', + model: 'gpt-4', + tools: [], + actions: [`api.example.com${actionDelimiter}act_original`], + versions: [{ name: 'Conflicted clone source', createdAt: new Date(), updatedAt: new Date() }], + }); + await mongoose.models.Action.create({ + user: userId, + action_id: 'act_original', + agent_id: originalAgentId, + metadata: { domain: 'api.example.com' }, + }); + const createAgent = jest + .spyOn(mongoose.models.Agent, 'create') + .mockRejectedValueOnce( + Object.assign(new Error('Code environment is being removed'), { statusCode: 409 }), + ); + const req = { + params: { id: agent.id }, + user: { id: userId.toString() }, + }; + const res = { + status: jest.fn().mockReturnThis(), + json: jest.fn(), + }; + + await duplicateAgent(req, res); + + expect(res.status).toHaveBeenCalledWith(409); + await expect(mongoose.models.Action.countDocuments()).resolves.toBe(1); + await expect( + mongoose.models.Action.findOne({ action_id: 'act_original', agent_id: originalAgentId }), + ).resolves.not.toBeNull(); + createAgent.mockRestore(); + }); + it('strips sensitive metadata fields from duplicated actions', async () => { const userId = new mongoose.Types.ObjectId(); const originalAgentId = 'agent_sensitive'; diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index f22217529b..20c1478697 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -918,6 +918,9 @@ const createAgentHandler = async (req, res) => { return res.status(400).json({ error: 'Invalid request data', details: error.errors }); } logger.error('[/Agents] Error creating agent', error); + if (error?.statusCode === 409) { + return res.status(409).json({ error: error.message }); + } res.status(500).json({ error: error.message }); } }; @@ -1568,7 +1571,18 @@ const duplicateAgentHandler = async (req, res) => { ); newAgentData.actions = agentActions; - const newAgent = await db.createAgent(newAgentData); + let newAgent; + try { + newAgent = await db.createAgent(newAgentData); + } catch (error) { + await db.deleteActions({ agent_id: newAgentId, user: userId }).catch((cleanupError) => { + logger.error( + '[/agents/:id/duplicate] Failed to clean up cloned Actions after Agent creation failed:', + cleanupError, + ); + }); + throw error; + } try { await Promise.all([ @@ -1605,7 +1619,9 @@ const duplicateAgentHandler = async (req, res) => { }); } catch (error) { logger.error('[/Agents/:id/duplicate] Error duplicating Agent:', error); - + if (error?.statusCode === 409) { + return res.status(409).json({ error: error.message }); + } res.status(500).json({ error: error.message }); } }; @@ -2134,6 +2150,9 @@ const revertAgentVersionHandler = async (req, res) => { return res.json(updatedAgent); } catch (error) { logger.error('[/agents/:id/revert] Error reverting Agent version', error); + if (error?.statusCode === 409) { + return res.status(409).json({ error: error.message }); + } res.status(500).json({ error: error.message }); } }; diff --git a/api/server/routes/admin/users.js b/api/server/routes/admin/users.js index ab8f5fec8b..b07c222c77 100644 --- a/api/server/routes/admin/users.js +++ b/api/server/routes/admin/users.js @@ -1,5 +1,6 @@ const express = require('express'); -const { createAdminUsersHandlers } = require('@librechat/api'); +const mongoose = require('mongoose'); +const { createAdminUsersHandlers, revokeUserCodeEnvironmentWorkers } = require('@librechat/api'); const { SystemCapabilities } = require('@librechat/data-schemas'); const { requireCapability } = require('~/server/middleware/roles/capabilities'); const { requireJwtAuth } = require('~/server/middleware'); @@ -10,7 +11,7 @@ const { purgeAgentTriggerDeliveriesForUser, } = require('~/server/services/Agents/triggers'); const db = require('~/models'); -const { invalidateCodeEnvironmentConfigCache } = require('~/server/services/Config'); +const { getAppConfig, invalidateCodeEnvironmentConfigCache } = require('~/server/services/Config'); const router = express.Router(); @@ -27,6 +28,12 @@ const handlers = createAdminUsersHandlers({ prepareAgentTriggerUserPurge, cancelAgentTriggerUserPurge, purgeAgentTriggerDeliveriesForUser, + revokeUserCodeEnvironmentWorkers: async (userId) => + revokeUserCodeEnvironmentWorkers({ + mongoose, + userId, + appConfig: await getAppConfig({ baseOnly: true }), + }), deleteUserById: db.deleteUserById, deleteUserCodeEnvironments: db.deleteUserCodeEnvironments, invalidateCodeEnvironmentConfigCache, diff --git a/api/server/routes/code-environments.js b/api/server/routes/code-environments.js index 9d66bdff4e..4e866dbe92 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -1,18 +1,25 @@ const express = require('express'); -const { createCodeEnvironmentHttpHandlers } = require('@librechat/api'); +const mongoose = require('mongoose'); +const { + createCodeEnvironmentHttpHandlers, + startCodeEnvironmentLifecycleReconciler, +} = require('@librechat/api'); const { SystemCapabilities } = require('@librechat/data-schemas'); const { requireCapability } = require('~/server/middleware/roles/capabilities'); const { codeEnvironmentPairingLimiter } = require('~/server/middleware/limiters/code'); const { getAppConfig, getCodeEnvironmentRegistry } = require('~/server/services/Config'); const { requireJwtAuth } = require('~/server/middleware'); +const db = require('~/models'); const router = express.Router(); +startCodeEnvironmentLifecycleReconciler({ mongoose }); let handlers; function getHandlers() { if (handlers == null) { handlers = createCodeEnvironmentHttpHandlers({ getAppConfig, registry: getCodeEnvironmentRegistry(), + principalIsActive: db.isAgentTriggerPrincipalActive, }); } return handlers; @@ -27,5 +34,6 @@ router.post('/pairings', codeEnvironmentPairingLimiter, (req, res, next) => router.post('/', requireCodeEnvironmentManage, (req, res, next) => getHandlers().register(req, res, next), ); +router.delete('/:environmentId', (req, res, next) => getHandlers().remove(req, res, next)); module.exports = router; diff --git a/api/server/routes/code-environments.test.js b/api/server/routes/code-environments.test.js index d483920040..eda05129a1 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -21,6 +21,7 @@ const mockHandlers = { list: jest.fn((_req, res) => res.status(200).json({ environments: [] })), register: jest.fn((_req, res) => res.status(201).json({ environment: { id: 'code-1' } })), pair: jest.fn((_req, res) => res.status(201).json({ environment: { id: 'code-1' } })), + remove: jest.fn((_req, res) => res.status(200).json({ environment: { id: 'code-1' } })), }; jest.mock('@librechat/data-schemas', () => ({ @@ -30,6 +31,7 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('@librechat/api', () => ({ createCodeEnvironmentRegistry: jest.fn(() => mockRegistry), createCodeEnvironmentHttpHandlers: jest.fn(() => mockHandlers), + startCodeEnvironmentLifecycleReconciler: jest.fn(), })); jest.mock('~/server/middleware/roles/capabilities', () => ({ @@ -44,6 +46,7 @@ jest.mock('~/server/services/Config', () => ({ getAppConfig: jest.fn(), getCodeEnvironmentRegistry: mockGetCodeEnvironmentRegistry, })); +jest.mock('~/models', () => ({ isAgentTriggerPrincipalActive: jest.fn() })); function createApp() { delete require.cache[require.resolve('./code-environments')]; @@ -92,4 +95,11 @@ describe('code environment routes', () => { expect(middlewareCalls).toEqual(['jwt', 'pairing-limit']); expect(mockHandlers.pair).toHaveBeenCalledTimes(1); }); + + it('allows an authenticated owner to remove an environment', async () => { + await request(createApp()).delete('/api/code-environments/code-1').expect(200); + + expect(middlewareCalls).toEqual(['jwt']); + expect(mockHandlers.remove).toHaveBeenCalledTimes(1); + }); }); diff --git a/client/src/components/Nav/Settings/CodeEnvironments.tsx b/client/src/components/Nav/Settings/CodeEnvironments.tsx new file mode 100644 index 0000000000..6b7204f2ba --- /dev/null +++ b/client/src/components/Nav/Settings/CodeEnvironments.tsx @@ -0,0 +1,213 @@ +import { useEffect, useState } from 'react'; +import { + Button, + Input, + Label, + OGDialog, + OGDialogTemplate, + OGDialogTrigger, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, + Spinner, + useToastContext, +} from '@librechat/client'; +import type { CodeEnvironmentPairingResponse } from '~/data-provider/CodeEnvironments'; +import { + useCodeEnvironmentsQuery, + useDeleteCodeEnvironmentMutation, + usePairCodeEnvironmentMutation, +} from '~/data-provider'; +import { useLocalize } from '~/hooks'; + +function WorkerCommand({ pairing }: { pairing: CodeEnvironmentPairingResponse['pairing'] }) { + const { showToast } = useToastContext(); + const localize = useLocalize(); + const shellQuote = (value: string) => `'${value.replace(/'/g, `'\\''`)}'`; + const command = `librechat-code pair ${shellQuote(pairing.endpoint)} ${shellQuote(pairing.code)} --worker-id ${shellQuote(pairing.workerId)}`; + return ( +
+

+ {localize('com_ui_code_environment_pairing_ready')} +

+

+ {localize('com_ui_code_environment_pairing_description')} +

+ 
+        {command}
+
+ +
+ ); +} + +export default function CodeEnvironments() { + const localize = useLocalize(); + const { showToast } = useToastContext(); + const query = useCodeEnvironmentsQuery(); + const pairMutation = usePairCodeEnvironmentMutation(); + const deleteMutation = useDeleteCodeEnvironmentMutation(); + const [name, setName] = useState(''); + const [controlPlaneId, setControlPlaneId] = useState(''); + const [pairing, setPairing] = useState(); + const controlPlanes = query.data?.controlPlanes ?? []; + const selectedControlPlane = controlPlaneId || controlPlanes[0]?.id || ''; + + useEffect(() => { + if (pairing == null) return; + const remaining = Date.parse(pairing.expiresAt) - Date.now(); + const delay = Number.isFinite(remaining) ? Math.max(0, remaining) : 0; + const timeout = window.setTimeout(() => setPairing(undefined), delay); + return () => window.clearTimeout(timeout); + }, [pairing]); + + const pair = () => { + const trimmedName = name.trim(); + if (!trimmedName || !selectedControlPlane) return; + pairMutation.mutate( + { name: trimmedName, controlPlaneId: selectedControlPlane }, + { + onSuccess: (result) => { + setName(''); + setPairing(result.pairing); + showToast({ + message: localize('com_ui_code_environment_pairing_created'), + status: 'success', + }); + }, + onError: () => + showToast({ + message: localize('com_ui_code_environment_pairing_error'), + status: 'error', + }), + }, + ); + }; + + let pairingControls = ( +

{localize('com_ui_code_environment_unavailable')}

+ ); + if (query.isLoading) { + pairingControls = ( +
+ + {localize('com_ui_loading')} +
+ ); + } else if (query.isError) { + pairingControls = ( +

+ {localize('com_ui_code_environment_load_error')} +

+ ); + } else if (controlPlanes.length > 0) { + pairingControls = ( +
+ setName(event.target.value)} + /> + + +
+ ); + } + + return ( +
+
+ +

+ {localize('com_ui_code_environments_description')} +

+
+ + {pairingControls} + + {pairing != null && } + +
+ {(query.data?.environments ?? []).map((environment) => ( +
+
+

{environment.name}

+

{environment.id}

+
+ {environment.canDelete && ( + + + + + + {localize('com_ui_code_environment_remove_description')} +

+ } + selection={{ + selectHandler: () => + deleteMutation.mutate(environment.id, { + onSuccess: () => + showToast({ + message: localize('com_ui_code_environment_removed'), + status: 'success', + }), + onError: () => + showToast({ + message: localize('com_ui_code_environment_remove_error'), + status: 'error', + }), + }), + selectClasses: + 'bg-surface-destructive text-text-on-status transition-all duration-200 hover:bg-surface-destructive-hover', + selectText: localize('com_ui_delete'), + }} + /> + + )} +
+ ))} +
+
+ ); +} diff --git a/client/src/components/Nav/Settings/registry.tsx b/client/src/components/Nav/Settings/registry.tsx index 2a42db8759..8da07921fa 100644 --- a/client/src/components/Nav/Settings/registry.tsx +++ b/client/src/components/Nav/Settings/registry.tsx @@ -47,6 +47,7 @@ import { showThinkingAtom } from '~/store/showThinking'; import ProviderKeys from '../SettingsTabs/ProviderKeys'; import { autoScrollAtom } from '~/store/autoScroll'; import Avatar from '../SettingsTabs/Account/Avatar'; +import CodeEnvironments from './CodeEnvironments'; import About from '../SettingsTabs/About/About'; import ApiKeys from '../SettingsTabs/ApiKeys'; import MemoryToggle from './MemoryToggle'; @@ -588,6 +589,15 @@ export const registry: SettingEntry[] = [ show: (ctx) => ctx.hasStatefulCodeSessions, Component: StatefulWorkspaceDefault, }, + { + id: 'codeEnvironments', + tab: DATA, + section: 'codeExecution', + labelKey: 'com_ui_code_environments', + keywords: ['bridge', 'cli', 'code', 'environment', 'sandbox', 'vm', 'worker'], + show: (ctx) => ctx.hasStatefulCodeSessions, + Component: CodeEnvironments, + }, // Data controls · Your data { id: 'importConversations', diff --git a/client/src/data-provider/CodeEnvironments.ts b/client/src/data-provider/CodeEnvironments.ts new file mode 100644 index 0000000000..07c3ca65ea --- /dev/null +++ b/client/src/data-provider/CodeEnvironments.ts @@ -0,0 +1,44 @@ +import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'; +import { MutationKeys, QueryKeys, dataService } from 'librechat-data-provider'; +import type { + TCodeEnvironmentPairingResponse, + TCodeEnvironmentsResponse, +} from 'librechat-data-provider'; + +export type CodeEnvironmentPairingResponse = TCodeEnvironmentPairingResponse; + +export function useCodeEnvironmentsQuery(enabled = true) { + return useQuery( + [QueryKeys.codeEnvironments], + () => dataService.getCodeEnvironments(), + { enabled }, + ); +} + +export function usePairCodeEnvironmentMutation() { + const queryClient = useQueryClient(); + return useMutation< + TCodeEnvironmentPairingResponse, + Error, + { name: string; controlPlaneId: string } + >([MutationKeys.pairCodeEnvironment], dataService.pairCodeEnvironment, { + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.codeEnvironments]); + queryClient.invalidateQueries([QueryKeys.endpoints]); + }, + }); +} + +export function useDeleteCodeEnvironmentMutation() { + const queryClient = useQueryClient(); + return useMutation< + { environment: TCodeEnvironmentsResponse['environments'][number] }, + Error, + string + >([MutationKeys.deleteCodeEnvironment], dataService.deleteCodeEnvironment, { + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.codeEnvironments]); + queryClient.invalidateQueries([QueryKeys.endpoints]); + }, + }); +} diff --git a/client/src/data-provider/index.ts b/client/src/data-provider/index.ts index 090c8b9689..38bb19a958 100644 --- a/client/src/data-provider/index.ts +++ b/client/src/data-provider/index.ts @@ -16,6 +16,7 @@ export * from './Subagents'; export * from './Tools'; export * from './connection'; export * from './Favorites'; +export * from './CodeEnvironments'; export * from './mutations'; export * from './prompts'; export * from './queries'; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 59cffed4f1..db284326f6 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -2183,6 +2183,21 @@ "com_ui_stateful_code_environment_conversation": "Conversation workspace", "com_ui_default_stateful_workspace": "Default stateful workspace", "com_ui_default_stateful_workspace_description": "Applied when you create a new agent. Existing agents keep their current workspace scope.", + "com_ui_code_environments": "Code environments", + "com_ui_code_environments_description": "Connect a sandboxed librechat-code worker running on a VM you control. The worker connects outbound and must be configured with its own sandbox and network policy.", + "com_ui_code_environment_name_placeholder": "My development VM", + "com_ui_code_environment_control_plane": "Code control plane", + "com_ui_code_environment_pair": "Connect VM", + "com_ui_code_environment_unavailable": "This LibreChat instance has not enabled self-service code workers.", + "com_ui_code_environment_load_error": "Code environments could not be loaded. Try again shortly.", + "com_ui_code_environment_pairing_ready": "One-time worker command", + "com_ui_code_environment_pairing_description": "Run this on your VM, then start librechat-code with LIBRECHAT_CODE_SANDBOX_ENDPOINT set to your sandbox service. The command expires shortly and is shown only once.", + "com_ui_code_environment_pairing_created": "Code environment created", + "com_ui_code_environment_pairing_error": "Could not create the code environment", + "com_ui_code_environment_remove_title": "Remove code environment?", + "com_ui_code_environment_remove_description": "This revokes the worker before removing it from LibreChat. Its sandbox files remain on your VM.", + "com_ui_code_environment_removed": "Code environment removed", + "com_ui_code_environment_remove_error": "Could not revoke the code environment", "com_ui_status_prefix": "Status:", "com_ui_steer": "Steer", "com_ui_steer_already_applied": "That steering message already reached the agent, so it was left in the response", diff --git a/config/__tests__/delete-user.spec.js b/config/__tests__/delete-user.spec.js index 282fe37767..3ea21b63e8 100644 --- a/config/__tests__/delete-user.spec.js +++ b/config/__tests__/delete-user.spec.js @@ -18,6 +18,7 @@ const mockMethods = { restoreUserSchedulesFromDeletion: jest.fn(), countActiveAgentTriggerDeliveriesByUser: jest.fn(), deleteSchedulesByUser: jest.fn(), + deleteUserCodeEnvironments: jest.fn(), deleteUserById: jest.fn(), deleteAgentTriggerDeliveriesByUser: jest.fn(), cancelAgentTriggerUserPurge: jest.fn(), @@ -28,6 +29,8 @@ const mockAbortJob = jest.fn(); const mockDestroy = jest.fn(); const mockSilentExit = jest.fn(); const mockAskQuestion = jest.fn(); +const mockGetAppConfig = jest.fn(); +const mockRevokeUserCodeEnvironmentWorkers = jest.fn(); jest.mock('../connect', () => jest.fn(async () => undefined)); jest.mock('mongoose', () => ({ disconnect: jest.fn(async () => undefined) })); @@ -46,8 +49,12 @@ jest.mock('@librechat/api', () => ({ abortJob: (...args) => mockAbortJob(...args), destroy: (...args) => mockDestroy(...args), }, + revokeUserCodeEnvironmentWorkers: (...args) => mockRevokeUserCodeEnvironmentWorkers(...args), })); jest.mock('~/cache/getLogStores', () => jest.fn()); +jest.mock('~/server/services/Config', () => ({ + getAppConfig: (...args) => mockGetAppConfig(...args), +})); jest.mock('../helpers', () => ({ ...jest.requireActual('../helpers'), askQuestion: mockAskQuestion, @@ -96,6 +103,7 @@ describe('Delete user CLI', () => { mockMethods.restoreUserSchedulesFromDeletion.mockReset().mockResolvedValue(undefined); mockMethods.countActiveAgentTriggerDeliveriesByUser.mockReset().mockResolvedValue(0); mockMethods.deleteSchedulesByUser.mockReset().mockResolvedValue(undefined); + mockMethods.deleteUserCodeEnvironments.mockReset().mockResolvedValue(0); mockMethods.deleteUserById.mockReset().mockResolvedValue({ deletedCount: 1 }); mockMethods.deleteAgentTriggerDeliveriesByUser.mockReset().mockResolvedValue(undefined); mockMethods.cancelAgentTriggerUserPurge.mockReset().mockResolvedValue(true); @@ -103,6 +111,8 @@ describe('Delete user CLI', () => { mockGetCleanupBlockingJobIdsForUser.mockReset().mockResolvedValue([]); mockAbortJob.mockReset().mockResolvedValue({ success: true }); mockDestroy.mockReset().mockResolvedValue(undefined); + mockGetAppConfig.mockReset().mockResolvedValue({}); + mockRevokeUserCodeEnvironmentWorkers.mockReset().mockResolvedValue(0); mockAskQuestion.mockReset().mockResolvedValueOnce('y').mockResolvedValueOnce('n'); }); @@ -200,6 +210,30 @@ describe('Delete user CLI', () => { expect(mockMethods.restoreUserSchedulesFromDeletion).not.toHaveBeenCalled(); }); + it('commits the account before revoking workers and removing environment records', async () => { + expect(await runCli()).toBe(0); + + expect(mockGetAppConfig).toHaveBeenCalledWith({ baseOnly: true }); + expect(mockRevokeUserCodeEnvironmentWorkers).toHaveBeenCalledWith( + expect.objectContaining({ userId: USER_ID }), + ); + expect(mockMethods.deleteUserCodeEnvironments).toHaveBeenCalledWith(USER_ID); + expect(mockMethods.deleteUserById.mock.invocationCallOrder[0]).toBeLessThan( + mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0], + ); + expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( + mockMethods.deleteUserCodeEnvironments.mock.invocationCallOrder[0], + ); + }); + + it('preserves code environment records when revocation marking fails', async () => { + mockRevokeUserCodeEnvironmentWorkers.mockRejectedValueOnce(new Error('mongo unavailable')); + + expect(await runCli()).toBe(0); + + expect(mockMethods.deleteUserCodeEnvironments).not.toHaveBeenCalled(); + }); + it('deletes nothing and releases both fences when provider drain fails', async () => { mockGetCleanupBlockingJobIdsForUser.mockResolvedValueOnce(['stream-1']); mockAbortJob.mockRejectedValueOnce(new Error('provider drain failed')); diff --git a/config/delete-user.js b/config/delete-user.js index 533a239f2f..333d50d9a7 100644 --- a/config/delete-user.js +++ b/config/delete-user.js @@ -34,9 +34,11 @@ require('module-alias')({ base: path.resolve(__dirname, '..', 'api') }); const { GenerationJobManager, createStreamServices, + revokeUserCodeEnvironmentWorkers, waitForKeyvRedisClient, } = require('@librechat/api'); const getLogStores = require('~/cache/getLogStores'); +const { getAppConfig } = require('~/server/services/Config'); const { askQuestion, silentExit } = require('./helpers'); const connect = require('./connect'); @@ -194,6 +196,8 @@ async function gracefulExit(code = 0) { ); } + const deletionAppConfig = await getAppConfig({ baseOnly: true }); + // 5) Build and run deletion tasks const tasks = [ Action.deleteMany({ user: uid }), @@ -234,6 +238,22 @@ async function gracefulExit(code = 0) { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; + let codeEnvironmentCleanupSafe = true; + try { + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: uid, + appConfig: deletionAppConfig, + }); + } catch (error) { + codeEnvironmentCleanupSafe = false; + console.error('Failed to revoke code environment workers after account deletion:', error); + } + if (codeEnvironmentCleanupSafe) { + await runAsSystem(() => methods.deleteUserCodeEnvironments(uid)).catch((error) => + console.error('Failed to delete code environment records after account deletion:', error), + ); + } await runAsSystem(() => methods.deleteAgentTriggerDeliveriesByUser(uid)); } finally { // RESTORE BEFORE RELEASING THE FENCE. While the user-deletion fence is still armed, new diff --git a/packages/api/src/admin/users.spec.ts b/packages/api/src/admin/users.spec.ts index 83ff375fe2..824dd72514 100644 --- a/packages/api/src/admin/users.spec.ts +++ b/packages/api/src/admin/users.spec.ts @@ -59,6 +59,7 @@ function createDeps(overrides: Partial = {}): AdminUsersDeps { prepareAgentTriggerUserPurge: jest.fn().mockResolvedValue(undefined), cancelAgentTriggerUserPurge: jest.fn().mockResolvedValue(true), purgeAgentTriggerDeliveriesForUser: jest.fn().mockResolvedValue(undefined), + revokeUserCodeEnvironmentWorkers: jest.fn().mockResolvedValue(0), deleteUserById: jest .fn() .mockResolvedValue({ deletedCount: 1, message: 'User was deleted successfully.' }), @@ -476,6 +477,19 @@ describe('createAdminUsersHandlers', () => { expect(json).toHaveBeenCalledWith({ message: 'User was deleted successfully.' }); }); + it('preserves code environment records when revocation marking fails', async () => { + const deps = createDeps({ + revokeUserCodeEnvironmentWorkers: jest.fn().mockRejectedValue(new Error('mongo down')), + }); + const handlers = createAdminUsersHandlers(deps); + const { req, res } = createReqRes({ params: { id: validUserId } }); + + await handlers.deleteUser(req, res); + + expect(deps.deleteUserById).toHaveBeenCalledWith(validUserId); + expect(deps.deleteUserCodeEnvironments).not.toHaveBeenCalled(); + }); + it('does not cascade when user is not found', async () => { const result: UserDeleteResult = { deletedCount: 0, message: '' }; const deps = createDeps({ deleteUserById: jest.fn().mockResolvedValue(result) }); @@ -584,12 +598,18 @@ describe('createAdminUsersHandlers', () => { const drainOrder = (deps.drainAgentTriggerDeliveriesForUser as jest.Mock).mock .invocationCallOrder[0]; const deleteOrder = (deps.deleteUserById as jest.Mock).mock.invocationCallOrder[0]; + const revokeCodeOrder = (deps.revokeUserCodeEnvironmentWorkers as jest.Mock).mock + .invocationCallOrder[0]; + const deleteCodeOrder = (deps.deleteUserCodeEnvironments as jest.Mock).mock + .invocationCallOrder[0]; const purgeOrder = (deps.purgeAgentTriggerDeliveriesForUser as jest.Mock).mock .invocationCallOrder[0]; expect(beginOrder).toBeLessThan(drainOrder); expect(beginOrder).toBeLessThan(prepareOrder); expect(prepareOrder).toBeLessThan(drainOrder); expect(drainOrder).toBeLessThan(deleteOrder); + expect(deleteOrder).toBeLessThan(revokeCodeOrder); + expect(revokeCodeOrder).toBeLessThan(deleteCodeOrder); expect(deleteOrder).toBeLessThan(purgeOrder); }); diff --git a/packages/api/src/admin/users.ts b/packages/api/src/admin/users.ts index 39e0243a79..0103c72696 100644 --- a/packages/api/src/admin/users.ts +++ b/packages/api/src/admin/users.ts @@ -37,6 +37,7 @@ export interface AdminUsersDeps { ) => Promise; cancelAgentTriggerUserPurge: (userId: string, fenceStartedAt: Date) => Promise; purgeAgentTriggerDeliveriesForUser: (userId: string) => Promise; + revokeUserCodeEnvironmentWorkers?: (userId: string) => Promise; /** * Thin data-layer delete — removes the User document only. * Full cascade of user-owned resources (conversations, messages, files, tokens, etc.) @@ -72,8 +73,9 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { prepareAgentTriggerUserPurge, cancelAgentTriggerUserPurge, purgeAgentTriggerDeliveriesForUser, - deleteUserById, + revokeUserCodeEnvironmentWorkers, deleteUserCodeEnvironments, + deleteUserById, invalidateCodeEnvironmentConfigCache, deleteConfig, deleteAclEntries, @@ -203,6 +205,13 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { return res.status(404).json({ error: 'User not found' }); } userDeleted = true; + let codeEnvironmentCleanupSafe = true; + try { + await revokeUserCodeEnvironmentWorkers?.(id); + } catch (error) { + codeEnvironmentCleanupSafe = false; + logger.error('[adminUsers] failed to revoke code environment workers:', id, error); + } await purgeAgentTriggerDeliveriesForUser(id); if (targetUser?.role === SystemRoles.ADMIN) { @@ -218,7 +227,7 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { const objectId = new Types.ObjectId(id); const cleanupResults = await Promise.allSettled([ deleteConfig(PrincipalType.USER, id), - deleteUserCodeEnvironments(objectId), + ...(codeEnvironmentCleanupSafe ? [deleteUserCodeEnvironments(objectId)] : []), deleteAclEntries({ principalType: PrincipalType.USER, principalId: objectId }), ]); for (const r of cleanupResults) { diff --git a/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts index 77b88d48fa..ff476cdee3 100644 --- a/packages/api/src/code/bridge.ts +++ b/packages/api/src/code/bridge.ts @@ -32,6 +32,16 @@ export class CodeBridgePairingError extends Error { } } +export class CodeBridgeLifecycleError extends Error { + constructor( + public readonly reason: 'rejected' | 'invalid' | 'timeout' | 'failed', + public readonly upstreamStatus?: number, + ) { + super(`Code bridge lifecycle request ${reason}`); + this.name = 'CodeBridgeLifecycleError'; + } +} + function validPairing(value: unknown, workerId: string): value is CodeBridgePairing { if (typeof value !== 'object' || value == null) return false; const pairing = value as Partial; @@ -87,6 +97,48 @@ export async function createCodeBridgePairing({ } } +export async function revokeCodeBridgeWorker({ + baseURL, + token, + workerId, + fetchImpl = fetch, +}: { + baseURL: string; + token: string; + workerId: string; + fetchImpl?: CodeBridgeFetch; +}): Promise { + try { + const response = await fetchImpl( + `${baseURL.trim().replace(/\/+$/, '')}/bridge/workers/${encodeURIComponent(workerId)}/revoke`, + { + method: 'POST', + headers: { Authorization: `Bearer ${token}` }, + redirect: 'error', + signal: AbortSignal.timeout(CODE_BRIDGE_REQUEST_TIMEOUT_MS), + }, + ); + if (!response.ok) { + throw new CodeBridgeLifecycleError('rejected', response.status); + } + const payload = (await response.json()) as unknown; + if ( + typeof payload !== 'object' || + payload == null || + (payload as { protocolVersion?: unknown }).protocolVersion !== 1 || + (payload as { revoked?: unknown }).revoked !== true + ) { + throw new CodeBridgeLifecycleError('invalid'); + } + } catch (error) { + if (error instanceof CodeBridgeLifecycleError) throw error; + if (error instanceof Error && error.name === 'TimeoutError') { + throw new CodeBridgeLifecycleError('timeout'); + } + throw new CodeBridgeLifecycleError('failed'); + } +} + export function readCodeBridgeSecret(name: string): string | undefined { return Object.prototype.hasOwnProperty.call(process.env, name) ? process.env[name] : undefined; } diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index cf30112c2b..47477bfe1d 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -1,7 +1,8 @@ import mongoose, { Types } from 'mongoose'; import { MongoMemoryServer } from 'mongodb-memory-server'; -import { createMethods, createModels } from '@librechat/data-schemas'; +import { createMethods, createModels, tenantStorage } from '@librechat/data-schemas'; import { AccessRoleIds, PrincipalType, ResourceType } from 'librechat-data-provider'; +import { reconcileCodeEnvironmentLifecycle, revokeUserCodeEnvironmentWorkers } from './lifecycle'; import { AccessControlService } from '~/acl/accessControlService'; import { createCodeEnvironmentRegistry } from './environments'; @@ -60,6 +61,7 @@ describe('code environment registry', () => { id: 'danny-vm', name: "Danny's VM", type: 'attached', + canDelete: true, }); await expect(registry.listRegisteredIds()).resolves.toEqual(['danny-vm']); await expect( @@ -142,7 +144,72 @@ describe('code environment registry', () => { role: 'CODE_USER', idOnTheSource: null, }), - ).resolves.toEqual([roleEnvironment, groupEnvironment]); + ).resolves.toEqual([ + { ...roleEnvironment, canDelete: false }, + { ...groupEnvironment, canDelete: false }, + ]); + }); + + test('computes delete permissions for an environment list in one batch', async () => { + const batchSpy = jest.spyOn(AccessControlService.prototype, 'getResourcePermissionsMap'); + const singleSpy = jest.spyOn(AccessControlService.prototype, 'checkPermission'); + const registry = createCodeEnvironmentRegistry(mongoose); + const ownerId = new Types.ObjectId(); + for (const id of ['batch-one', 'batch-two']) { + await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id, + name: id, + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + }); + } + + await expect( + registry.listAccessible({ userId: ownerId, role: 'USER', idOnTheSource: null }), + ).resolves.toHaveLength(2); + expect(batchSpy).toHaveBeenCalledTimes(1); + expect(singleSpy).not.toHaveBeenCalled(); + + batchSpy.mockRestore(); + singleSpy.mockRestore(); + }); + + test('atomically limits concurrent environment registrations for one owner', async () => { + await mongoose.models.CodeEnvironment.createCollection(); + await expect(mongoose.models.CodeEnvironment.collection.indexes()).resolves.toEqual([ + expect.objectContaining({ name: '_id_' }), + ]); + const registry = createCodeEnvironmentRegistry(mongoose); + const ownerId = new Types.ObjectId(); + + const results = await Promise.allSettled( + ['quota-one', 'quota-two', 'quota-three'].map((id) => + registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id, + name: id, + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + maxOwned: 2, + } as never), + ), + ); + + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(2); + expect(results.filter(({ status }) => status === 'rejected')).toHaveLength(1); + expect(results.find(({ status }) => status === 'rejected')).toMatchObject({ + reason: { name: 'CodeEnvironmentLimitError' }, + }); + await expect( + mongoose.models.CodeEnvironment.countDocuments({ createdBy: ownerId }), + ).resolves.toBe(2); }); test('keeps a user-bound worker private even if its ACL is granted to a role', async () => { @@ -157,8 +224,8 @@ describe('code environment registry', () => { name: 'Owner worker', type: 'attached', baseURL: 'https://code.example.com', - workerId: 'owner-worker', controlPlaneId: 'shared-code-api', + workerId: 'owner-worker', workerPrincipal: { type: 'user', id: ownerId.toString() }, }, }); @@ -179,6 +246,364 @@ describe('code environment registry', () => { ).resolves.toEqual([environment]); }); + test('deletes an owner environment and its ACL after lifecycle cleanup succeeds', async () => { + const registry = createCodeEnvironmentRegistry(mongoose); + const ownerId = new Types.ObjectId(); + const environment = await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'remove-me', + name: 'Remove me', + type: 'attached', + baseURL: 'https://code.example.com', + workerId: 'remove-me', + controlPlaneId: 'self-service', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + const beforeDelete = jest.fn().mockResolvedValue(undefined); + + await expect( + registry.remove({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environmentId: 'remove-me', + beforeDelete, + }), + ).resolves.toEqual(environment); + expect(beforeDelete).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'remove-me', + workerId: 'remove-me', + controlPlaneId: 'self-service', + }), + ); + await expect( + mongoose.models.AclEntry.countDocuments({ resourceId: environment.resourceId }), + ).resolves.toBe(0); + await expect( + registry.listAccessible({ userId: ownerId, role: 'USER', idOnTheSource: null }), + ).resolves.toEqual([]); + }); + + test('reconciles an interrupted removal after the remote lifecycle has started', async () => { + const registry = createCodeEnvironmentRegistry(mongoose); + const ownerId = new Types.ObjectId(); + await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'interrupted-removal', + name: 'Interrupted removal', + type: 'attached', + baseURL: 'https://code.example.com/v1', + workerId: 'interrupted-removal', + controlPlaneId: 'self-service', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + const beforeDelete = jest.fn().mockResolvedValue(undefined); + const commit = jest + .spyOn(mongoose.models.CodeEnvironment, 'updateOne') + .mockRejectedValueOnce(new Error('mongo unavailable after revoke')); + + await expect( + registry.remove({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environmentId: 'interrupted-removal', + beforeDelete, + }), + ).rejects.toThrow('mongo unavailable after revoke'); + expect(beforeDelete).toHaveBeenCalledTimes(1); + await expect( + registry.listAccessible({ userId: ownerId, role: 'USER', idOnTheSource: null }), + ).resolves.toEqual([]); + + commit.mockRestore(); + await mongoose.models.CodeEnvironment.updateOne( + { environmentId: 'interrupted-removal' }, + { $set: { deletionLeaseExpiresAt: new Date(Date.now() - 1_000) } }, + ); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + await reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl, + }); + + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.example.com/v1/bridge/workers/interrupted-removal/revoke', + expect.objectContaining({ method: 'POST' }), + ); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'interrupted-removal' }), + ).resolves.toBeNull(); + }); + + test('fences removal while an agent write is reserving the environment', async () => { + const registry = createCodeEnvironmentRegistry(mongoose); + const methods = createMethods(mongoose); + const ownerId = new Types.ObjectId(); + await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'agent-write-race', + name: 'Agent write race', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + }); + const Agent = mongoose.models.Agent; + const createAgent = Agent.create.bind(Agent); + let enteredCreate!: () => void; + let releaseCreate!: () => void; + const entered = new Promise((resolve) => (enteredCreate = resolve)); + const release = new Promise((resolve) => (releaseCreate = resolve)); + const createSpy = jest.spyOn(Agent, 'create').mockImplementationOnce(async (input) => { + enteredCreate(); + await release; + return await createAgent(input); + }); + + const pendingAgent = methods.createAgent({ + id: 'agent_write_race', + name: 'Agent write race', + author: ownerId, + model: 'test-model', + provider: 'test-provider', + code_environment_id: 'agent-write-race', + }); + await entered; + + await expect( + registry.remove({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environmentId: 'agent-write-race', + }), + ).rejects.toMatchObject({ name: 'CodeEnvironmentInUseError' }); + + releaseCreate(); + await expect(pendingAgent).resolves.toMatchObject({ + code_environment_id: 'agent-write-race', + }); + createSpy.mockRestore(); + }); + + test('revokes every user-bound worker before account deletion', async () => { + const ownerId = new Types.ObjectId(); + await createCodeEnvironmentRegistry(mongoose).register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'account-worker', + name: 'Account worker', + type: 'attached', + baseURL: 'https://code.example.com/v1', + workerId: 'account-worker', + controlPlaneId: 'self-service', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + + await expect( + revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: ownerId.toString(), + appConfig: { + endpoints: { + agents: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [], + }, + }, + }, + } as never, + readSecret: () => 'administrator-token', + fetchImpl, + }), + ).resolves.toBe(1); + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.example.com/v1/bridge/workers/account-worker/revoke', + expect.objectContaining({ method: 'POST' }), + ); + }); + + test('reports successful revocations without aborting after another worker fails', async () => { + const ownerId = new Types.ObjectId(); + const registry = createCodeEnvironmentRegistry(mongoose); + for (const id of ['worker-ok', 'worker-unreachable']) { + await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id, + name: id, + type: 'attached', + baseURL: `https://${id}.example.com/v1`, + controlPlaneId: 'shared-code-api', + workerId: id, + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + } + const fetchImpl = jest.fn(async (input: string | URL | Request) => { + const url = String(input); + if (url.includes('unreachable')) { + throw new Error('control plane unavailable'); + } + return { + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + } as Response; + }); + + await expect( + revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: ownerId.toString(), + appConfig: {} as never, + readSecret: () => 'administrator-token', + fetchImpl, + }), + ).resolves.toBe(1); + expect(fetchImpl).toHaveBeenCalledTimes(2); + await expect(createMethods(mongoose).deleteUserCodeEnvironments(ownerId)).resolves.toBe(0); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-ok' }).lean(), + ).resolves.toMatchObject({ deletionCommittedAt: expect.any(Date) }); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-unreachable' }).lean(), + ).resolves.toMatchObject({ + revocationPendingAt: expect.any(Date), + revocationAttempts: 1, + revocationLastError: 'Code bridge lifecycle request failed', + }); + + const reconcileFailure = jest.fn().mockRejectedValue(new Error('control plane still down')); + await reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl: reconcileFailure, + }); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-unreachable' }).lean(), + ).resolves.toMatchObject({ revocationReconcileAfter: expect.any(Date) }); + + await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'worker-later', + name: 'worker-later', + type: 'attached', + baseURL: 'https://worker-later.example.com/v1', + controlPlaneId: 'shared-code-api', + workerId: 'worker-later', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + await mongoose.models.CodeEnvironment.updateOne( + { environmentId: 'worker-later' }, + { $set: { revocationPendingAt: new Date() } }, + ); + + const retryFetch = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + await reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl: retryFetch, + limit: 1, + }); + expect(retryFetch).toHaveBeenCalledWith( + 'https://worker-later.example.com/v1/bridge/workers/worker-later/revoke', + expect.objectContaining({ method: 'POST' }), + ); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-later' }), + ).resolves.toBeNull(); + retryFetch.mockClear(); + await mongoose.models.CodeEnvironment.updateOne( + { environmentId: 'worker-unreachable' }, + { $set: { revocationReconcileAfter: new Date(Date.now() - 1_000) } }, + ); + await reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl: retryFetch, + }); + expect(retryFetch).toHaveBeenCalledWith( + 'https://worker-unreachable.example.com/v1/bridge/workers/worker-unreachable/revoke', + expect.objectContaining({ method: 'POST' }), + ); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-unreachable' }), + ).resolves.toBeNull(); + }); + + test('reconciles an unmarked user worker after its owner document is deleted', async () => { + const ownerId = new Types.ObjectId(); + await mongoose.models.User.collection.insertOne({ + _id: ownerId, + email: 'departed-worker-owner@example.com', + }); + await createCodeEnvironmentRegistry(mongoose).register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'unmarked-orphan-worker', + name: 'Unmarked orphan worker', + type: 'attached', + baseURL: 'https://code.example.com/v1', + controlPlaneId: 'shared-code-api', + workerId: 'unmarked-orphan-worker', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + await createCodeEnvironmentRegistry(mongoose).register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'unmarked-local-environment', + name: 'Unmarked local environment', + type: 'attached', + baseURL: 'https://local.example.com/v1', + controlPlaneId: 'deployment-code-api', + }, + }); + await mongoose.models.User.deleteOne({ _id: ownerId }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + + await reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl, + }); + + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.example.com/v1/bridge/workers/unmarked-orphan-worker/revoke', + expect.objectContaining({ method: 'POST' }), + ); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'unmarked-orphan-worker' }), + ).resolves.toBeNull(); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'unmarked-local-environment' }), + ).resolves.toBeNull(); + }); + test('removes creator-owned environment records and grants when the user is deleted', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); @@ -230,6 +655,37 @@ describe('code environment registry', () => { ]); }); + test('does not reuse a cached configuration after removal is fenced', async () => { + const cache = createSharedCache(); + const registry = createCodeEnvironmentRegistry(mongoose, { configurationCache: cache }); + const ownerId = new Types.ObjectId(); + const actor = { userId: ownerId, role: 'USER', idOnTheSource: null }; + const environment = await registry.register({ + actor, + environment: { + id: 'cached-removal-vm', + name: 'Cached removal VM', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + }); + await expect(registry.listAccessibleConfigurations(actor)).resolves.toHaveLength(1); + + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment.resourceId }, + { + $set: { + deletionStartedAt: new Date(), + deletionLeaseId: 'in-flight-removal', + deletionLeaseExpiresAt: new Date(Date.now() + 60_000), + }, + }, + ); + + await expect(registry.listAccessibleConfigurations(actor)).resolves.toEqual([]); + }); + test('caches registered environment ids behind the shared tenant revision', async () => { const cache = createSharedCache(); const firstWorker = createCodeEnvironmentRegistry(mongoose, { configurationCache: cache }); @@ -402,4 +858,321 @@ describe('code environment registry', () => { mongoose.models.AclEntry.countDocuments({ resourceType: ResourceType.CODE_ENVIRONMENT }), ).resolves.toBe(0); }); + + test('reconciles a pending registration when ACL rollback fails', async () => { + const rollback = jest + .spyOn(AccessControlService.prototype, 'removeAllPermissions') + .mockRejectedValueOnce(new Error('acl store unavailable')); + const registry = createCodeEnvironmentRegistry(mongoose, { + configurationCache: { + get: jest.fn(), + set: jest.fn().mockRejectedValue(new Error('redis unavailable')), + }, + }); + const ownerId = new Types.ObjectId(); + + await expect( + registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'pending-rollback-vm', + name: 'Pending rollback VM', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + workerId: 'shared-deployment-worker', + workerPrincipal: { type: 'deployment', id: 'shared-code-api' }, + }, + }), + ).rejects.toThrow('redis unavailable'); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'pending-rollback-vm' }).lean(), + ).resolves.toMatchObject({ registrationPendingAt: expect.any(Date) }); + + rollback.mockRestore(); + await mongoose.models.CodeEnvironment.updateOne( + { environmentId: 'pending-rollback-vm' }, + { $set: { registrationPendingAt: new Date(Date.now() - 10 * 60_000) } }, + ); + await reconcileCodeEnvironmentLifecycle({ mongoose }); + + await expect( + mongoose.models.CodeEnvironment.countDocuments({ environmentId: 'pending-rollback-vm' }), + ).resolves.toBe(0); + await expect( + mongoose.models.AclEntry.countDocuments({ resourceType: ResourceType.CODE_ENVIRONMENT }), + ).resolves.toBe(0); + }); + + test('claims a stale registration before compensating its worker', async () => { + const methods = createMethods(mongoose); + const ownerId = new Types.ObjectId(); + const environment = await methods.createCodeEnvironment({ + environmentId: 'stale-registration-race', + name: 'Stale registration race', + type: 'attached', + baseURL: 'https://code.example.com/v1', + controlPlaneId: 'shared-code-api', + createdBy: ownerId, + workerId: 'stale-registration-race', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }); + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment._id }, + { $set: { registrationPendingAt: new Date(Date.now() - 10 * 60_000) } }, + ); + let enteredRevoke!: () => void; + let releaseRevoke!: () => void; + const entered = new Promise((resolve) => (enteredRevoke = resolve)); + const release = new Promise((resolve) => (releaseRevoke = resolve)); + const fetchImpl = jest.fn(async () => { + enteredRevoke(); + await release; + return { + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + } as Response; + }); + + const reconciliation = reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret: () => 'administrator-token', + fetchImpl, + }); + await entered; + await expect(methods.completeCodeEnvironmentRegistration(environment._id)).rejects.toThrow( + 'registration could not be committed', + ); + releaseRevoke(); + await reconciliation; + + await expect(mongoose.models.CodeEnvironment.findById(environment._id)).resolves.toBeNull(); + }); + + test('preserves a creator-owned environment referenced by another surviving agent', async () => { + const registry = createCodeEnvironmentRegistry(mongoose); + const methods = createMethods(mongoose); + const ownerId = new Types.ObjectId(); + const teammateId = new Types.ObjectId(); + const environment = await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'shared-deployment-worker', + name: 'Shared deployment worker', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + workerPrincipal: { type: 'deployment', id: 'shared-control-plane' }, + }, + }); + await mongoose.models.Agent.create({ + id: 'agent_survives_owner', + name: 'Surviving agent', + author: teammateId, + model: 'test-model', + provider: 'test-provider', + code_environment_id: environment.id, + }); + + await expect(methods.deleteUserCodeEnvironments(ownerId)).resolves.toBe(0); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: environment.id }), + ).resolves.not.toBeNull(); + await expect( + mongoose.models.AclEntry.countDocuments({ resourceId: environment.resourceId }), + ).resolves.toBeGreaterThan(0); + }); + + test('recovers expired agent reservations and removal leases', async () => { + const methods = createMethods(mongoose); + const ownerId = new Types.ObjectId(); + const environment = await createCodeEnvironmentRegistry(mongoose).register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'expired-lifecycle-leases', + name: 'Expired lifecycle leases', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + }); + const expiredAt = new Date(Date.now() - 1_000); + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment.resourceId }, + { + $set: { + deletionStartedAt: expiredAt, + deletionLeaseId: 'abandoned-removal', + deletionLeaseExpiresAt: expiredAt, + pendingAgentReferences: [{ reservationId: 'abandoned-reference', expiresAt: expiredAt }], + }, + }, + ); + + const claimed = await methods.beginCodeEnvironmentRemoval(environment.resourceId); + + expect(claimed).toMatchObject({ + deletionLeaseId: expect.any(String), + pendingAgentReferences: [], + }); + expect(claimed?.deletionLeaseId).not.toBe('abandoned-removal'); + }); + + test('preserves persisted agent references during interrupted-removal recovery', async () => { + const ownerId = new Types.ObjectId(); + const registry = createCodeEnvironmentRegistry(mongoose); + const environment = await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'referenced-interrupted-removal', + name: 'Referenced interrupted removal', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + }, + }); + await mongoose.models.Agent.create({ + id: 'agent_references_interrupted_removal', + name: 'Referenced environment agent', + author: ownerId, + model: 'test-model', + provider: 'test-provider', + code_environment_id: environment.id, + }); + const expiredAt = new Date(Date.now() - 1_000); + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment.resourceId }, + { + $set: { + deletionStartedAt: expiredAt, + deletionLeaseId: 'abandoned-removal', + deletionLeaseExpiresAt: expiredAt, + }, + }, + ); + + await reconcileCodeEnvironmentLifecycle({ mongoose }); + + await expect( + mongoose.models.CodeEnvironment.findById(environment.resourceId).lean(), + ).resolves.toMatchObject({ environmentId: environment.id }); + await expect( + mongoose.models.CodeEnvironment.findById(environment.resourceId).lean(), + ).resolves.not.toHaveProperty('deletionStartedAt'); + }); + + test('finishes an interrupted local-only removal', async () => { + const ownerId = new Types.ObjectId(); + const registry = createCodeEnvironmentRegistry(mongoose); + const environment = await registry.register({ + actor: { userId: ownerId, role: 'USER', idOnTheSource: null }, + environment: { + id: 'local-interrupted-removal', + name: 'Local interrupted removal', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + workerPrincipal: { type: 'deployment', id: 'shared-code-api' }, + }, + }); + const expiredAt = new Date(Date.now() - 1_000); + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment.resourceId }, + { + $set: { + deletionStartedAt: expiredAt, + deletionLeaseId: 'abandoned-local-removal', + deletionLeaseExpiresAt: expiredAt, + }, + }, + ); + + await reconcileCodeEnvironmentLifecycle({ mongoose }); + + await expect( + mongoose.models.CodeEnvironment.findById(environment.resourceId), + ).resolves.toBeNull(); + await expect( + mongoose.models.AclEntry.countDocuments({ resourceId: environment.resourceId }), + ).resolves.toBe(0); + }); + + test('scopes retired environment ids to their tenant', async () => { + const ownerId = new Types.ObjectId(); + const input = { + environmentId: 'tenant-reusable-id', + name: 'Tenant VM', + type: 'attached' as const, + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + createdBy: ownerId, + }; + + await tenantStorage.run({ tenantId: 'tenant-a' }, async () => { + const methods = createMethods(mongoose); + const created = await methods.createCodeEnvironment(input); + await methods.deleteCodeEnvironmentById(created._id); + await expect(methods.createCodeEnvironment(input)).rejects.toThrow( + 'Code environment id was previously retired', + ); + }); + + await tenantStorage.run({ tenantId: 'tenant-b' }, async () => { + await expect(createMethods(mongoose).createCodeEnvironment(input)).resolves.toMatchObject({ + environmentId: 'tenant-reusable-id', + tenantId: 'tenant-b', + }); + }); + }); + + test('scopes recovered agent references to the environment tenant', async () => { + const ownerId = new Types.ObjectId(); + let tenantAEnvironmentId!: Types.ObjectId; + await tenantStorage.run({ tenantId: 'tenant-a' }, async () => { + const environment = await createMethods(mongoose).createCodeEnvironment({ + environmentId: 'shared-tenant-environment-id', + name: 'Tenant A environment', + type: 'attached', + baseURL: 'https://tenant-a.example.com', + controlPlaneId: 'shared-code-api', + createdBy: ownerId, + }); + tenantAEnvironmentId = environment._id; + await mongoose.models.CodeEnvironment.updateOne( + { _id: environment._id }, + { + $unset: { registrationPendingAt: 1 }, + $set: { + deletionStartedAt: new Date(Date.now() - 10_000), + deletionLeaseId: 'expired-tenant-a-removal', + deletionLeaseExpiresAt: new Date(Date.now() - 1_000), + }, + }, + ); + }); + await tenantStorage.run({ tenantId: 'tenant-b' }, async () => { + await createMethods(mongoose).createAgent({ + id: 'agent_tenant_b_shared_environment_id', + name: 'Tenant B agent', + author: ownerId, + model: 'test-model', + provider: 'test-provider', + code_environment_id: 'shared-tenant-environment-id', + }); + }); + + await reconcileCodeEnvironmentLifecycle({ mongoose }); + + await tenantStorage.run({ tenantId: 'tenant-a' }, async () => { + await expect( + mongoose.models.CodeEnvironment.findById(tenantAEnvironmentId), + ).resolves.toBeNull(); + }); + await tenantStorage.run({ tenantId: 'tenant-b' }, async () => { + await expect( + mongoose.models.Agent.findOne({ id: 'agent_tenant_b_shared_environment_id' }), + ).resolves.not.toBeNull(); + }); + }); }); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index 0625b93ce2..0720644d87 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -23,6 +23,7 @@ export type CodeEnvironmentSummary = { id: string; name: string; type: 'managed' | 'attached'; + canDelete: boolean; }; export type CodeEnvironmentRegistration = { @@ -32,6 +33,7 @@ export type CodeEnvironmentRegistration = { baseURL: string; controlPlaneId: string; workerId?: string; + revocationTokenEnv?: string; workerPrincipal?: { type: 'deployment' | 'tenant' | 'user' | 'role' | 'group'; id: string; @@ -52,6 +54,21 @@ type CachedAccessibleCodeEnvironmentConfiguration = AccessibleCodeEnvironmentCon resourceId: string; }; +export type CodeEnvironmentLifecycleTarget = CodeEnvironmentSummary & { + baseURL: string; + workerId?: string; + controlPlaneId?: string; + revocationTokenEnv?: string; + workerPrincipal?: CodeEnvironmentRegistration['workerPrincipal']; +}; + +function agentReferenceFilter(environmentId: string, tenantId?: string) { + return { + code_environment_id: environmentId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + }; +} + const ENVIRONMENT_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; const WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; const WORKER_PRINCIPAL_ID_PATTERN = /^\S(?:.{0,254}\S)?$/; @@ -85,6 +102,13 @@ export class CodeEnvironmentValidationError extends Error { } } +export class CodeEnvironmentLimitError extends Error { + constructor() { + super('Personal code environment limit reached'); + this.name = 'CodeEnvironmentLimitError'; + } +} + function normalizeRegistration(input: CodeEnvironmentRegistration): CodeEnvironmentRegistration { const id = input.id.trim(); const name = normalizeCodeEnvironmentName(input.name); @@ -117,17 +141,21 @@ function normalizeRegistration(input: CodeEnvironmentRegistration): CodeEnvironm return { ...input, id, name, baseURL, controlPlaneId, workerId }; } -function toSummary(environment: { - _id: Types.ObjectId; - environmentId: string; - name: string; - type: 'managed' | 'attached'; -}): CodeEnvironmentSummary { +function toSummary( + environment: { + _id: Types.ObjectId; + environmentId: string; + name: string; + type: 'managed' | 'attached'; + }, + canDelete = false, +): CodeEnvironmentSummary { return { resourceId: environment._id.toString(), id: environment.environmentId, name: environment.name, type: environment.type, + canDelete, }; } @@ -138,6 +166,7 @@ export function createCodeEnvironmentRegistry( register: (params: { actor: CodeEnvironmentPrincipalContext; environment: CodeEnvironmentRegistration; + maxOwned?: number; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; listAccessibleConfigurations: ( @@ -145,6 +174,12 @@ export function createCodeEnvironmentRegistry( ) => Promise; listRegisteredIds: () => Promise; invalidateAccessibleConfigurations: (tenantId?: string) => Promise; + markRevocationPending: (environmentId: string) => Promise; + remove: (params: { + actor: CodeEnvironmentPrincipalContext; + environmentId: string; + beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; + }) => Promise; } { const methods = createMethods(mongoose); const access = new AccessControlService(mongoose); @@ -189,21 +224,31 @@ export function createCodeEnvironmentRegistry( async function register({ actor, environment: input, + maxOwned, }: { actor: CodeEnvironmentPrincipalContext; environment: CodeEnvironmentRegistration; + maxOwned?: number; }): Promise { const environment = normalizeRegistration(input); - const created = await methods.createCodeEnvironment({ + const createInput = { environmentId: environment.id, name: environment.name, type: environment.type, baseURL: environment.baseURL, controlPlaneId: environment.controlPlaneId, workerId: environment.workerId, + revocationTokenEnv: environment.revocationTokenEnv, workerPrincipal: environment.workerPrincipal, createdBy: new Types.ObjectId(actor.userId), - }); + }; + const created = + maxOwned == null + ? await methods.createCodeEnvironment(createInput) + : await methods.createCodeEnvironmentWithinOwnerLimit(createInput, maxOwned); + if (created == null) { + throw new CodeEnvironmentLimitError(); + } try { const permission = await access.grantPermission({ principalType: PrincipalType.USER, @@ -216,20 +261,26 @@ export function createCodeEnvironmentRegistry( if (permission == null) { throw new Error('Unable to grant code environment ownership'); } - const summary = toSummary(created); + const summary = toSummary(created, true); await invalidateAccessibleConfigurations(); + await methods.completeCodeEnvironmentRegistration(created._id); return summary; } catch (error) { - const cleanup = await Promise.allSettled([ - access.removeAllPermissions({ + let permissionsRemoved = false; + try { + await access.removeAllPermissions({ resourceType: ResourceType.CODE_ENVIRONMENT, resourceId: created._id, - }), - methods.deleteCodeEnvironmentById(created._id), - ]); - for (const result of cleanup) { - if (result.status === 'rejected') { - logger.error('[codeEnvironments] registration rollback failed:', result.reason); + }); + permissionsRemoved = true; + } catch (cleanupError) { + logger.error('[codeEnvironments] registration ACL rollback failed:', cleanupError); + } + if (permissionsRemoved) { + try { + await methods.discardCodeEnvironmentById(created._id); + } catch (cleanupError) { + logger.error('[codeEnvironments] registration record rollback failed:', cleanupError); } } throw error; @@ -251,7 +302,10 @@ export function createCodeEnvironmentRegistry( const userId = actor.userId.toString(); return environments.filter( (environment) => - environment.workerPrincipal?.type !== 'user' || environment.workerPrincipal.id === userId, + environment.registrationPendingAt == null && + environment.deletionStartedAt == null && + environment.deletionCommittedAt == null && + (environment.workerPrincipal?.type !== 'user' || environment.workerPrincipal.id === userId), ); } @@ -259,7 +313,32 @@ export function createCodeEnvironmentRegistry( actor: CodeEnvironmentPrincipalContext, ): Promise { const environments = await findAccessible(actor); - return environments.map(toSummary); + const permissions = await access.getResourcePermissionsMap({ + userId: actor.userId, + role: actor.role ?? '', + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceIds: environments.map((environment) => environment._id), + }); + return environments.map((environment) => { + const permission = permissions.get(environment._id.toString()) ?? 0; + return toSummary(environment, (permission & PermissionBits.DELETE) === PermissionBits.DELETE); + }); + } + + async function markRevocationPending(environmentId: string): Promise { + const CodeEnvironment = mongoose.models.CodeEnvironment; + const result = await CodeEnvironment.updateOne( + { environmentId, deletionCommittedAt: { $exists: false } }, + { + $set: { revocationPendingAt: new Date() }, + $inc: { revocationAttempts: 1 }, + $unset: { revocationLastError: 1, revocationReconcileAfter: 1 }, + }, + ); + if (result.matchedCount !== 1) { + throw new Error('Code environment cleanup target is unavailable'); + } + await invalidateAccessibleConfigurations(); } async function listAccessibleConfigurations( @@ -281,28 +360,41 @@ export function createCodeEnvironmentRegistry( ...configuration }: CachedAccessibleCodeEnvironmentConfiguration): AccessibleCodeEnvironmentConfiguration => configuration; - const load = async (): Promise => { + const load = async (): Promise<{ + configurations: CachedAccessibleCodeEnvironmentConfiguration[]; + hasPendingRegistration: boolean; + }> => { const ids = await findAccessibleResourceIds(principals); const environments = await methods.findCodeEnvironmentsByIds(ids); const userId = actor.userId.toString(); - return environments - .filter( - (environment) => - environment.workerPrincipal?.type !== 'user' || - environment.workerPrincipal.id === userId, - ) - .map((environment) => ({ - resourceId: environment._id.toString(), - id: environment.environmentId, - name: environment.name, - type: environment.type, - baseURL: environment.baseURL, - controlPlaneId: environment.controlPlaneId, - owner: 'principal', - workerId: environment.workerId, - })); + return { + hasPendingRegistration: environments.some( + (environment) => environment.registrationPendingAt != null, + ), + configurations: environments + .filter( + (environment) => + environment.registrationPendingAt == null && + environment.deletionStartedAt == null && + environment.deletionCommittedAt == null && + (environment.workerPrincipal?.type !== 'user' || + environment.workerPrincipal.id === userId), + ) + .map((environment) => ({ + resourceId: environment._id.toString(), + id: environment.environmentId, + name: environment.name, + type: environment.type, + baseURL: environment.baseURL, + controlPlaneId: environment.controlPlaneId, + owner: 'principal', + workerId: environment.workerId, + })), + }; }; - if (configurationCache == null) return (await load()).map(toPublicConfiguration); + if (configurationCache == null) { + return (await load()).configurations.map(toPublicConfiguration); + } const tenant = tenantCacheKey(); const revision = String((await configurationCache.get(revisionKey())) ?? '0'); @@ -326,25 +418,130 @@ export function createCodeEnvironmentRegistry( const accessibleIds = new Set( (await findAccessibleResourceIds(principals)).map((id) => id.toString()), ); - return (cached as CachedAccessibleCodeEnvironmentConfiguration[]) - .filter(({ resourceId }) => accessibleIds.has(resourceId)) + const cachedConfigurations = cached as CachedAccessibleCodeEnvironmentConfiguration[]; + const liveEnvironments = await methods.findCodeEnvironmentsByIds( + cachedConfigurations.map(({ resourceId }) => resourceId), + ); + const userId = actor.userId.toString(); + const liveIds = new Set( + liveEnvironments + .filter( + (environment) => + environment.registrationPendingAt == null && + environment.deletionStartedAt == null && + environment.deletionCommittedAt == null && + (environment.workerPrincipal?.type !== 'user' || + environment.workerPrincipal.id === userId), + ) + .map((environment) => environment._id.toString()), + ); + return cachedConfigurations + .filter(({ resourceId }) => accessibleIds.has(resourceId) && liveIds.has(resourceId)) .map(toPublicConfiguration); } - const configurations = await load(); + const { configurations, hasPendingRegistration } = await load(); const currentRevision = String((await configurationCache.get(revisionKey())) ?? '0'); if (currentRevision !== revision) { return await listAccessibleConfigurations(actor); } + // Registration commits after its first revision write so a failed rollback remains hidden and + // recoverable. Do not cache that transient empty view under the new revision. + if (hasPendingRegistration) { + return configurations.map(toPublicConfiguration); + } await configurationCache.set(key, configurations, CONFIGURATION_CACHE_TTL_MS); return configurations.map(toPublicConfiguration); } + async function remove({ + actor, + environmentId, + beforeDelete, + }: { + actor: CodeEnvironmentPrincipalContext; + environmentId: string; + beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; + }): Promise { + const environment = await methods.findCodeEnvironmentByEnvironmentId(environmentId); + if (environment == null) return null; + if ( + environment.workerPrincipal?.type === 'user' && + environment.workerPrincipal.id !== actor.userId.toString() + ) { + return null; + } + const allowed = await access.checkPermission({ + userId: actor.userId.toString(), + role: actor.role, + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + requiredPermission: PermissionBits.DELETE, + }); + if (!allowed) return null; + const removal = await methods.beginCodeEnvironmentRemoval(environment._id); + if (removal == null) { + throw new CodeEnvironmentInUseError(environmentId); + } + const removalLeaseId = removal.deletionLeaseId; + if (removalLeaseId == null) { + throw new Error('Code environment removal lease is unavailable'); + } + let deletionCommitted = false; + let externalLifecycleStarted = false; + try { + const Agent = mongoose.models.Agent; + if ( + Agent != null && + (await Agent.exists(agentReferenceFilter(environmentId, environment.tenantId))) != null + ) { + throw new CodeEnvironmentInUseError(environmentId); + } + if (beforeDelete != null) { + // From this point onward the remote outcome may be committed even if the request fails. + // Keep the local fence for idempotent reconciliation instead of reopening a dead worker. + externalLifecycleStarted = true; + await beforeDelete({ + ...toSummary(environment), + baseURL: environment.baseURL, + workerId: environment.workerId, + controlPlaneId: environment.controlPlaneId, + revocationTokenEnv: environment.revocationTokenEnv, + workerPrincipal: environment.workerPrincipal, + }); + } + await methods.commitCodeEnvironmentRemoval(environment._id, removalLeaseId); + deletionCommitted = true; + await access.removeAllPermissions({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + const deleted = await methods.deleteCodeEnvironmentById(environment._id); + if (deleted == null) return null; + await invalidateAccessibleConfigurations(); + return toSummary(deleted, true); + } catch (error) { + if (!deletionCommitted && !externalLifecycleStarted) { + await methods.cancelCodeEnvironmentRemoval(environment._id, removalLeaseId); + } + throw error; + } + } + return { register, + markRevocationPending, listAccessible, listAccessibleConfigurations, listRegisteredIds, invalidateAccessibleConfigurations, + remove, }; } + +export class CodeEnvironmentInUseError extends Error { + constructor(public readonly environmentId: string) { + super(`Code environment is still referenced by an agent: ${environmentId}`); + this.name = 'CodeEnvironmentInUseError'; + } +} diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index bd3d9b332c..4e4efd131e 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -1,6 +1,6 @@ import { EModelEndpoint } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; -import { CodeEnvironmentValidationError } from './environments'; +import { CodeEnvironmentLimitError, CodeEnvironmentValidationError } from './environments'; import { createCodeEnvironmentHttpHandlers } from './http'; function response() { @@ -20,11 +20,85 @@ function response() { } describe('code environment HTTP handlers', () => { + test('does not advertise principal pairing when Code API principal auth is disabled', async () => { + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'principal-workers', + name: 'Principal workers', + type: 'attached', + baseURL: 'https://code.example.com/v1', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { + register: jest.fn(), + listAccessible: jest.fn().mockResolvedValue([]), + remove: jest.fn(), + }, + principalAuthEnabled: () => false, + }); + const res = response(); + + await handlers.list( + { user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' } } as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(res.body).toEqual({ environments: [], controlPlanes: [] }); + }); + + test('lists only principal control planes present in the caller effective policy', async () => { + const getAppConfig = jest.fn().mockResolvedValue({ + endpoints: { [EModelEndpoint.agents]: { statefulCodeSessions: { environments: [] } } }, + } as unknown as AppConfig); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig, + registry: { + register: jest.fn(), + listAccessible: jest.fn().mockResolvedValue([]), + remove: jest.fn(), + }, + principalAuthEnabled: () => true, + }); + const res = response(); + + await handlers.list( + { + user: { + id: '68b2f0c498f24c1e78fa0001', + role: 'USER', + tenantId: 'tenant-1', + }, + } as never, + res as never, + ); + + expect(res.body).toEqual({ environments: [], controlPlanes: [] }); + expect(getAppConfig).toHaveBeenCalledWith({ + role: 'USER', + userId: '68b2f0c498f24c1e78fa0001', + idOnTheSource: undefined, + tenantId: 'tenant-1', + failClosed: true, + }); + }); + test('returns 400 when registration has no request body', async () => { const register = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn(), - registry: { register, listAccessible: jest.fn() }, + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, }); const res = response(); @@ -77,7 +151,7 @@ describe('code environment HTTP handlers', () => { } as AppConfig; const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn().mockResolvedValue(appConfig), - registry: { register, listAccessible: jest.fn() }, + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, createEnvironmentId: () => 'code-generated', readSecret: jest.fn(() => 'administrator-token'), resolveTenantId: jest.fn(() => 'tenant-1'), @@ -119,27 +193,368 @@ describe('code environment HTTP handlers', () => { role: 'USER', idOnTheSource: null, }, + maxOwned: 5, environment: { id: 'code-generated', name: 'Personal VM', - type: 'attached', + type: 'attached' as const, baseURL: 'https://code.librechat.example/v1', workerId: 'code-generated', controlPlaneId: 'shared-code-api', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, }, }); expect(res.body).toEqual({ environment: expect.objectContaining({ id: 'code-generated' }), - pairing: expect.objectContaining({ workerId: 'code-generated', code: 'a'.repeat(32) }), + pairing: expect.objectContaining({ + workerId: 'code-generated', + code: 'a'.repeat(32), + endpoint: 'https://code.librechat.example/v1', + }), }); }); + test('revokes an upstream pairing when the atomic owner quota is exhausted', async () => { + const fetchImpl = jest.fn( + async (input: string | URL | Request) => + ({ + ok: true, + json: async () => + String(input).endsWith('/bridge/pairings') + ? { + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + } + : { protocolVersion: 1, revoked: true }, + }) as Response, + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { + register: jest.fn().mockRejectedValue(new CodeEnvironmentLimitError()), + listAccessible: jest.fn(), + remove: jest.fn(), + }, + createEnvironmentId: () => 'code-generated', + readSecret: () => 'administrator-token', + principalAuthEnabled: () => true, + principalAuthReady: jest.fn().mockResolvedValue(undefined), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(409); + expect(res.body).toEqual({ error: 'Personal code environment limit reached' }); + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.librechat.example/v1/bridge/workers/code-generated/revoke', + expect.objectContaining({ method: 'POST' }), + ); + }); + + test('returns 503 without issuing a pairing when the initial principal check is unavailable', async () => { + const fetchImpl = jest.fn(); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, + createEnvironmentId: () => 'code-generated', + readSecret: () => 'administrator-token', + principalAuthEnabled: () => true, + principalAuthReady: jest.fn(), + principalIsActive: jest.fn().mockRejectedValue(new Error('user store unavailable')), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(503); + expect(res.body).toEqual({ error: 'Account status could not be confirmed' }); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + test('revokes through the registry fence when the principal becomes inactive after registration', async () => { + const remove = jest.fn( + async ({ beforeDelete }: { beforeDelete?: (target: never) => Promise }) => { + await beforeDelete?.({} as never); + return { id: 'code-generated' } as never; + }, + ); + const principalIsActive = jest + .fn() + .mockResolvedValueOnce(true) + .mockResolvedValueOnce(true) + .mockResolvedValueOnce(false); + const fetchImpl = jest.fn( + async (input: string | URL | Request) => + ({ + ok: true, + json: async () => + String(input).endsWith('/bridge/pairings') + ? { + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + } + : { protocolVersion: 1, revoked: true }, + }) as Response, + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { + register: jest.fn().mockResolvedValue({ id: 'code-generated' }), + listAccessible: jest.fn(), + remove, + }, + createEnvironmentId: () => 'code-generated', + readSecret: () => 'administrator-token', + principalAuthEnabled: () => true, + principalAuthReady: jest.fn(), + principalIsActive, + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(409); + expect(remove).toHaveBeenCalledWith( + expect.objectContaining({ + environmentId: 'code-generated', + beforeDelete: expect.any(Function), + }), + ); + expect(fetchImpl).toHaveBeenLastCalledWith( + 'https://code.librechat.example/v1/bridge/workers/code-generated/revoke', + expect.objectContaining({ method: 'POST' }), + ); + }); + + test('persists cleanup intent before compensating a registered pairing after removal fails', async () => { + const markRevocationPending = jest.fn().mockResolvedValue(undefined); + const remove = jest.fn().mockRejectedValue(new Error('registry removal failed')); + const principalIsActive = jest + .fn() + .mockResolvedValueOnce(true) + .mockResolvedValueOnce(true) + .mockResolvedValueOnce(false); + const fetchImpl = jest.fn( + async (input: string | URL | Request) => + ({ + ok: true, + json: async () => + String(input).endsWith('/bridge/pairings') + ? { + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + } + : { protocolVersion: 1, revoked: true }, + }) as Response, + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { + register: jest.fn().mockResolvedValue({ id: 'code-generated' }), + listAccessible: jest.fn(), + remove, + markRevocationPending, + }, + createEnvironmentId: () => 'code-generated', + readSecret: () => 'administrator-token', + principalAuthEnabled: () => true, + principalAuthReady: jest.fn(), + principalIsActive, + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(500); + expect(markRevocationPending).toHaveBeenCalledWith('code-generated'); + expect(markRevocationPending.mock.invocationCallOrder[0]).toBeLessThan( + fetchImpl.mock.invocationCallOrder[fetchImpl.mock.invocationCallOrder.length - 1] ?? 0, + ); + }); + + test('revokes an issued pairing when the post-issue principal check is unavailable', async () => { + const register = jest.fn(); + const principalIsActive = jest + .fn() + .mockResolvedValueOnce(true) + .mockRejectedValueOnce(new Error('user store unavailable')); + const fetchImpl = jest.fn( + async (input: string | URL | Request) => + ({ + ok: true, + json: async () => + String(input).endsWith('/bridge/pairings') + ? { + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + } + : { protocolVersion: 1, revoked: true }, + }) as Response, + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, + createEnvironmentId: () => 'code-generated', + readSecret: () => 'administrator-token', + principalAuthEnabled: () => true, + principalAuthReady: jest.fn(), + principalIsActive, + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(503); + expect(register).not.toHaveBeenCalled(); + expect(fetchImpl).toHaveBeenLastCalledWith( + 'https://code.librechat.example/v1/bridge/workers/code-generated/revoke', + expect.objectContaining({ method: 'POST' }), + ); + }); + test('rejects self-service pairing without principal-aware Code API auth', async () => { const fetchImpl = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn(), - registry: { register: jest.fn(), listAccessible: jest.fn() }, + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, principalAuthEnabled: jest.fn(() => false), fetchImpl, }); @@ -185,7 +600,7 @@ describe('code environment HTTP handlers', () => { ); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig, - registry: { register: jest.fn(), listAccessible: jest.fn() }, + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, readSecret: jest.fn(() => 'administrator-token'), principalAuthEnabled: jest.fn(() => true), principalAuthReady: jest.fn(), @@ -219,7 +634,7 @@ describe('code environment HTTP handlers', () => { if (options.baseOnly === true) return {} as AppConfig; throw new Error('authorization unavailable'); }), - registry: { register: jest.fn(), listAccessible: jest.fn() }, + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, principalAuthEnabled: jest.fn(() => true), principalAuthReady: jest.fn(), fetchImpl, @@ -294,6 +709,7 @@ describe('code environment HTTP handlers', () => { registry: { register: jest.fn().mockResolvedValue({ id: 'code-generated' }), listAccessible: jest.fn(), + remove: jest.fn(), }, createEnvironmentId: () => 'code-generated', readSecret, @@ -325,7 +741,7 @@ describe('code environment HTTP handlers', () => { const fetchImpl = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn(), - registry: { register: jest.fn(), listAccessible: jest.fn() }, + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, principalAuthEnabled: jest.fn(() => true), principalAuthReady: jest.fn(() => { throw new Error('invalid signing key'); @@ -374,7 +790,7 @@ describe('code environment HTTP handlers', () => { } as AppConfig; const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn().mockResolvedValue(appConfig), - registry: { register, listAccessible: jest.fn() }, + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, createEnvironmentId: () => 'personal-vm', }); const req = { @@ -409,6 +825,96 @@ describe('code environment HTTP handlers', () => { }); }); + test('removes a fixed environment when its principal becomes inactive after registration', async () => { + const register = jest.fn().mockResolvedValue({ + resourceId: '68b2f0c498f24c1e78fa0111', + id: 'personal-vm', + name: 'Personal VM', + type: 'attached', + }); + const remove = jest.fn().mockResolvedValue({ id: 'personal-vm' }); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example', + owner: 'deployment', + pairing: { workerId: 'deployment-worker' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register, listAccessible: jest.fn(), remove }, + createEnvironmentId: () => 'personal-vm', + principalIsActive: jest.fn().mockResolvedValueOnce(true).mockResolvedValueOnce(false), + }); + const res = response(); + + await handlers.register( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(409); + expect(remove).toHaveBeenCalledWith({ + actor: { + userId: '68b2f0c498f24c1e78fa0001', + role: 'USER', + idOnTheSource: null, + }, + environmentId: 'personal-vm', + }); + }); + + test('returns 503 without fixed registration when the principal check is unavailable', async () => { + const register = jest.fn(); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'shared-code-api', + name: 'Shared Code API', + type: 'attached', + baseURL: 'https://code.librechat.example', + owner: 'deployment', + pairing: { workerId: 'deployment-worker' }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, + principalIsActive: jest.fn().mockRejectedValue(new Error('user store unavailable')), + }); + const res = response(); + + await handlers.register( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(503); + expect(register).not.toHaveBeenCalled(); + }); + test('does not register a fixed environment on a self-service-only control plane', async () => { const register = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ @@ -434,7 +940,7 @@ describe('code environment HTTP handlers', () => { }, }, } as AppConfig), - registry: { register, listAccessible: jest.fn() }, + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, }); const res = response(); @@ -475,7 +981,7 @@ describe('code environment HTTP handlers', () => { }, }, } as AppConfig), - registry: { register: jest.fn(), listAccessible: jest.fn() }, + registry: { register: jest.fn(), listAccessible: jest.fn(), remove: jest.fn() }, readSecret: jest.fn(() => 'administrator-token'), principalAuthEnabled: jest.fn(() => true), principalAuthReady: jest.fn(), @@ -495,6 +1001,79 @@ describe('code environment HTTP handlers', () => { expect(fetchImpl).not.toHaveBeenCalled(); }); + test('revokes a user-bound worker before deleting its registry entry', async () => { + const remove = jest.fn(async ({ beforeDelete }) => { + await beforeDelete({ + resourceId: '68b2f0c498f24c1e78fa0111', + id: 'code-generated', + name: 'Personal VM', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + workerId: 'code-generated', + controlPlaneId: 'self-service', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, + }); + return { + resourceId: '68b2f0c498f24c1e78fa0111', + id: 'code-generated', + name: 'Personal VM', + type: 'attached' as const, + canDelete: true, + }; + }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + owner: 'deployment', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register: jest.fn(), listAccessible: jest.fn(), remove }, + readSecret: jest.fn(() => 'administrator-token'), + fetchImpl, + }); + const res = response(); + + await handlers.remove( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + params: { environmentId: 'code-generated' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(200); + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.librechat.example/v1/bridge/workers/code-generated/revoke', + expect.objectContaining({ + method: 'POST', + headers: { Authorization: 'Bearer administrator-token' }, + }), + ); + expect(remove).toHaveBeenCalled(); + }); + test.each([ { error: new Error('mongo connection details'), @@ -526,7 +1105,11 @@ describe('code environment HTTP handlers', () => { }, }, } as AppConfig), - registry: { register: jest.fn().mockRejectedValue(error), listAccessible: jest.fn() }, + registry: { + register: jest.fn().mockRejectedValue(error), + listAccessible: jest.fn(), + remove: jest.fn(), + }, }); const res = response(); diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index b6c485759d..60c0766a8c 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -3,6 +3,7 @@ import { EModelEndpoint } from 'librechat-data-provider'; import { logger, type AppConfig } from '@librechat/data-schemas'; import type { Response } from 'express'; import type { + CodeEnvironmentLifecycleTarget, CodeEnvironmentPrincipalContext, CodeEnvironmentRegistration, CodeEnvironmentSummary, @@ -10,21 +11,39 @@ import type { import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; import type { CodeBridgeFetch } from './bridge'; +import { + CodeEnvironmentInUseError, + CodeEnvironmentLimitError, + CodeEnvironmentValidationError, + normalizeCodeEnvironmentName, +} from './environments'; +import { + CodeBridgeLifecycleError, + CodeBridgePairingError, + createCodeBridgePairing, + readCodeBridgeSecret, + revokeCodeBridgeWorker, +} from './bridge'; import { assertCodeApiJwtSigningReady, getCodeApiTenantId, isCodeApiJwtAuthEnabled, } from '~/auth/codeapi'; -import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; -import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; import { getAppConfigOptionsFromUser } from '~/app/service'; type Registry = { register: (params: { actor: CodeEnvironmentPrincipalContext; environment: CodeEnvironmentRegistration; + maxOwned?: number; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; + remove: (params: { + actor: CodeEnvironmentPrincipalContext; + environmentId: string; + beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; + }) => Promise; + markRevocationPending?: (environmentId: string) => Promise; }; type StatefulCodeConfig = NonNullable< @@ -42,6 +61,8 @@ export interface CodeEnvironmentHttpDeps { resolveTenantId?: (req: ServerRequest) => string; principalAuthEnabled?: () => boolean; principalAuthReady?: () => Promise | void; + principalIsActive?: (userId: string) => Promise; + maxPrincipalEnvironments?: number; fetchImpl?: CodeBridgeFetch; } @@ -80,6 +101,28 @@ function configuredPrincipalControlPlane( ); } +function principalControlPlanes(appConfig: AppConfig): Array<{ id: string; name: string }> { + return ( + appConfig.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments + ?.filter( + (environment) => + environment.type === 'attached' && + environment.owner === 'deployment' && + environment.pairing?.allowPrincipalWorkers === true, + ) + .map(({ id, name }) => ({ id, name })) ?? [] + ); +} + +class CodeEnvironmentLifecycleHttpError extends Error { + constructor( + public readonly status: number, + message: string, + ) { + super(message); + } +} + function pairingErrorResponse(error: unknown, res: Response): Response { if (!(error instanceof CodeBridgePairingError)) { return res.status(502).json({ error: 'Code API pairing request failed' }); @@ -105,20 +148,36 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) list: (req: ServerRequest, res: Response) => Promise; register: (req: ServerRequest, res: Response) => Promise; pair: (req: ServerRequest, res: Response) => Promise; + remove: (req: ServerRequest, res: Response) => Promise; } { const createEnvironmentId = deps.createEnvironmentId ?? (() => `code-${nanoid(20)}`); const readSecret = deps.readSecret ?? readCodeBridgeSecret; const resolveTenantId = deps.resolveTenantId ?? getCodeApiTenantId; const principalAuthEnabled = deps.principalAuthEnabled ?? isCodeApiJwtAuthEnabled; const principalAuthReady = deps.principalAuthReady ?? assertCodeApiJwtSigningReady; + const principalIsActive = deps.principalIsActive ?? (async () => true); + const maxPrincipalEnvironments = deps.maxPrincipalEnvironments ?? 5; async function list(req: ServerRequest, res: Response): Promise { const principal = actor(req); if (principal == null) { return res.status(401).json({ error: 'Authentication required' }); } - const environments = await deps.registry.listAccessible(principal); - return res.status(200).json({ environments }); + let environments: CodeEnvironmentSummary[]; + let appConfig: AppConfig; + try { + [environments, appConfig] = await Promise.all([ + deps.registry.listAccessible(principal), + deps.getAppConfig({ ...getAppConfigOptionsFromUser(req.user), failClosed: true }), + ]); + } catch (error) { + logger.error('[codeEnvironments] discovery policy resolution failed:', error); + return res.status(503).json({ error: 'Code environment policy is unavailable' }); + } + return res.status(200).json({ + environments, + controlPlanes: principalAuthEnabled() ? principalControlPlanes(appConfig) : [], + }); } async function register(req: ServerRequest, res: Response): Promise { @@ -165,6 +224,17 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) return res.status(404).json({ error: 'Code control plane was not found' }); } + let activeBeforeRegistration: boolean; + try { + activeBeforeRegistration = await principalIsActive(principal.userId.toString()); + } catch (error) { + logger.error('[codeEnvironments] pre-registration principal check failed:', error); + return res.status(503).json({ error: 'Account status could not be confirmed' }); + } + if (!activeBeforeRegistration) { + return res.status(409).json({ error: 'Account deletion is already in progress' }); + } + try { const environment = await deps.registry.register({ actor: principal, @@ -173,11 +243,25 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) name, type: 'attached', baseURL: controlPlane.baseURL, - controlPlaneId, workerId: controlPlane.pairing?.workerId, + controlPlaneId: controlPlane.id, workerPrincipal: { type: 'deployment', id: controlPlane.id }, }, }); + let activeAfterRegistration = false; + let principalCheckUnavailable = false; + try { + activeAfterRegistration = await principalIsActive(principal.userId.toString()); + } catch (error) { + principalCheckUnavailable = true; + logger.error('[codeEnvironments] post-registration principal check failed:', error); + } + if (!activeAfterRegistration) { + await deps.registry.remove({ actor: principal, environmentId: environment.id }); + return principalCheckUnavailable + ? res.status(503).json({ error: 'Account status could not be confirmed' }) + : res.status(409).json({ error: 'Account deletion is already in progress' }); + } return res.status(201).json({ environment }); } catch (error) { const duplicate = @@ -253,6 +337,16 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } const workerId = createEnvironmentId(); + let activeBeforePairing: boolean; + try { + activeBeforePairing = await principalIsActive(principal.userId.toString()); + } catch (error) { + logger.error('[codeEnvironments] pre-pairing principal check failed:', error); + return res.status(503).json({ error: 'Account status could not be confirmed' }); + } + if (!activeBeforePairing) { + return res.status(409).json({ error: 'Account deletion is already in progress' }); + } let pairing; try { pairing = await createCodeBridgePairing({ @@ -269,9 +363,36 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) return pairingErrorResponse(error, res); } + let activeAfterPairing = false; + let principalCheckUnavailable = false; + try { + activeAfterPairing = await principalIsActive(principal.userId.toString()); + } catch (error) { + principalCheckUnavailable = true; + logger.error('[codeEnvironments] post-pairing principal check failed:', error); + } + if (!activeAfterPairing) { + try { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + } catch (error) { + logger.error('[codeEnvironments] orphaned pairing compensation failed:', error); + return res.status(502).json({ error: 'Code worker pairing could not be compensated' }); + } + return principalCheckUnavailable + ? res.status(503).json({ error: 'Account status could not be confirmed' }) + : res.status(409).json({ error: 'Account deletion is already in progress' }); + } + + let registrationCommitted = false; try { const environment = await deps.registry.register({ actor: principal, + maxOwned: maxPrincipalEnvironments, environment: { id: workerId, name, @@ -279,18 +400,73 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) baseURL: controlPlane.baseURL, workerId, controlPlaneId: controlPlane.id, + revocationTokenEnv: tokenEnv, workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); + registrationCommitted = true; + let activeAfterRegistration = false; + try { + activeAfterRegistration = await principalIsActive(principal.userId.toString()); + } catch (error) { + logger.error('[codeEnvironments] post-registration principal check failed:', error); + } + if (!activeAfterRegistration) { + const removed = await deps.registry.remove({ + actor: principal, + environmentId: workerId, + beforeDelete: async () => { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + }, + }); + if (removed == null) { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + } + return res.status(409).json({ error: 'Account is unavailable for code worker pairing' }); + } return res.status(201).json({ environment, pairing: { workerId: pairing.workerId, code: pairing.code, expiresAt: pairing.expiresAt, + endpoint: controlPlane.baseURL, }, }); } catch (error) { + if (registrationCommitted) { + try { + if (deps.registry.markRevocationPending == null) { + throw new Error('Code environment cleanup scheduling is unavailable'); + } + await deps.registry.markRevocationPending(workerId); + } catch (markerError) { + logger.error('[codeEnvironments] failed to persist pairing cleanup intent:', markerError); + return res.status(503).json({ error: 'Code environment cleanup could not be scheduled' }); + } + } + try { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + } catch { + return res.status(502).json({ + error: 'Code environment registration failed and its pairing could not be revoked', + }); + } const duplicate = typeof error === 'object' && error != null && @@ -299,6 +475,9 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (duplicate) { return res.status(409).json({ error: 'Code environment already exists' }); } + if (error instanceof CodeEnvironmentLimitError) { + return res.status(409).json({ error: error.message }); + } if (error instanceof CodeEnvironmentValidationError) { return res.status(400).json({ error: error.message }); } @@ -307,5 +486,58 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } } - return { list, register, pair }; + async function remove(req: ServerRequest, res: Response): Promise { + const principal = actor(req); + if (principal == null) { + return res.status(401).json({ error: 'Authentication required' }); + } + const params = req.params as Record; + const environmentId = typeof params.environmentId === 'string' ? params.environmentId : ''; + if (!environmentId) { + return res.status(400).json({ error: 'Code environment id is required' }); + } + try { + const environment = await deps.registry.remove({ + actor: principal, + environmentId, + beforeDelete: async (target) => { + if (target.workerPrincipal?.type !== 'user' || target.workerId == null) return; + const tokenEnv = target.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token) { + throw new CodeEnvironmentLifecycleHttpError( + 503, + 'Code environment revocation is not configured', + ); + } + await revokeCodeBridgeWorker({ + baseURL: target.baseURL, + token, + workerId: target.workerId, + fetchImpl: deps.fetchImpl, + }); + }, + }); + if (environment == null) { + return res.status(404).json({ error: 'Code environment was not found' }); + } + return res.status(200).json({ environment }); + } catch (error) { + if (error instanceof CodeEnvironmentLifecycleHttpError) { + return res.status(error.status).json({ error: error.message }); + } + if (error instanceof CodeEnvironmentInUseError) { + return res.status(409).json({ error: error.message }); + } + if (error instanceof CodeBridgeLifecycleError) { + return res.status(error.reason === 'timeout' ? 504 : 502).json({ + error: 'Code environment worker could not be revoked', + ...(error.upstreamStatus != null ? { upstreamStatus: error.upstreamStatus } : {}), + }); + } + throw error; + } + } + + return { list, register, pair, remove }; } diff --git a/packages/api/src/code/index.ts b/packages/api/src/code/index.ts index 0e405e9f19..787a58b1de 100644 --- a/packages/api/src/code/index.ts +++ b/packages/api/src/code/index.ts @@ -2,3 +2,4 @@ export * from './environments'; export * from './http'; export * from './config'; export * from './bridge'; +export * from './lifecycle'; diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts new file mode 100644 index 0000000000..f82de80e0b --- /dev/null +++ b/packages/api/src/code/lifecycle.ts @@ -0,0 +1,440 @@ +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 { CodeBridgeFetch } from './bridge'; +import { readCodeBridgeSecret, revokeCodeBridgeWorker } from './bridge'; + +const RECONCILE_INTERVAL_MS = 60_000; +const RECONCILE_LEASE_MS = 2 * 60_000; +const REGISTRATION_STALE_MS = 5 * 60_000; +const REGISTRATION_RETRY_MS = 5 * 60_000; +const REVOCATION_RETRY_MS = 5 * 60_000; +let reconcileTimer: NodeJS.Timeout | undefined; +let reconcileInFlight: Promise | undefined; + +function agentReferenceFilter(environmentId: string, tenantId?: string) { + return { + code_environment_id: environmentId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + }; +} + +export async function reconcileCodeEnvironmentLifecycle({ + mongoose, + readSecret = readCodeBridgeSecret, + fetchImpl, + limit = 25, +}: { + mongoose: typeof import('mongoose'); + readSecret?: (name: string) => string | undefined; + fetchImpl?: CodeBridgeFetch; + limit?: number; +}): Promise { + await runAsSystem(async () => { + const CodeEnvironment = mongoose.models.CodeEnvironment; + const AclEntry = mongoose.models.AclEntry; + if (CodeEnvironment == null || AclEntry == null) return; + const methods = createMethods(mongoose); + const now = new Date(); + await CodeEnvironment.updateMany( + {}, + { $pull: { pendingAgentReferences: { expiresAt: { $lte: now } } } }, + ); + const expiredRemovals = await CodeEnvironment.find({ + deletionCommittedAt: { $exists: false }, + deletionLeaseExpiresAt: { $lte: now }, + }) + .sort({ _id: 1 }) + .limit(limit) + .select('_id') + .lean>>(); + for (const candidate of expiredRemovals) { + const environment = await methods.beginCodeEnvironmentRemoval(candidate._id); + if (environment == null || environment.deletionLeaseId == null) continue; + const leaseId = environment.deletionLeaseId; + const Agent = mongoose.models.Agent; + if ( + Agent != null && + (await Agent.exists( + agentReferenceFilter(environment.environmentId, environment.tenantId), + )) != null + ) { + await methods.cancelCodeEnvironmentRemoval(environment._id, leaseId); + continue; + } + try { + if (environment.workerPrincipal?.type === 'user') { + const tokenEnv = environment.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token || environment.workerId == null) continue; + await revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }); + } + await methods.commitCodeEnvironmentRemoval(environment._id, leaseId); + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + await methods.deleteCodeEnvironmentById(environment._id); + } catch (error) { + logger.error('[code-environments] interrupted removal reconciliation failed:', error); + } + } + const staleRegistration = new Date(Date.now() - REGISTRATION_STALE_MS); + const registrationCandidates = await CodeEnvironment.find({ + registrationPendingAt: { $lte: staleRegistration }, + $and: [ + { + $or: [ + { registrationReconcileAfter: { $exists: false } }, + { registrationReconcileAfter: { $lte: now } }, + ], + }, + { + $or: [ + { registrationLeaseExpiresAt: { $exists: false } }, + { registrationLeaseExpiresAt: { $lte: now } }, + ], + }, + ], + }) + .sort({ registrationReconcileAfter: 1, _id: 1 }) + .limit(limit) + .select('_id') + .lean>>(); + for (const candidate of registrationCandidates) { + const leaseId = new mongoose.Types.ObjectId().toHexString(); + const leaseNow = new Date(); + const environment = await CodeEnvironment.findOneAndUpdate( + { + _id: candidate._id, + registrationPendingAt: { $lte: staleRegistration }, + $and: [ + { + $or: [ + { registrationReconcileAfter: { $exists: false } }, + { registrationReconcileAfter: { $lte: leaseNow } }, + ], + }, + { + $or: [ + { registrationLeaseExpiresAt: { $exists: false } }, + { registrationLeaseExpiresAt: { $lte: leaseNow } }, + ], + }, + ], + }, + { + $set: { + registrationLeaseId: leaseId, + registrationLeaseExpiresAt: new Date(leaseNow.getTime() + RECONCILE_LEASE_MS), + }, + }, + { new: true }, + ).lean(); + if (environment == null) continue; + const deferRegistration = async (): Promise => { + await CodeEnvironment.updateOne( + { _id: environment._id, registrationLeaseId: leaseId }, + { + $set: { registrationReconcileAfter: new Date(Date.now() + REGISTRATION_RETRY_MS) }, + $unset: { registrationLeaseId: 1, registrationLeaseExpiresAt: 1 }, + }, + ); + }; + const tokenEnv = environment.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (environment.workerId != null && environment.workerPrincipal?.type !== 'deployment') { + if (!token) { + await deferRegistration(); + continue; + } + try { + await revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }); + } catch { + await deferRegistration(); + continue; + } + } + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + await CodeEnvironment.deleteOne({ + _id: environment._id, + registrationLeaseId: leaseId, + }); + } + + const committedDeletions = await CodeEnvironment.find({ + deletionCommittedAt: { $exists: true }, + }) + .limit(limit) + .lean(); + for (const environment of committedDeletions) { + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + await methods.deleteCodeEnvironmentById(environment._id); + } + + /** A failed marker write after an account deletion must not strand live worker + * credentials. The missing owner is itself durable retry intent, so discover + * unmarked orphans independently of the account-deletion request path. */ + const orphanedEnvironments = await CodeEnvironment.aggregate<{ + _id: CodeEnvironmentDocument['_id']; + createdBy: CodeEnvironmentDocument['createdBy']; + workerId?: string; + workerPrincipal?: CodeEnvironmentDocument['workerPrincipal']; + }>([ + { + $match: { + revocationPendingAt: { $exists: false }, + deletionCommittedAt: { $exists: false }, + }, + }, + { + $lookup: { + from: mongoose.models.User.collection.name, + localField: 'createdBy', + foreignField: '_id', + as: 'owner', + }, + }, + { $match: { owner: { $size: 0 } } }, + { $sort: { _id: 1 } }, + { $limit: limit }, + { $project: { _id: 1, createdBy: 1, workerId: 1, workerPrincipal: 1 } }, + ]); + const revocationTargets = orphanedEnvironments.filter( + ({ workerId, workerPrincipal }) => workerId != null && workerPrincipal?.type === 'user', + ); + if (revocationTargets.length > 0) { + await CodeEnvironment.updateMany( + { _id: { $in: revocationTargets.map(({ _id }) => _id) } }, + { + $set: { revocationPendingAt: now }, + $inc: { revocationAttempts: 1 }, + $unset: { revocationLastError: 1, revocationReconcileAfter: 1 }, + }, + ); + } + for (const creatorId of new Set( + orphanedEnvironments.map(({ createdBy }) => createdBy.toHexString()), + )) { + await methods.deleteUserCodeEnvironments(creatorId); + } + + const candidates = await CodeEnvironment.find({ + revocationPendingAt: { $exists: true }, + $and: [ + { + $or: [ + { revocationReconcileAfter: { $exists: false } }, + { revocationReconcileAfter: { $lte: now } }, + ], + }, + { + $or: [ + { revocationLeaseExpiresAt: { $exists: false } }, + { revocationLeaseExpiresAt: { $lte: now } }, + ], + }, + ], + }) + .sort({ revocationReconcileAfter: 1, _id: 1 }) + .limit(limit) + .select('_id') + .lean>>(); + for (const candidate of candidates) { + const now = new Date(); + const leaseId = new mongoose.Types.ObjectId().toHexString(); + const environment = await CodeEnvironment.findOneAndUpdate( + { + _id: candidate._id, + revocationPendingAt: { $exists: true }, + $and: [ + { + $or: [ + { revocationReconcileAfter: { $exists: false } }, + { revocationReconcileAfter: { $lte: now } }, + ], + }, + { + $or: [ + { revocationLeaseExpiresAt: { $exists: false } }, + { revocationLeaseExpiresAt: { $lte: now } }, + ], + }, + ], + }, + { + $set: { + revocationLeaseId: leaseId, + revocationLeaseExpiresAt: new Date(now.getTime() + RECONCILE_LEASE_MS), + }, + $inc: { revocationAttempts: 1 }, + }, + { new: true }, + ).lean(); + if (environment == null || environment.workerId == null) continue; + const tokenEnv = environment.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + try { + if (!token) throw new Error('Code environment revocation token is unavailable'); + await revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }); + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + await methods.deleteCodeEnvironmentById(environment._id); + } catch (error) { + const message = (error instanceof Error ? error.message : 'Worker revocation failed').slice( + 0, + 500, + ); + await CodeEnvironment.updateOne( + { _id: environment._id, revocationLeaseId: leaseId }, + { + $set: { + revocationLastError: message, + revocationReconcileAfter: new Date(Date.now() + REVOCATION_RETRY_MS), + }, + $unset: { revocationLeaseId: 1, revocationLeaseExpiresAt: 1 }, + }, + ); + } + } + }); +} + +export function startCodeEnvironmentLifecycleReconciler( + options: Parameters[0], +): void { + if (reconcileTimer != null) return; + const run = (): void => { + if (reconcileInFlight != null) return; + const current = reconcileCodeEnvironmentLifecycle(options) + .catch((error) => { + logger.error('[code-environments] lifecycle reconciliation failed:', error); + }) + .finally(() => { + if (reconcileInFlight === current) reconcileInFlight = undefined; + }); + reconcileInFlight = current; + }; + run(); + reconcileTimer = setInterval(run, RECONCILE_INTERVAL_MS); + reconcileTimer.unref(); +} + +export async function revokeUserCodeEnvironmentWorkers({ + mongoose, + userId, + appConfig, + readSecret = readCodeBridgeSecret, + fetchImpl, +}: { + mongoose: typeof import('mongoose'); + userId: string; + appConfig: AppConfig; + readSecret?: (name: string) => string | undefined; + fetchImpl?: CodeBridgeFetch; +}): Promise { + const environments = await createMethods(mongoose).findCodeEnvironmentsByCreator(userId); + const controlPlanes = + appConfig.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments ?? []; + const targets = environments.filter( + (environment) => environment.workerPrincipal?.type === 'user' && environment.workerId != null, + ); + const CodeEnvironment = mongoose.models.CodeEnvironment; + if (targets.length > 0 && CodeEnvironment != null) { + await CodeEnvironment.updateMany( + { _id: { $in: targets.map(({ _id }) => _id) } }, + { + $set: { revocationPendingAt: new Date() }, + $inc: { revocationAttempts: 1 }, + $unset: { revocationLastError: 1 }, + }, + ); + } + const results = await Promise.allSettled( + targets.map(async (environment) => { + const workerId = environment.workerId; + if (workerId == null) { + throw new Error( + `Code environment worker id is unavailable for ${environment.environmentId}`, + ); + } + const fallbackControlPlane = controlPlanes.find( + (candidate) => + candidate.id === environment.controlPlaneId && + candidate.owner === 'deployment' && + candidate.type === 'attached' && + candidate.pairing?.allowPrincipalWorkers === true && + candidate.baseURL === environment.baseURL, + ); + const tokenEnv = environment.revocationTokenEnv ?? fallbackControlPlane?.pairing?.tokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token) { + throw new Error( + `Code environment revocation is unavailable for ${environment.environmentId}`, + ); + } + await revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId, + fetchImpl, + }); + return environment.environmentId; + }), + ); + let revoked = 0; + for (const [index, result] of results.entries()) { + const target = targets[index]; + if (result.status === 'fulfilled') { + revoked += 1; + await CodeEnvironment?.updateOne( + { _id: target._id }, + { + $set: { deletionCommittedAt: new Date() }, + $unset: { + revocationPendingAt: 1, + revocationAttempts: 1, + revocationLastError: 1, + revocationReconcileAfter: 1, + }, + }, + ); + continue; + } + const failureMessage = ( + result.reason instanceof Error ? result.reason.message : 'Worker revocation failed' + ).slice(0, 500); + await CodeEnvironment?.updateOne( + { _id: target._id }, + { $set: { revocationLastError: failureMessage } }, + ); + logger.error( + `[code-environments] worker revocation failed during committed user deletion: ${userId}`, + result.reason, + ); + } + return revoked; +} diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index c958c8d299..72edbed6d8 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -49,6 +49,10 @@ export const balance = () => `${BASE_URL}/api/balance`; export const userPlugins = () => `${BASE_URL}/api/user/plugins`; export const deleteUser = () => `${BASE_URL}/api/user/delete`; +export const codeEnvironments = () => `${BASE_URL}/api/code-environments`; +export const codeEnvironmentPairings = () => `${codeEnvironments()}/pairings`; +export const codeEnvironmentById = (id: string) => + `${codeEnvironments()}/${encodeURIComponent(id)}`; const messagesRoot = `${BASE_URL}/api/messages`; diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 7b0e8194e8..51a2e87c5b 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -68,6 +68,23 @@ export function deleteUser(payload?: t.TDeleteUserRequest): Promise { return request.deleteWithOptions(endpoints.deleteUser(), { data: payload }); } +export function getCodeEnvironments(): Promise { + return request.get(endpoints.codeEnvironments()); +} + +export function pairCodeEnvironment(payload: { + name: string; + controlPlaneId: string; +}): Promise { + return request.post(endpoints.codeEnvironmentPairings(), payload); +} + +export function deleteCodeEnvironment( + id: string, +): Promise<{ environment: t.TCodeEnvironmentSummary }> { + return request.delete(endpoints.codeEnvironmentById(id)); +} + export function getFavorites(): Promise { return request.get(`${endpoints.apiBaseUrl()}/api/user/settings/favorites`); } diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index d58856d066..36c7d61b62 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -95,6 +95,7 @@ export enum QueryKeys { schedule = 'schedule', parentSubagents = 'parentSubagents', subagentThread = 'subagentThread', + codeEnvironments = 'codeEnvironments', agentQueuedTurns = 'agentQueuedTurns', } @@ -151,4 +152,6 @@ export enum MutationKeys { updateSchedule = 'updateSchedule', deleteSchedule = 'deleteSchedule', runSchedule = 'runSchedule', + pairCodeEnvironment = 'pairCodeEnvironment', + deleteCodeEnvironment = 'deleteCodeEnvironment', } diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 8f28e0a987..31f92b1e68 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -534,6 +534,34 @@ export type TPublicCodeEnvironment = { pairingAvailable?: boolean; }; +export type TCodeEnvironmentSummary = { + resourceId: string; + id: string; + name: string; + type: 'managed' | 'attached'; + canDelete: boolean; +}; + +export type TCodeControlPlane = { + id: string; + name: string; +}; + +export type TCodeEnvironmentsResponse = { + environments: TCodeEnvironmentSummary[]; + controlPlanes: TCodeControlPlane[]; +}; + +export type TCodeEnvironmentPairingResponse = { + environment: TCodeEnvironmentSummary; + pairing: { + workerId: string; + code: string; + expiresAt: string; + endpoint: string; + }; +}; + export type TConfig = { order: number; type?: EModelEndpoint; diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index 139fdb71c2..017713cf24 100644 --- a/packages/data-schemas/src/methods/agent.spec.ts +++ b/packages/data-schemas/src/methods/agent.spec.ts @@ -19,7 +19,8 @@ import type { UpdateQuery, Model, } from 'mongoose'; -import type { IAgent, IAclEntry, IUser, IAccessRole } from '..'; +import type { IAgent, IAclEntry, IUser, IAccessRole, CodeEnvironmentDocument } from '..'; +import { withCodeEnvironmentReference } from './codeEnvironment'; import { createAgentMethods, type AgentMethods } from './agent'; import { tenantStorage } from '~/config/tenantContext'; import { createAclEntryMethods } from './aclEntry'; @@ -2595,6 +2596,152 @@ describe('Agent Methods', () => { expect(revertedAgent.skill_authoring_enabled).toBeUndefined(); }); + test('renews a code environment reference until the guarded write settles', async () => { + const environmentId = `environment_${uuidv4()}`; + await mongoose.models.CodeEnvironment.create({ + environmentId, + name: 'Slow writer environment', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + createdBy: new mongoose.Types.ObjectId(), + }); + let settleWrite: (() => void) | undefined; + const guardedWrite = withCodeEnvironmentReference( + mongoose, + environmentId, + async () => + await new Promise((resolve) => { + settleWrite = resolve; + }), + 10, + ); + let initial: CodeEnvironmentDocument | null = null; + for (let attempt = 0; attempt < 20; attempt++) { + initial = await mongoose.models.CodeEnvironment.findOne({ + environmentId, + }).lean(); + if (initial?.pendingAgentReferences?.length) break; + await new Promise((resolve) => setTimeout(resolve, 5)); + } + const initialExpiry = initial?.pendingAgentReferences?.[0]?.expiresAt; + + await new Promise((resolve) => setTimeout(resolve, 30)); + + const renewed = await mongoose.models.CodeEnvironment.findOne({ + environmentId, + }).lean(); + expect(initialExpiry).toBeInstanceOf(Date); + expect(renewed?.pendingAgentReferences?.[0]?.expiresAt.getTime()).toBeGreaterThan( + initialExpiry?.getTime() ?? 0, + ); + settleWrite?.(); + await guardedWrite; + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId }).lean(), + ).resolves.toMatchObject({ pendingAgentReferences: [] }); + }); + + test('rejects a guarded write when its code environment reference lease is lost', async () => { + const environmentId = `environment_${uuidv4()}`; + await mongoose.models.CodeEnvironment.create({ + environmentId, + name: 'Lost lease environment', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + createdBy: new mongoose.Types.ObjectId(), + }); + let settleWrite: (() => void) | undefined; + const compensate = jest.fn().mockResolvedValue(undefined); + const guardedWrite = withCodeEnvironmentReference( + mongoose, + environmentId, + async () => + await new Promise((resolve) => { + settleWrite = resolve; + }), + 10, + compensate, + ); + for (let attempt = 0; attempt < 20; attempt++) { + const environment = await mongoose.models.CodeEnvironment.findOne({ + environmentId, + pendingAgentReferences: { $exists: true }, + }); + if (environment != null) break; + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await mongoose.models.CodeEnvironment.updateOne( + { environmentId }, + { $set: { deletionStartedAt: new Date() } }, + ); + await new Promise((resolve) => setTimeout(resolve, 25)); + settleWrite?.(); + + await expect(guardedWrite).rejects.toMatchObject({ + name: 'CodeEnvironmentReferenceError', + }); + expect(compensate).toHaveBeenCalledWith(undefined); + }); + + test('does not overwrite an intervening Agent update when acquiring an environment reference', async () => { + const agentId = `agent_${uuidv4()}`; + const environmentId = `environment_${uuidv4()}`; + const authorId = new mongoose.Types.ObjectId(); + await createAgent({ + id: agentId, + name: 'Original agent name', + author: authorId, + model: 'test-model', + provider: 'test-provider', + }); + await mongoose.models.CodeEnvironment.create({ + environmentId, + name: 'Concurrent update environment', + type: 'attached', + baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', + createdBy: authorId, + }); + const CodeEnvironment = mongoose.models.CodeEnvironment; + const reserve = CodeEnvironment.findOneAndUpdate.bind(CodeEnvironment); + let enteredReserve!: () => void; + let releaseReserve!: () => void; + const entered = new Promise((resolve) => (enteredReserve = resolve)); + const release = new Promise((resolve) => (releaseReserve = resolve)); + const reserveSpy = jest.spyOn(CodeEnvironment, 'findOneAndUpdate').mockImplementationOnce( + (...args: Parameters) => + ({ + lean: async () => { + enteredReserve(); + await release; + return await reserve(...args).lean(); + }, + }) as ReturnType, + ); + const guardedUpdate = updateAgent( + { id: agentId }, + { name: 'Guarded update name', code_environment_id: environmentId }, + ); + await entered; + await Agent.updateOne( + { id: agentId }, + { $set: { description: 'Intervening update survived' } }, + ); + releaseReserve(); + + await expect(guardedUpdate).resolves.toBeNull(); + await expect(Agent.findOne({ id: agentId }).lean()).resolves.toMatchObject({ + name: 'Original agent name', + description: 'Intervening update survived', + }); + await expect(Agent.findOne({ id: agentId }).lean()).resolves.not.toHaveProperty( + 'code_environment_id', + ); + reserveSpy.mockRestore(); + }); + test('should prune deleted skill ids when reverting to an older version', async () => { const agentId = `agent_${uuidv4()}`; const authorId = new mongoose.Types.ObjectId(); diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index 2dcff45400..8e5e138564 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -9,6 +9,7 @@ import { import type { FilterQuery, Model, PipelineStage, ProjectionType, Types } from 'mongoose'; import type { AgentToolResources } from 'librechat-data-provider'; import type { IAgent, IAclEntry } from '~/types'; +import { withCodeEnvironmentReference } from './codeEnvironment'; import { filterExistingSkillIds } from './skill'; import logger from '~/config/winston'; @@ -548,6 +549,34 @@ export function createAgentMethods( } { const { removeAllPermissions, getActions, getSoleOwnedResourceIds, isExternalSkillId } = deps; + async function restoreAgentAfterReferenceLoss( + Agent: Model, + agentAfterWrite: IAgent | null, + originalAgent: IAgent, + lostEnvironmentId: string, + ): Promise { + if (agentAfterWrite == null) return; + const { updatedAt } = agentAfterWrite as IAgent & { updatedAt: Date }; + const restored = await Agent.replaceOne( + { + _id: agentAfterWrite._id, + code_environment_id: lostEnvironmentId, + updatedAt, + }, + originalAgent, + { timestamps: false }, + ); + if (restored.matchedCount === 0) { + /** A concurrent writer may have changed the document after the guarded + * write. Never erase that writer, but still remove the lost reference if + * it remains active. */ + await Agent.updateOne( + { _id: agentAfterWrite._id, code_environment_id: lostEnvironmentId }, + { $unset: { code_environment_id: 1 } }, + ); + } + } + /** * Create an agent with the provided data. */ @@ -585,7 +614,15 @@ export function createAgentMethods( extractMCPServerNames(agentData.tools as string[] | undefined), }; - return (await Agent.create(initialAgentData)).toObject() as IAgent; + return await withCodeEnvironmentReference( + mongoose, + typeof agentData.code_environment_id === 'string' ? agentData.code_environment_id : undefined, + async () => (await Agent.create(initialAgentData)).toObject() as IAgent, + undefined, + async (createdAgent) => { + await Agent.deleteOne({ _id: createdAgent._id }); + }, + ); } /** @@ -685,6 +722,7 @@ export function createAgentMethods( let suppressedVersionEntry = false; const currentAgent = await Agent.findOne(searchParameter); + const currentRevision = (currentAgent as (IAgent & { updatedAt: Date }) | null)?.updatedAt; if (currentAgent) { const currentObject = currentAgent.toObject() as unknown as Record; const { __v, _id, id: __id, versions, author: _author, ...versionData } = currentObject; @@ -828,11 +866,40 @@ export function createAgentMethods( } } - const updatedAgent = (await Agent.findOneAndUpdate( - searchParameter, - updateData, - mongoOptions, - ).lean()) as IAgent | null; + const directEnvironmentId = updateData.code_environment_id; + const setEnvironmentId = + typeof updateData.$set === 'object' && updateData.$set != null + ? (updateData.$set as { code_environment_id?: unknown }).code_environment_id + : undefined; + let nextEnvironmentId: string | undefined; + if (typeof directEnvironmentId === 'string') { + nextEnvironmentId = directEnvironmentId; + } else if (typeof setEnvironmentId === 'string') { + nextEnvironmentId = setEnvironmentId; + } + const updatedAgent = await withCodeEnvironmentReference( + mongoose, + nextEnvironmentId, + async () => + (await Agent.findOneAndUpdate( + currentAgent == null || nextEnvironmentId == null + ? searchParameter + : { ...searchParameter, _id: currentAgent._id, updatedAt: currentRevision }, + updateData, + mongoOptions, + ).lean()) as IAgent | null, + undefined, + async (agentAfterUpdate) => { + if (agentAfterUpdate == null || nextEnvironmentId == null) return; + if (currentAgent == null) return; + await restoreAgentAfterReferenceLoss( + Agent, + agentAfterUpdate, + currentAgent.toObject() as IAgent, + nextEnvironmentId, + ); + }, + ); /** `version` is a response-only field holding the count of `versions`. It is reported * here so a suppressed entry keeps the shape callers saw before the write was fixed. @@ -1224,6 +1291,7 @@ export function createAgentMethods( } const revertToVersion = { ...(agent.versions[versionIndex] as Record) }; + const originalRevision = (agent as unknown as IAgent & { updatedAt: Date }).updatedAt; delete revertToVersion._id; delete revertToVersion.id; delete revertToVersion.versions; @@ -1256,9 +1324,30 @@ export function createAgentMethods( Object.keys(unsetOnRestore).length > 0 ? { $set: revertToVersion, $unset: unsetOnRestore } : { $set: revertToVersion }; - const revertedAgent = await Agent.findOneAndUpdate(searchParameter, revertUpdate, { - new: true, - }).lean(); + const revertedAgent = await withCodeEnvironmentReference( + mongoose, + typeof revertToVersion.code_environment_id === 'string' + ? revertToVersion.code_environment_id + : undefined, + async () => + await Agent.findOneAndUpdate( + { ...searchParameter, _id: agent._id, updatedAt: originalRevision }, + revertUpdate, + { new: true }, + ).lean(), + undefined, + async (agentAfterRevert) => { + if (agentAfterRevert == null || typeof revertToVersion.code_environment_id !== 'string') { + return; + } + await restoreAgentAfterReferenceLoss( + Agent, + agentAfterRevert, + agent.toObject() as IAgent, + revertToVersion.code_environment_id, + ); + }, + ); if (!revertedAgent) { throw new Error('Agent not found'); } diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index d9a4899303..cb6a563293 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -1,17 +1,202 @@ import { Types } from 'mongoose'; +import { createHash } from 'crypto'; import { ResourceType } from 'librechat-data-provider'; import type { Model } from 'mongoose'; import type { CodeEnvironmentDocument } from '~/types'; import type { IAclEntry } from '~/types'; +import { getTenantId, SYSTEM_TENANT_ID } from '~/config/tenantContext'; +import logger from '~/config/winston'; + +const CODE_ENVIRONMENT_TOMBSTONES = 'code_environment_tombstones'; +const REFERENCE_LEASE_MS = 2 * 60_000; +const REFERENCE_RENEW_INTERVAL_MS = REFERENCE_LEASE_MS / 3; +const REMOVAL_LEASE_MS = 2 * 60_000; + +type CodeEnvironmentTombstone = { + _id: string; + tenantId: string; + environmentId: string; + deletedAt: Date; +}; + +function tenantScope(tenantId?: string): string { + const resolved = tenantId ?? getTenantId(); + return resolved == null || resolved === SYSTEM_TENANT_ID ? '' : resolved; +} + +function tombstoneId(environmentId: string, tenantId?: string): string { + return createHash('sha256') + .update(`librechat-code-environment-tombstone\0${tenantScope(tenantId)}\0${environmentId}`) + .digest('hex'); +} + +function tombstones(mongoose: typeof import('mongoose')) { + const db = mongoose.connection.db; + if (db == null) { + throw new Error('MongoDB is not connected'); + } + return db.collection(CODE_ENVIRONMENT_TOMBSTONES); +} + +function ownerSlotId(ownerId: Types.ObjectId, slot: number): Types.ObjectId { + const hex = createHash('sha256') + .update(`librechat-code-environment-slot\0${ownerId.toHexString()}\0${slot}`) + .digest('hex') + .slice(0, 24); + return new Types.ObjectId(hex); +} + +function agentReferenceFilter(environmentId: string, tenantId?: string) { + return { + code_environment_id: environmentId, + ...(tenantId == null ? { tenantId: { $exists: false } } : { tenantId }), + }; +} + +export type CodeEnvironmentReferenceReservation = { + environmentId: string; + reservationId: string; +}; + +export class CodeEnvironmentReferenceError extends Error { + readonly statusCode = 409; + constructor(environmentId: string) { + super(`Code environment is being removed: ${environmentId}`); + this.name = 'CodeEnvironmentReferenceError'; + } +} + +export async function reserveCodeEnvironmentReference( + mongoose: typeof import('mongoose'), + environmentId: string | undefined, +): Promise { + if (environmentId == null || environmentId.length === 0) return undefined; + const CodeEnvironment = mongoose.models.CodeEnvironment as + | Model + | undefined; + if (CodeEnvironment == null) return undefined; + const reservationId = new Types.ObjectId().toHexString(); + const now = new Date(); + const expiresAt = new Date(now.getTime() + REFERENCE_LEASE_MS); + const reserved = await CodeEnvironment.findOneAndUpdate( + { + environmentId, + registrationPendingAt: { $exists: false }, + deletionCommittedAt: { $exists: false }, + deletionStartedAt: { $exists: false }, + }, + { $push: { pendingAgentReferences: { reservationId, expiresAt } } }, + { new: true }, + ).lean(); + if (reserved != null) return { environmentId, reservationId }; + + const [persisted, tombstone] = await Promise.all([ + CodeEnvironment.exists({ environmentId }), + tombstones(mongoose).findOne({ _id: tombstoneId(environmentId) }), + ]); + if (persisted != null || tombstone != null) { + throw new CodeEnvironmentReferenceError(environmentId); + } + // Deployment-owned configuration does not have a registry row. + return undefined; +} + +export async function releaseCodeEnvironmentReference( + mongoose: typeof import('mongoose'), + reservation: CodeEnvironmentReferenceReservation | undefined, +): Promise { + if (reservation == null) return; + const CodeEnvironment = mongoose.models.CodeEnvironment as + | Model + | undefined; + if (CodeEnvironment == null) return; + await CodeEnvironment.updateOne( + { environmentId: reservation.environmentId }, + { $pull: { pendingAgentReferences: { reservationId: reservation.reservationId } } }, + ); +} + +async function renewCodeEnvironmentReference( + mongoose: typeof import('mongoose'), + reservation: CodeEnvironmentReferenceReservation, +): Promise { + const CodeEnvironment = mongoose.models.CodeEnvironment as + | Model + | undefined; + if (CodeEnvironment == null) return; + const result = await CodeEnvironment.updateOne( + { + environmentId: reservation.environmentId, + deletionStartedAt: { $exists: false }, + 'pendingAgentReferences.reservationId': reservation.reservationId, + }, + { + $set: { + 'pendingAgentReferences.$.expiresAt': new Date(Date.now() + REFERENCE_LEASE_MS), + }, + }, + ); + if (result.matchedCount !== 1) { + throw new CodeEnvironmentReferenceError(reservation.environmentId); + } +} + +export async function withCodeEnvironmentReference( + mongoose: typeof import('mongoose'), + environmentId: string | undefined, + operation: () => Promise, + renewIntervalMs: number = REFERENCE_RENEW_INTERVAL_MS, + onReferenceLoss?: (result: T) => Promise, +): Promise { + const reservation = await reserveCodeEnvironmentReference(mongoose, environmentId); + if (reservation == null) return await operation(); + + let renewal = Promise.resolve(); + let renewalError: unknown; + const timer = setInterval(() => { + renewal = renewal + .then(() => renewCodeEnvironmentReference(mongoose, reservation)) + .catch((error) => { + renewalError ??= error; + logger.error('[code-environments] agent reference lease renewal failed:', error); + }); + }, renewIntervalMs); + timer.unref(); + try { + const result = await operation(); + /** Freeze the renewal chain before awaiting it. Otherwise the interval can + * append one last renewal after this await has captured an earlier promise. */ + clearInterval(timer); + await renewal; + if (renewalError != null) { + try { + await renewCodeEnvironmentReference(mongoose, reservation); + renewalError = undefined; + } catch (error) { + await onReferenceLoss?.(result); + throw error; + } + } + return result; + } finally { + clearInterval(timer); + await renewal; + await releaseCodeEnvironmentReference(mongoose, reservation); + } +} type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' > & - Pick, 'workerId' | 'workerPrincipal'>; + Pick, 'workerId' | 'revocationTokenEnv' | 'workerPrincipal'>; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; + createCodeEnvironmentWithinOwnerLimit: ( + input: CreateCodeEnvironmentInput, + maxOwned: number, + ) => Promise; findCodeEnvironmentsByIds: ( ids: Array, ) => Promise; @@ -19,9 +204,19 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') environmentId: string, ) => Promise; listCodeEnvironmentIds: () => Promise; + findCodeEnvironmentsByCreator: ( + userId: string | Types.ObjectId, + ) => Promise; deleteCodeEnvironmentById: ( id: string | Types.ObjectId, ) => Promise; + discardCodeEnvironmentById: (id: string | Types.ObjectId) => Promise; + completeCodeEnvironmentRegistration: (id: string | Types.ObjectId) => Promise; + beginCodeEnvironmentRemoval: ( + id: string | Types.ObjectId, + ) => Promise; + cancelCodeEnvironmentRemoval: (id: string | Types.ObjectId, leaseId: string) => Promise; + commitCodeEnvironmentRemoval: (id: string | Types.ObjectId, leaseId: string) => Promise; deleteUserCodeEnvironments: (userId: string | Types.ObjectId) => Promise; } { const model = () => mongoose.models.CodeEnvironment as Model; @@ -29,7 +224,64 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') async function createCodeEnvironment( input: CreateCodeEnvironmentInput, ): Promise { - return (await model().create(input)).toObject() as CodeEnvironmentDocument; + if ((await tombstones(mongoose).findOne({ _id: tombstoneId(input.environmentId) })) != null) { + throw new Error(`Code environment id was previously retired: ${input.environmentId}`); + } + return ( + await model().create({ ...input, registrationPendingAt: new Date() }) + ).toObject() as CodeEnvironmentDocument; + } + + async function createCodeEnvironmentWithinOwnerLimit( + input: CreateCodeEnvironmentInput, + maxOwned: number, + ): Promise { + if ((await tombstones(mongoose).findOne({ _id: tombstoneId(input.environmentId) })) != null) { + throw new Error(`Code environment id was previously retired: ${input.environmentId}`); + } + const existing = await model() + .find({ createdBy: input.createdBy }) + .select('ownerSlot') + .lean>>(); + if (existing.length >= maxOwned) return null; + const reserved = new Set( + existing + .map(({ ownerSlot }) => ownerSlot) + .filter( + (slot): slot is number => + typeof slot === 'number' && Number.isSafeInteger(slot) && slot >= 0 && slot < maxOwned, + ), + ); + for (let ownerSlot = 0; reserved.size < existing.length && ownerSlot < maxOwned; ownerSlot++) { + reserved.add(ownerSlot); + } + for (let ownerSlot = 0; ownerSlot < maxOwned; ownerSlot++) { + if (reserved.has(ownerSlot)) continue; + try { + return ( + await model().create({ + ...input, + _id: ownerSlotId(input.createdBy, ownerSlot), + ownerSlot, + registrationPendingAt: new Date(), + }) + ).toObject() as CodeEnvironmentDocument; + } catch (error) { + const duplicateOwnerSlot = + typeof error === 'object' && + error != null && + 'code' in error && + error.code === 11000 && + 'keyPattern' in error && + typeof error.keyPattern === 'object' && + error.keyPattern != null && + ('_id' in error.keyPattern || 'ownerSlot' in error.keyPattern); + if (!duplicateOwnerSlot) { + throw error; + } + } + } + return null; } async function findCodeEnvironmentsByIds( @@ -52,36 +304,193 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') return await model().distinct('environmentId'); } + async function findCodeEnvironmentsByCreator( + userId: string | Types.ObjectId, + ): Promise { + const creatorId = typeof userId === 'string' ? new Types.ObjectId(userId) : userId; + return await model() + .find({ createdBy: creatorId }) + .sort({ createdAt: 1, _id: 1 }) + .lean(); + } + async function deleteCodeEnvironmentById( id: string | Types.ObjectId, ): Promise { + const environment = await model().findById(id).lean(); + if (environment == null) return null; + await tombstones(mongoose).updateOne( + { _id: tombstoneId(environment.environmentId, environment.tenantId) }, + { + $setOnInsert: { + tenantId: tenantScope(environment.tenantId), + environmentId: environment.environmentId, + deletedAt: new Date(), + }, + }, + { upsert: true }, + ); return await model().findByIdAndDelete(id).lean(); } - /** User-attached environments represent processes enrolled by their creator. - * Delete the registry records and every grant when that user is removed; - * sharing an environment does not transfer its lifecycle ownership. */ + async function discardCodeEnvironmentById(id: string | Types.ObjectId): Promise { + await model().deleteOne({ _id: id }); + } + + async function completeCodeEnvironmentRegistration(id: string | Types.ObjectId): Promise { + const result = await model().updateOne( + { + _id: id, + registrationPendingAt: { $exists: true }, + registrationLeaseId: { $exists: false }, + }, + { $unset: { registrationPendingAt: 1, registrationReconcileAfter: 1 } }, + ); + if (result.matchedCount !== 1) { + throw new Error('Code environment registration could not be committed'); + } + } + + async function beginCodeEnvironmentRemoval( + id: string | Types.ObjectId, + ): Promise { + const now = new Date(); + const leaseId = new Types.ObjectId().toHexString(); + return await model() + .findOneAndUpdate( + { + _id: id, + deletionCommittedAt: { $exists: false }, + $and: [ + { + $or: [ + { deletionStartedAt: { $exists: false } }, + { deletionLeaseExpiresAt: { $lte: now } }, + ], + }, + { + $or: [ + { pendingAgentReferences: { $exists: false } }, + { pendingAgentReferences: { $size: 0 } }, + { + pendingAgentReferences: { + $not: { $elemMatch: { expiresAt: { $gt: now } } }, + }, + }, + ], + }, + ], + }, + { + $set: { + deletionStartedAt: now, + deletionLeaseId: leaseId, + deletionLeaseExpiresAt: new Date(now.getTime() + REMOVAL_LEASE_MS), + }, + $pull: { pendingAgentReferences: { expiresAt: { $lte: now } } }, + }, + { new: true }, + ) + .lean(); + } + + async function cancelCodeEnvironmentRemoval( + id: string | Types.ObjectId, + leaseId: string, + ): Promise { + await model().updateOne( + { _id: id, deletionLeaseId: leaseId, deletionCommittedAt: { $exists: false } }, + { + $unset: { + deletionStartedAt: 1, + deletionLeaseId: 1, + deletionLeaseExpiresAt: 1, + }, + }, + ); + } + + async function commitCodeEnvironmentRemoval( + id: string | Types.ObjectId, + leaseId: string, + ): Promise { + const result = await model().updateOne( + { _id: id, deletionLeaseId: leaseId }, + { + $set: { deletionCommittedAt: new Date() }, + $unset: { deletionLeaseExpiresAt: 1 }, + }, + ); + if (result.matchedCount !== 1) { + throw new Error('Code environment removal lease was lost'); + } + } + + /** Delete unreferenced creator-owned records and grants after user removal. + * Failed worker revocations and environments selected by surviving agents + * remain durable for reconciliation instead of becoming dangling refs. */ async function deleteUserCodeEnvironments(userId: string | Types.ObjectId): Promise { const creatorId = typeof userId === 'string' ? new Types.ObjectId(userId) : userId; - const environments = await model().find({ createdBy: creatorId }).select('_id').lean(); + const environments = await model() + .find({ + createdBy: creatorId, + revocationPendingAt: { $exists: false }, + deletionCommittedAt: { $exists: false }, + }) + .select('_id') + .lean>>(); if (environments.length === 0) return 0; - const resourceIds = environments.map((environment) => environment._id); + const Agent = mongoose.models.Agent; const AclEntry = mongoose.models.AclEntry as Model; - await AclEntry.deleteMany({ - resourceType: ResourceType.CODE_ENVIRONMENT, - resourceId: { $in: resourceIds }, - }); - const result = await model().deleteMany({ _id: { $in: resourceIds } }); - return result.deletedCount ?? 0; + let deleted = 0; + for (const candidate of environments) { + const claimed = await beginCodeEnvironmentRemoval(candidate._id); + if (claimed == null || claimed.deletionLeaseId == null) continue; + const referenced = + Agent != null && + (await Agent.exists(agentReferenceFilter(claimed.environmentId, claimed.tenantId))) != null; + if (referenced) { + await cancelCodeEnvironmentRemoval(candidate._id, claimed.deletionLeaseId); + continue; + } + await tombstones(mongoose).updateOne( + { _id: tombstoneId(claimed.environmentId, claimed.tenantId) }, + { + $setOnInsert: { + tenantId: tenantScope(claimed.tenantId), + environmentId: claimed.environmentId, + deletedAt: new Date(), + }, + }, + { upsert: true }, + ); + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: candidate._id, + }); + const result = await model().deleteOne({ + _id: candidate._id, + deletionLeaseId: claimed.deletionLeaseId, + }); + deleted += result.deletedCount ?? 0; + } + return deleted; } return { createCodeEnvironment, + createCodeEnvironmentWithinOwnerLimit, findCodeEnvironmentsByIds, findCodeEnvironmentByEnvironmentId, listCodeEnvironmentIds, + findCodeEnvironmentsByCreator, deleteCodeEnvironmentById, + discardCodeEnvironmentById, + completeCodeEnvironmentRegistration, + beginCodeEnvironmentRemoval, + cancelCodeEnvironmentRemoval, + commitCodeEnvironmentRemoval, deleteUserCodeEnvironments, }; } diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index fb297e5d35..ea320b3b86 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -45,9 +45,51 @@ const codeEnvironmentSchema: Schema = new Schema; + deletionStartedAt?: Date; + deletionLeaseId?: string; + deletionLeaseExpiresAt?: Date; + deletionCommittedAt?: Date; + registrationPendingAt?: Date; + registrationLeaseId?: string; + registrationLeaseExpiresAt?: Date; + registrationReconcileAfter?: Date; + revocationPendingAt?: Date; + revocationAttempts?: number; + revocationLastError?: string; + revocationReconcileAfter?: Date; + revocationLeaseId?: string; + revocationLeaseExpiresAt?: Date; workerId?: string; + revocationTokenEnv?: string; workerPrincipal?: { type: 'deployment' | 'tenant' | 'user' | 'role' | 'group'; id: string;