From e0861d633f7d0236edc4d6bfcba77f50a5b4fa62 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 13:06:56 -0400 Subject: [PATCH 01/29] feat: add self-service code worker pairing --- api/app/clients/tools/util/handleTools.js | 7 +- api/server/routes/code-environments.js | 1 + api/server/routes/code-environments.test.js | 11 ++ api/server/services/ToolService.js | 10 +- e2e/setup/fake-model.js | 17 ++- librechat.example.yaml | 14 +- packages/api/src/admin/code.spec.ts | 6 +- packages/api/src/admin/code.ts | 6 +- packages/api/src/agents/execution.spec.ts | 38 +++++- packages/api/src/agents/execution.ts | 22 +++- packages/api/src/auth/codeapi.ts | 13 +- packages/api/src/code/bridge.ts | 92 +++++++++++++ .../src/code/environments.integration.spec.ts | 2 + packages/api/src/code/environments.ts | 2 + packages/api/src/code/http.spec.ts | 94 ++++++++++++++ packages/api/src/code/http.ts | 122 +++++++++++++++++- packages/api/src/code/index.ts | 1 + .../src/endpoints/config/endpoints.spec.ts | 5 + .../specs/config-schemas.spec.ts | 43 ++++++ packages/data-provider/src/config.ts | 20 ++- 20 files changed, 506 insertions(+), 20 deletions(-) create mode 100644 packages/api/src/code/bridge.ts diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index a982d698e1d..13fd231c61d 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -21,6 +21,7 @@ const { ASK_USER_QUESTION_TOOL_NAME, resolveWebSearchSSRFAgents, buildWebSearchDynamicContext, + codeExecutionAuthHeaders, resolveCodeExecutionContext, } = require('@librechat/api'); const { @@ -379,7 +380,11 @@ const loadTools = async ({ return createCodeExecutionTool({ user_id: user, files, - authHeaders: () => getCodeApiAuthHeaders(options.req), + authHeaders: () => + codeExecutionAuthHeaders( + () => getCodeApiAuthHeaders(options.req), + codeExecutionContext, + ), ...codeExecutionContext, }); }; diff --git a/api/server/routes/code-environments.js b/api/server/routes/code-environments.js index 316ffe1d09a..f0f84edaeee 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -20,6 +20,7 @@ const requireCodeEnvironmentManage = requireCapability(SystemCapabilities.MANAGE router.use(requireJwtAuth); router.get('/', (req, res, next) => getHandlers().list(req, res, next)); +router.post('/pairings', (req, res, next) => getHandlers().pair(req, res, next)); router.post('/', requireCodeEnvironmentManage, (req, res, next) => getHandlers().register(req, res, next), ); diff --git a/api/server/routes/code-environments.test.js b/api/server/routes/code-environments.test.js index 1be0bbca3a4..ed0d1f28fba 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -16,6 +16,7 @@ const mockGetCodeEnvironmentRegistry = jest.fn(() => mockRegistry); 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' } })), }; jest.mock('@librechat/data-schemas', () => ({ @@ -74,4 +75,14 @@ describe('code environment routes', () => { expect(middlewareCalls).toEqual(['jwt', 'manage:code_environments']); expect(mockHandlers.register).toHaveBeenCalledTimes(1); }); + + it('allows an authenticated user to pair through an opted-in control plane', async () => { + await request(createApp()) + .post('/api/code-environments/pairings') + .send({ name: 'Personal VM', controlPlaneId: 'shared-code-api' }) + .expect(201); + + expect(middlewareCalls).toEqual(['jwt']); + expect(mockHandlers.pair).toHaveBeenCalledTimes(1); + }); }); diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index f5655088a85..8f1956ed60b 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -44,6 +44,7 @@ const { isNormalizationSensitiveName, AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE, isFatalAgentInitializationError, + codeExecutionAuthHeaders, resolveCodeExecutionContext, resolveCallerCapabilityProjectionSnapshot, getTransactionsConfig, @@ -1694,7 +1695,8 @@ async function loadAgentTools({ deferredToolsEnabled, programmaticToolsEnabled, codeExecutionEnabled, - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), codeExecutionContext, }); @@ -2079,7 +2081,8 @@ async function loadToolsForExecution({ */ for (const name of ptcToolNames) { const ptcTool = createBashProgrammaticToolCallingTool({ - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), baseUrl: codeExecutionContext.baseUrl, executionProfile: codeExecutionContext.executionProfile, runtimeSessionHint: codeExecutionContext.runtimeSessionHint, @@ -2105,7 +2108,8 @@ async function loadToolsForExecution({ if (isBashTool) { try { const bashTool = createBashExecutionTool({ - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), ...codeExecutionContext, }); allLoadedTools.push(bashTool); diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index fcaaa27a759..a804535d2e8 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -1665,19 +1665,20 @@ function statefulCodeResponses(operation, toolNames) { return { responses: [`E2E stateful code failed: unsupported operation ${operation}`] }; } + const toolCallId = `call_e2e_stateful_code_${operation}`; return { responses: ['', ''], toolCalls: [ { - id: `call_e2e_stateful_code_${operation}`, + id: toolCallId, name: BASH_TOOL_NAME, args: { command }, type: 'tool_call', }, ], resolveOnStream: (streamMessages) => { - const toolText = findLastToolMessageText(streamMessages, STATEFUL_CODE_VALUE); - if (!toolText) { + const toolMessage = findLastToolMessage(streamMessages, toolCallId); + if (!getContentText(toolMessage?.content).includes(STATEFUL_CODE_VALUE)) { return null; } return { responses: [`E2E stateful code ${operation} observed ${STATEFUL_CODE_VALUE}`] }; @@ -1835,6 +1836,16 @@ function findToolMessage(messages, toolCallId) { ); } +function findLastToolMessage(messages, toolCallId) { + for (let index = (messages?.length ?? 0) - 1; index >= 0; index--) { + const message = messages[index]; + if (messageType(message) === 'tool' && message?.tool_call_id === toolCallId) { + return message; + } + } + return undefined; +} + function deferredHitlCallId(label, phase) { return `call_e2e_deferred_hitl_${phase}_${label}`; } diff --git a/librechat.example.yaml b/librechat.example.yaml index 4bc675a8e71..55e759ea448 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -578,12 +578,22 @@ endpoints: # name: Engineering VM # type: attached # baseURL: https://code-bridge.example.com/v1 - # # Enables the admin pairing API without placing the administrator - # # secret in YAML or returning its environment-variable name to users. + # # Pins an administrator-managed environment to one worker. # owner: deployment # pairing: # workerId: engineering-vm # tokenEnv: CODE_BRIDGE_ADMIN_TOKEN + # - id: personal-workers + # name: Personal Code Workers + # type: attached + # baseURL: https://code-bridge.example.com/v1 + # # Lets authenticated users create principal-bound environments and + # # receive a one-time code for an outbound @librechat/code worker. + # # The administrator secret stays in the named environment variable. + # owner: deployment + # pairing: + # allowPrincipalWorkers: true + # tokenEnv: CODE_BRIDGE_ADMIN_TOKEN # # (optional) Trusted origin for internal event delivery. By default, LibreChat # # uses its own bound listener; override only for an internal TLS/front-door route. # eventDriven: diff --git a/packages/api/src/admin/code.spec.ts b/packages/api/src/admin/code.spec.ts index 38b22c7deb6..56508cfa518 100644 --- a/packages/api/src/admin/code.spec.ts +++ b/packages/api/src/admin/code.spec.ts @@ -106,7 +106,11 @@ describe('createAdminCodeEnvironmentHandlers', () => { throw new Error('Expected the test code environment'); } overriddenEnvironment.baseURL = 'https://attacker.example.com/v1'; - overriddenEnvironment.pairing = { workerId: 'vm-1', tokenEnv: 'DATABASE_URL' }; + overriddenEnvironment.pairing = { + workerId: 'vm-1', + allowPrincipalWorkers: false, + tokenEnv: 'DATABASE_URL', + }; const getAppConfig = jest.fn(async (options: { baseOnly?: boolean }) => options.baseOnly === true ? config() : writableOverride, ); diff --git a/packages/api/src/admin/code.ts b/packages/api/src/admin/code.ts index f84c45a177c..2b05c7d4993 100644 --- a/packages/api/src/admin/code.ts +++ b/packages/api/src/admin/code.ts @@ -55,7 +55,11 @@ function pairingConfig(environment: ConfiguredCodeEnvironment): if (environment.type !== 'attached' || environment.owner !== 'deployment') { return undefined; } - return environment.pairing; + const pairing = environment.pairing; + if (pairing?.workerId == null) { + return undefined; + } + return { workerId: pairing.workerId, tokenEnv: pairing.tokenEnv }; } function bridgeUrl(environment: ConfiguredCodeEnvironment, path: string): string { diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts index 859e231d33e..8dc74c6f13a 100644 --- a/packages/api/src/agents/execution.spec.ts +++ b/packages/api/src/agents/execution.spec.ts @@ -1,4 +1,8 @@ -import { resolveCodeExecutionContext } from './execution'; +import { + codeExecutionAuthHeaders, + codeExecutionHeaders, + resolveCodeExecutionContext, +} from './execution'; jest.mock('@librechat/agents', () => ({ Constants: { EXECUTE_CODE: 'execute_code' }, @@ -113,6 +117,7 @@ describe('resolveCodeExecutionContext', () => { name: 'My VM', type: 'attached', baseURL: 'https://bridge.example/v1/', + workerId: 'opaque-worker-id', owner: 'deployment', }, ], @@ -126,12 +131,43 @@ describe('resolveCodeExecutionContext', () => { environmentId: 'my-vm', environmentType: 'attached', executionProfile: 'stateful', + bridgeWorkerId: 'opaque-worker-id', }), ); expect(context.runtimeSessionHint).toMatch(/^v3:[a-f0-9]{12}:agent-user:/); expect(context.executionRouteKey).toMatch(/^stateful:[a-f0-9]{32}$/); }); + it('adds the server-selected worker to execution auth without replacing authentication', async () => { + const context = resolveCodeExecutionContext({ + statefulSessions: true, + environmentId: 'my-vm', + environments: [ + { + id: 'my-vm', + name: 'My VM', + type: 'attached', + baseURL: 'https://bridge.example/v1', + owner: 'principal', + workerId: 'opaque-worker-id', + }, + ], + userId: 'user-1', + }); + + expect(codeExecutionHeaders(context)).toEqual({ + 'X-CodeAPI-Expected-Profile': 'stateful', + 'X-LibreChat-Code-Worker-ID': 'opaque-worker-id', + }); + await expect( + codeExecutionAuthHeaders(async () => ({ Authorization: 'Bearer user-token' }), context), + ).resolves.toEqual({ + Authorization: 'Bearer user-token', + 'X-CodeAPI-Expected-Profile': 'stateful', + 'X-LibreChat-Code-Worker-ID': 'opaque-worker-id', + }); + }); + it('namespaces configured deployments independently of the shared wire profile', () => { const environment = (id: string, baseURL: string) => ({ id, diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index 9bab6cb5510..dbac27b810b 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -3,6 +3,7 @@ import { Constants, getCodeBaseURL } from '@librechat/agents'; import type { StatefulCodeEnvironment, TAgentsEndpoint } from 'librechat-data-provider'; export const CODE_API_EXPECTED_PROFILE_HEADER = 'X-CodeAPI-Expected-Profile'; +export const CODE_API_BRIDGE_WORKER_HEADER = 'X-LibreChat-Code-Worker-ID'; export type CodeExecutionProfile = 'default' | 'stateful'; export type CodeEnvironmentConfig = NonNullable< @@ -21,6 +22,7 @@ export interface CodeExecutionContext { statefulSessions: boolean; environmentId?: string; environmentType?: CodeEnvironmentConfig['type']; + bridgeWorkerId?: string; } export function createCodeExecutionRouteKey( @@ -152,11 +154,27 @@ export function resolveCodeExecutionContext(params: { statefulSessions: true, environmentId: configuredEnvironment?.id, environmentType: configuredEnvironment?.type, + bridgeWorkerId: configuredEnvironment?.workerId, }; } export function codeExecutionHeaders( - context: Pick, + context: Pick, ): Record { - return { [CODE_API_EXPECTED_PROFILE_HEADER]: context.executionProfile }; + return { + [CODE_API_EXPECTED_PROFILE_HEADER]: context.executionProfile, + ...(context.bridgeWorkerId != null + ? { [CODE_API_BRIDGE_WORKER_HEADER]: context.bridgeWorkerId } + : {}), + }; +} + +export async function codeExecutionAuthHeaders( + authHeaders: () => Promise> | Record, + context: Pick, +): Promise> { + return { + ...(await authHeaders()), + ...codeExecutionHeaders(context), + }; } diff --git a/packages/api/src/auth/codeapi.ts b/packages/api/src/auth/codeapi.ts index 24bf26bf391..a4a25edde77 100644 --- a/packages/api/src/auth/codeapi.ts +++ b/packages/api/src/auth/codeapi.ts @@ -218,6 +218,14 @@ function resolveTenantId(user: CodeApiUserContext): string | undefined { return resolveSingleTenantId(); } +export function getCodeApiTenantId(req: ServerRequest): string { + const tenantId = resolveTenantId(resolveUser(req)); + if (!tenantId) { + throw new Error('Code API JWT auth requires tenant context'); + } + return tenantId; +} + function isManagedCodeApiJwtMode(): boolean { const provider = process.env.CODEAPI_AUTH_PROVIDER; return provider === 'librechat-jwt' || provider === 'both'; @@ -258,10 +266,7 @@ function canonicalContextHash(input: { function buildClaims(req: ServerRequest, config: SigningConfig, now: number): CodeApiClaims { const user = resolveUser(req); const userId = resolveUserId(user); - const tenantId = resolveTenantId(user); - if (!tenantId) { - throw new Error('Code API JWT auth requires tenant context'); - } + const tenantId = getCodeApiTenantId(req); const role = user.role ?? 'USER'; const principalSource = resolvePrincipalSource(req); diff --git a/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts new file mode 100644 index 00000000000..77b88d48fa8 --- /dev/null +++ b/packages/api/src/code/bridge.ts @@ -0,0 +1,92 @@ +const CODE_BRIDGE_REQUEST_TIMEOUT_MS = 10_000; + +export type CodeBridgePrincipalType = 'deployment' | 'tenant' | 'user' | 'role' | 'group'; + +export type CodeBridgeWorkerBinding = { + tenantId: string; + principal: { + type: CodeBridgePrincipalType; + id: string; + }; +}; + +export type CodeBridgePairing = { + protocolVersion: 1; + workerId: string; + code: string; + expiresAt: string; +}; + +export type CodeBridgeFetch = ( + input: Parameters[0], + init?: Parameters[1], +) => ReturnType; + +export class CodeBridgePairingError extends Error { + constructor( + public readonly reason: 'rejected' | 'invalid' | 'timeout' | 'failed', + public readonly upstreamStatus?: number, + ) { + super(`Code bridge pairing ${reason}`); + this.name = 'CodeBridgePairingError'; + } +} + +function validPairing(value: unknown, workerId: string): value is CodeBridgePairing { + if (typeof value !== 'object' || value == null) return false; + const pairing = value as Partial; + const expiresAt = typeof pairing.expiresAt === 'string' ? Date.parse(pairing.expiresAt) : NaN; + return ( + pairing.protocolVersion === 1 && + pairing.workerId === workerId && + typeof pairing.code === 'string' && + /^[A-Za-z0-9_-]{32}$/.test(pairing.code) && + Number.isFinite(expiresAt) && + expiresAt > Date.now() + ); +} + +export async function createCodeBridgePairing({ + baseURL, + token, + workerId, + binding, + fetchImpl = fetch, +}: { + baseURL: string; + token: string; + workerId: string; + binding?: CodeBridgeWorkerBinding; + fetchImpl?: CodeBridgeFetch; +}): Promise { + try { + const response = await fetchImpl(`${baseURL.trim().replace(/\/+$/, '')}/bridge/pairings`, { + method: 'POST', + headers: { + Authorization: `Bearer ${token}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId, ...(binding != null ? { binding } : {}) }), + redirect: 'error', + signal: AbortSignal.timeout(CODE_BRIDGE_REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + throw new CodeBridgePairingError('rejected', response.status); + } + const payload = (await response.json()) as unknown; + if (!validPairing(payload, workerId)) { + throw new CodeBridgePairingError('invalid'); + } + return payload; + } catch (error) { + if (error instanceof CodeBridgePairingError) throw error; + if (error instanceof Error && error.name === 'TimeoutError') { + throw new CodeBridgePairingError('timeout'); + } + throw new CodeBridgePairingError('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 ae52b21bbad..eb71da86eb4 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -50,6 +50,7 @@ describe('code environment registry', () => { type: 'attached', baseURL: 'https://code.example.com', controlPlaneId: 'shared-code-api', + workerId: 'danny-worker', }, }); @@ -81,6 +82,7 @@ describe('code environment registry', () => { baseURL: 'https://code.example.com', controlPlaneId: 'shared-code-api', owner: 'principal', + workerId: 'danny-worker', }, ]); }); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index aced59bb9df..acf29d3ad96 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -41,6 +41,7 @@ export type AccessibleCodeEnvironmentConfiguration = { baseURL: string; controlPlaneId: string; owner: 'principal'; + workerId?: string; }; type CachedAccessibleCodeEnvironmentConfiguration = AccessibleCodeEnvironmentConfiguration & { @@ -266,6 +267,7 @@ export function createCodeEnvironmentRegistry( baseURL: environment.baseURL, controlPlaneId: environment.controlPlaneId, owner: 'principal', + workerId: environment.workerId, })); }; if (configurationCache == null) return (await load()).map(toPublicConfiguration); diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index b9b4df55e51..7ff2e89e4c0 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -37,6 +37,100 @@ describe('code environment HTTP handlers', () => { expect(register).not.toHaveBeenCalled(); }); + test('pairs a generated worker to the authenticated user and persists its private route', async () => { + const register = jest.fn().mockResolvedValue({ + resourceId: '68b2f0c498f24c1e78fa0111', + id: 'code-generated', + name: 'Personal VM', + type: 'attached', + }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: jest.fn().mockResolvedValue({ + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + }); + const appConfig = { + 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; + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn().mockResolvedValue(appConfig), + registry: { register, listAccessible: jest.fn() }, + createEnvironmentId: () => 'code-generated', + readSecret: jest.fn(() => 'administrator-token'), + resolveTenantId: jest.fn(() => 'tenant-1'), + fetchImpl, + }); + const req = { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { + name: 'Personal VM', + controlPlaneId: 'shared-code-api', + workerId: 'attacker-worker', + baseURL: 'https://attacker.example', + }, + }; + const res = response(); + + await handlers.pair(req as never, res as never); + + expect(res.statusCode).toBe(201); + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.librechat.example/v1/bridge/pairings', + expect.objectContaining({ + method: 'POST', + headers: expect.objectContaining({ Authorization: 'Bearer administrator-token' }), + body: JSON.stringify({ + workerId: 'code-generated', + binding: { + tenantId: 'tenant-1', + principal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, + }, + }), + }), + ); + expect(register).toHaveBeenCalledWith({ + actor: { + userId: '68b2f0c498f24c1e78fa0001', + role: 'USER', + idOnTheSource: null, + }, + environment: { + id: 'code-generated', + name: 'Personal VM', + type: 'attached', + baseURL: 'https://code.librechat.example/v1', + workerId: 'code-generated', + }, + }); + expect(res.body).toEqual({ + environment: expect.objectContaining({ id: 'code-generated' }), + pairing: expect.objectContaining({ workerId: 'code-generated', code: 'a'.repeat(32) }), + }); + }); + test('registers against an operator-configured control plane and ignores client URLs', async () => { const register = jest.fn().mockResolvedValue({ resourceId: '68b2f0c498f24c1e78fa0111', diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 9cd73544889..c6016799290 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -9,7 +9,10 @@ import type { } from './environments'; import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; +import type { CodeBridgeFetch } from './bridge'; +import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; import { CodeEnvironmentValidationError } from './environments'; +import { getCodeApiTenantId } from '~/auth/codeapi'; type Registry = { register: (params: { @@ -30,6 +33,9 @@ export interface CodeEnvironmentHttpDeps { getAppConfig: (options: GetAppConfigOptions) => Promise; registry: Registry; createEnvironmentId?: () => string; + readSecret?: (name: string) => string | undefined; + resolveTenantId?: (req: ServerRequest) => string; + fetchImpl?: CodeBridgeFetch; } function actor(req: ServerRequest): CodeEnvironmentPrincipalContext | null { @@ -54,11 +60,48 @@ function configuredControlPlane( ); } +function configuredPrincipalControlPlane( + appConfig: AppConfig, + controlPlaneId: string, +): ConfiguredCodeEnvironment | undefined { + return appConfig.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments?.find( + (environment) => + environment.id === controlPlaneId && + environment.type === 'attached' && + environment.owner === 'deployment' && + environment.pairing?.allowPrincipalWorkers === true, + ); +} + +function pairingErrorResponse(error: unknown, res: Response): Response { + if (!(error instanceof CodeBridgePairingError)) { + return res.status(502).json({ error: 'Code API pairing request failed' }); + } + if (error.reason === 'timeout') { + return res.status(504).json({ error: 'Code API pairing request timed out' }); + } + if (error.reason === 'rejected') { + return res.status(502).json({ + error: 'Code API rejected the pairing request', + upstreamStatus: error.upstreamStatus, + }); + } + return res.status(502).json({ + error: + error.reason === 'invalid' + ? 'Code API returned an invalid pairing response' + : 'Code API pairing request failed', + }); +} + export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps): { list: (req: ServerRequest, res: Response) => Promise; register: (req: ServerRequest, res: Response) => Promise; + pair: (req: ServerRequest, res: Response) => Promise; } { const createEnvironmentId = deps.createEnvironmentId ?? (() => `code-${nanoid(20)}`); + const readSecret = deps.readSecret ?? readCodeBridgeSecret; + const resolveTenantId = deps.resolveTenantId ?? getCodeApiTenantId; async function list(req: ServerRequest, res: Response): Promise { const principal = actor(req); @@ -125,5 +168,82 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } } - return { list, register }; + async function pair(req: ServerRequest, res: Response): Promise { + const principal = actor(req); + if (principal == null) { + return res.status(401).json({ error: 'Authentication required' }); + } + const body = req.body as unknown as Record; + const name = typeof body.name === 'string' ? body.name.trim() : ''; + const controlPlaneId = + typeof body.controlPlaneId === 'string' ? body.controlPlaneId.trim() : ''; + if (!name || !controlPlaneId) { + return res.status(400).json({ error: 'name and controlPlaneId are required' }); + } + + const appConfig = await deps.getAppConfig({ baseOnly: true }); + const controlPlane = configuredPrincipalControlPlane(appConfig, controlPlaneId); + if (controlPlane == null) { + return res.status(404).json({ error: 'Principal code control plane was not found' }); + } + const tokenEnv = controlPlane.pairing?.tokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token) { + return res.status(503).json({ error: 'Code environment pairing is not configured' }); + } + + const workerId = createEnvironmentId(); + let pairing; + try { + pairing = await createCodeBridgePairing({ + baseURL: controlPlane.baseURL, + token, + workerId, + binding: { + tenantId: resolveTenantId(req), + principal: { type: 'user', id: principal.userId.toString() }, + }, + fetchImpl: deps.fetchImpl, + }); + } catch (error) { + return pairingErrorResponse(error, res); + } + + try { + const environment = await deps.registry.register({ + actor: principal, + environment: { + id: workerId, + name, + type: 'attached', + baseURL: controlPlane.baseURL, + workerId, + }, + }); + return res.status(201).json({ + environment, + pairing: { + workerId: pairing.workerId, + code: pairing.code, + expiresAt: pairing.expiresAt, + }, + }); + } catch (error) { + const duplicate = + typeof error === 'object' && + error != null && + 'code' in error && + (error as { code?: number }).code === 11000; + if (duplicate) { + return res.status(409).json({ error: 'Code environment already exists' }); + } + if (error instanceof CodeEnvironmentValidationError) { + return res.status(400).json({ error: error.message }); + } + logger.error('[codeEnvironments] pairing registration failed:', error); + return res.status(500).json({ error: 'Code environment registration failed' }); + } + } + + return { list, register, pair }; } diff --git a/packages/api/src/code/index.ts b/packages/api/src/code/index.ts index 494b7431a33..0e405e9f190 100644 --- a/packages/api/src/code/index.ts +++ b/packages/api/src/code/index.ts @@ -1,3 +1,4 @@ export * from './environments'; export * from './http'; export * from './config'; +export * from './bridge'; diff --git a/packages/api/src/endpoints/config/endpoints.spec.ts b/packages/api/src/endpoints/config/endpoints.spec.ts index d0671f757a1..bff43536885 100644 --- a/packages/api/src/endpoints/config/endpoints.spec.ts +++ b/packages/api/src/endpoints/config/endpoints.spec.ts @@ -194,6 +194,11 @@ describe('createEndpointsConfigService', () => { name: 'Attached VM', type: 'attached', baseURL: 'https://internal-code.example.com/v1', + workerId: 'private-worker-route', + pairing: { + workerId: 'private-worker-route', + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, default: true, }, ], diff --git a/packages/data-provider/specs/config-schemas.spec.ts b/packages/data-provider/specs/config-schemas.spec.ts index 4476ac2ce66..057438f72c6 100644 --- a/packages/data-provider/specs/config-schemas.spec.ts +++ b/packages/data-provider/specs/config-schemas.spec.ts @@ -576,6 +576,29 @@ describe('agentsEndpointSchema', () => { expect(result.success).toBe(true); }); + it('accepts a principal-worker control plane without a singleton worker', () => { + const result = agentsEndpointSchema.safeParse({ + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'personal-code-control-plane', + name: 'Personal Code', + type: 'attached', + baseURL: 'https://bridge.example.com/v1', + default: true, + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }); + + expect(result.success).toBe(true); + }); + it('rejects pairing metadata on a principal-owned environment', () => { const result = agentsEndpointSchema.safeParse({ statefulCodeSessions: { @@ -599,6 +622,26 @@ describe('agentsEndpointSchema', () => { expect(result.success).toBe(false); }); + it('rejects pairing configuration without a singleton or principal workers', () => { + const result = agentsEndpointSchema.safeParse({ + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'invalid-control-plane', + name: 'Invalid', + type: 'attached', + baseURL: 'https://bridge.example.com/v1', + default: true, + pairing: { tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN' }, + }, + ], + }, + }); + + expect(result.success).toBe(false); + }); + it('rejects pairing configuration for a managed environment', () => { const result = agentsEndpointSchema.safeParse({ statefulCodeSessions: { diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 42e39f10e58..5cf4c5a9bc8 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1084,6 +1084,11 @@ export const agentsEndpointSchema = baseEndpointSchema type: z.enum(['managed', 'attached']), baseURL: codeEnvironmentBaseURLSchema, default: z.boolean().optional(), + /** Server-only outbound worker route. Removed from public config. */ + workerId: z + .string() + .regex(/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/) + .optional(), /** Distinguishes operator policy from a principal-authorized * environment merged into request-scoped server config. */ owner: z.enum(['deployment', 'principal']).optional().default('deployment'), @@ -1091,9 +1096,22 @@ export const agentsEndpointSchema = baseEndpointSchema * environment variable and never contains the token itself. */ pairing: z .object({ - workerId: z.string().regex(/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/), + workerId: z + .string() + .regex(/^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/) + .optional(), + allowPrincipalWorkers: z.boolean().optional().default(false), tokenEnv: z.string().regex(/^[A-Za-z_][A-Za-z0-9_]*$/), }) + .superRefine((pairing, pairingContext) => { + if (pairing.workerId != null || pairing.allowPrincipalWorkers === true) { + return; + } + pairingContext.addIssue({ + code: z.ZodIssueCode.custom, + message: 'Pairing requires a workerId or principal workers', + }); + }) .optional(), }), ) From 4b5295a8d7c68010043914b4c3791d4e72e16a49 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 13:49:31 -0400 Subject: [PATCH 02/29] fix: enforce code worker principal boundaries --- packages/api/src/agents/execution.spec.ts | 27 ++++++ packages/api/src/agents/execution.ts | 2 +- .../src/code/environments.integration.spec.ts | 33 +++++++ packages/api/src/code/environments.ts | 56 +++++++++--- packages/api/src/code/http.spec.ts | 89 +++++++++++++++---- packages/api/src/code/http.ts | 20 ++++- .../src/methods/codeEnvironment.ts | 2 +- .../src/schema/codeEnvironment.ts | 18 ++++ .../data-schemas/src/types/codeEnvironment.ts | 4 + 9 files changed, 219 insertions(+), 32 deletions(-) diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts index 8dc74c6f13a..c2c96c1f3cc 100644 --- a/packages/api/src/agents/execution.spec.ts +++ b/packages/api/src/agents/execution.spec.ts @@ -138,6 +138,33 @@ describe('resolveCodeExecutionContext', () => { expect(context.executionRouteKey).toMatch(/^stateful:[a-f0-9]{32}$/); }); + it('routes a deployment worker declared in pairing metadata', () => { + const context = resolveCodeExecutionContext({ + statefulSessions: true, + environmentId: 'deployment-vm', + environments: [ + { + id: 'deployment-vm', + name: 'Deployment VM', + type: 'attached', + baseURL: 'https://bridge.example/v1', + owner: 'deployment', + pairing: { + workerId: 'deployment-worker', + allowPrincipalWorkers: false, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ], + userId: 'user-1', + }); + + expect(context.bridgeWorkerId).toBe('deployment-worker'); + expect(codeExecutionHeaders(context)).toMatchObject({ + 'X-LibreChat-Code-Worker-ID': 'deployment-worker', + }); + }); + it('adds the server-selected worker to execution auth without replacing authentication', async () => { const context = resolveCodeExecutionContext({ statefulSessions: true, diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index dbac27b810b..953425922bf 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -154,7 +154,7 @@ export function resolveCodeExecutionContext(params: { statefulSessions: true, environmentId: configuredEnvironment?.id, environmentType: configuredEnvironment?.type, - bridgeWorkerId: configuredEnvironment?.workerId, + bridgeWorkerId: configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId, }; } diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index eb71da86eb4..9031b44eb1e 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -145,6 +145,39 @@ describe('code environment registry', () => { ).resolves.toEqual([roleEnvironment, groupEnvironment]); }); + test('keeps a user-bound worker private even if its ACL is granted to a role', async () => { + const registry = createCodeEnvironmentRegistry(mongoose); + const access = new AccessControlService(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: 'owner-worker', + name: 'Owner worker', + type: 'attached', + baseURL: 'https://code.example.com', + workerId: 'owner-worker', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + await access.grantPermission({ + principalType: PrincipalType.ROLE, + principalId: 'CODE_USER', + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment.resourceId, + accessRoleId: AccessRoleIds.CODE_ENVIRONMENT_VIEWER, + grantedBy: ownerId, + }); + + await expect( + registry.listAccessible({ userId: teammateId, role: 'CODE_USER', idOnTheSource: null }), + ).resolves.toEqual([]); + await expect( + registry.listAccessible({ userId: ownerId, role: 'USER', idOnTheSource: null }), + ).resolves.toEqual([environment]); + }); + test('removes creator-owned environment records and grants when the user is deleted', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index acf29d3ad96..0625b93ce23 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -32,6 +32,10 @@ export type CodeEnvironmentRegistration = { baseURL: string; controlPlaneId: string; workerId?: string; + workerPrincipal?: { + type: 'deployment' | 'tenant' | 'user' | 'role' | 'group'; + id: string; + }; }; export type AccessibleCodeEnvironmentConfiguration = { @@ -50,6 +54,7 @@ type CachedAccessibleCodeEnvironmentConfiguration = AccessibleCodeEnvironmentCon 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)?$/; const CONFIGURATION_CACHE_TTL_MS = 5_000; const CONFIGURATION_CACHE_REVISION_PREFIX = 'revision'; const CONFIGURATION_CACHE_REGISTERED_PREFIX = 'registered'; @@ -65,6 +70,14 @@ type CodeEnvironmentRegistryOptions = { configurationCache?: CodeEnvironmentConfigurationCache; }; +export function normalizeCodeEnvironmentName(input: string): string { + const name = input.trim(); + if (name.length < 1 || name.length > 100) { + throw new Error('Code environment name must contain between 1 and 100 characters'); + } + return name; +} + export class CodeEnvironmentValidationError extends Error { constructor(message: string) { super(message); @@ -74,7 +87,7 @@ export class CodeEnvironmentValidationError extends Error { function normalizeRegistration(input: CodeEnvironmentRegistration): CodeEnvironmentRegistration { const id = input.id.trim(); - const name = input.name.trim(); + const name = normalizeCodeEnvironmentName(input.name); const baseURL = input.baseURL.trim().replace(/\/+$/, ''); const controlPlaneId = input.controlPlaneId.trim(); const workerId = input.workerId?.trim(); @@ -95,6 +108,12 @@ function normalizeRegistration(input: CodeEnvironmentRegistration): CodeEnvironm if (workerId != null && !WORKER_ID_PATTERN.test(workerId)) { throw new CodeEnvironmentValidationError('Code environment worker id is invalid'); } + if ( + input.workerPrincipal != null && + !WORKER_PRINCIPAL_ID_PATTERN.test(input.workerPrincipal.id) + ) { + throw new CodeEnvironmentValidationError('Code environment worker principal is invalid'); + } return { ...input, id, name, baseURL, controlPlaneId, workerId }; } @@ -182,6 +201,7 @@ export function createCodeEnvironmentRegistry( baseURL: environment.baseURL, controlPlaneId: environment.controlPlaneId, workerId: environment.workerId, + workerPrincipal: environment.workerPrincipal, createdBy: new Types.ObjectId(actor.userId), }); try { @@ -227,7 +247,12 @@ export function createCodeEnvironmentRegistry( async function findAccessible(actor: CodeEnvironmentPrincipalContext) { const principals = actor.principals ?? (await methods.getUserPrincipals(actor)); const ids = await findAccessibleResourceIds(principals); - return await methods.findCodeEnvironmentsByIds(ids); + const environments = await methods.findCodeEnvironmentsByIds(ids); + const userId = actor.userId.toString(); + return environments.filter( + (environment) => + environment.workerPrincipal?.type !== 'user' || environment.workerPrincipal.id === userId, + ); } async function listAccessible( @@ -259,16 +284,23 @@ export function createCodeEnvironmentRegistry( const load = async (): Promise => { const ids = await findAccessibleResourceIds(principals); const environments = await methods.findCodeEnvironmentsByIds(ids); - return environments.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, - })); + 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, + })); }; if (configurationCache == null) return (await load()).map(toPublicConfiguration); diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 7ff2e89e4c0..c7a583308f7 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -123,6 +123,7 @@ describe('code environment HTTP handlers', () => { type: 'attached', baseURL: 'https://code.librechat.example/v1', workerId: 'code-generated', + workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, }, }); expect(res.body).toEqual({ @@ -189,25 +190,30 @@ describe('code environment HTTP handlers', () => { baseURL: 'https://code.librechat.example', controlPlaneId: 'shared-code-api', workerId: 'deployment-worker', + workerPrincipal: { type: 'deployment', id: 'shared-code-api' }, }, }); }); - test('returns a generic 500 for operational registration failures', async () => { - const register = jest.fn().mockRejectedValue(new Error('mongo connection details')); + test('does not register a fixed environment on a self-service-only control plane', async () => { + const register = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn().mockResolvedValue({ endpoints: { [EModelEndpoint.agents]: { statefulCodeSessions: { + allowedEnvironments: ['user'], environments: [ { - id: 'shared-code-api', - name: 'Shared Code API', + id: 'self-service', + name: 'Self-service', type: 'attached', baseURL: 'https://code.librechat.example', owner: 'deployment', - pairing: { workerId: 'deployment-worker' }, + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, }, ], }, @@ -220,20 +226,71 @@ describe('code environment HTTP handlers', () => { await handlers.register( { - user: { id: '68b2f0c498f24c1e78fa0001' }, - body: { name: 'Personal VM', controlPlaneId: 'shared-code-api' }, + user: { id: '68b2f0c498f24c1e78fa0001', role: 'ADMIN' }, + body: { name: 'Invalid fixed route', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(404); + expect(register).not.toHaveBeenCalled(); + }); + + test('validates a pairing name before creating upstream state', async () => { + const fetchImpl = jest.fn(); + 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', + owner: 'deployment', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + } as AppConfig), + registry: { register: jest.fn(), listAccessible: jest.fn() }, + readSecret: jest.fn(() => 'administrator-token'), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'x'.repeat(101), controlPlaneId: 'self-service' }, } as never, res as never, ); - expect(res.statusCode).toBe(500); - expect(res.body).toEqual({ error: 'Code environment registration failed' }); + expect(res.statusCode).toBe(400); + expect(fetchImpl).not.toHaveBeenCalled(); }); - test('returns 400 for validated registration input failures', async () => { - const register = jest - .fn() - .mockRejectedValue(new CodeEnvironmentValidationError('Code environment id is invalid')); + test.each([ + { + error: new Error('mongo connection details'), + status: 500, + body: { error: 'Code environment registration failed' }, + }, + { + error: new CodeEnvironmentValidationError('Code environment id is invalid'), + status: 400, + body: { error: 'Code environment id is invalid' }, + }, + ])('classifies registration failure with status $status', async ({ error, status, body }) => { const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn().mockResolvedValue({ endpoints: { @@ -253,7 +310,7 @@ describe('code environment HTTP handlers', () => { }, }, } as AppConfig), - registry: { register, listAccessible: jest.fn() }, + registry: { register: jest.fn().mockRejectedValue(error), listAccessible: jest.fn() }, }); const res = response(); @@ -265,7 +322,7 @@ describe('code environment HTTP handlers', () => { res as never, ); - expect(res.statusCode).toBe(400); - expect(res.body).toEqual({ error: 'Code environment id is invalid' }); + expect(res.statusCode).toBe(status); + expect(res.body).toEqual(body); }); }); diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index c6016799290..0774a25a299 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -11,7 +11,7 @@ import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; import type { CodeBridgeFetch } from './bridge'; import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; -import { CodeEnvironmentValidationError } from './environments'; +import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; import { getCodeApiTenantId } from '~/auth/codeapi'; type Registry = { @@ -56,7 +56,7 @@ function configuredControlPlane( environment.id === controlPlaneId && environment.type === 'attached' && environment.owner === 'deployment' && - environment.pairing != null, + environment.pairing?.workerId != null, ); } @@ -129,6 +129,13 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) error: 'name and controlPlaneId are required', }); } + try { + normalizeCodeEnvironmentName(name); + } catch (error) { + return res.status(400).json({ + error: error instanceof Error ? error.message : 'Code environment name is invalid', + }); + } /** Control-plane destinations are deployment policy. Client-provided URLs * are deliberately ignored to prevent an authenticated SSRF primitive. */ @@ -148,6 +155,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) baseURL: controlPlane.baseURL, controlPlaneId, workerId: controlPlane.pairing?.workerId, + workerPrincipal: { type: 'deployment', id: controlPlane.id }, }, }); return res.status(201).json({ environment }); @@ -180,6 +188,13 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (!name || !controlPlaneId) { return res.status(400).json({ error: 'name and controlPlaneId are required' }); } + try { + normalizeCodeEnvironmentName(name); + } catch (error) { + return res.status(400).json({ + error: error instanceof Error ? error.message : 'Code environment name is invalid', + }); + } const appConfig = await deps.getAppConfig({ baseOnly: true }); const controlPlane = configuredPrincipalControlPlane(appConfig, controlPlaneId); @@ -218,6 +233,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) type: 'attached', baseURL: controlPlane.baseURL, workerId, + workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); return res.status(201).json({ diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 43eff6ec9e1..d9a48993034 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -8,7 +8,7 @@ type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' > & - Pick, 'workerId'>; + Pick, 'workerId' | 'workerPrincipal'>; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index 13786c63cd6..fb297e5d350 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -1,6 +1,21 @@ import { Schema } from 'mongoose'; import type { CodeEnvironmentDocument } from '~/types'; +const workerPrincipalSchema = new Schema( + { + type: { + type: String, + enum: ['deployment', 'tenant', 'user', 'role', 'group'], + required: true, + }, + id: { + type: String, + required: true, + }, + }, + { _id: false }, +); + const codeEnvironmentSchema: Schema = new Schema( { environmentId: { @@ -33,6 +48,9 @@ const codeEnvironmentSchema: Schema = new Schema Date: Sun, 30 Aug 2026 14:43:59 -0400 Subject: [PATCH 03/29] fix: separate pairing control planes from execution --- packages/api/src/agents/execution.spec.ts | 36 ++++++++++++++ packages/api/src/agents/execution.ts | 14 +++++- packages/api/src/code/http.spec.ts | 24 ++++++++++ packages/api/src/code/http.ts | 14 +++++- .../src/endpoints/config/endpoints.spec.ts | 39 +++++++++++++++ .../api/src/endpoints/config/endpoints.ts | 15 ++++-- .../specs/config-schemas.spec.ts | 48 ++++++++++++++++++- packages/data-provider/src/config.ts | 30 +++++++++++- 8 files changed, 209 insertions(+), 11 deletions(-) diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts index c2c96c1f3cc..93453056963 100644 --- a/packages/api/src/agents/execution.spec.ts +++ b/packages/api/src/agents/execution.spec.ts @@ -165,6 +165,42 @@ describe('resolveCodeExecutionContext', () => { }); }); + it('does not execute a pairing-only control plane', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + const environments = [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached' as const, + baseURL: 'https://bridge.example/v1', + default: true, + owner: 'deployment' as const, + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ]; + + expect( + resolveCodeExecutionContext({ statefulSessions: true, environments, userId: 'user-1' }), + ).toEqual( + expect.objectContaining({ + baseUrl: 'http://code-stateful.test/v1', + environmentId: undefined, + bridgeWorkerId: undefined, + }), + ); + expect(() => + resolveCodeExecutionContext({ + statefulSessions: true, + environmentId: 'self-service', + environments, + userId: 'user-1', + }), + ).toThrow('Stateful code environment "self-service" is not configured'); + }); + it('adds the server-selected worker to execution auth without replacing authentication', async () => { const context = resolveCodeExecutionContext({ statefulSessions: true, diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index 953425922bf..1f7e3a042d4 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -104,14 +104,24 @@ function resolveConfiguredEnvironment(params: { environments?: readonly CodeEnvironmentConfig[]; }): CodeEnvironmentConfig | undefined { const { environmentId, environments } = params; + const executableEnvironments = environments?.filter( + (environment) => + !( + environment.pairing?.allowPrincipalWorkers === true && + environment.pairing.workerId == null && + environment.workerId == null + ), + ); if (environmentId) { - const configured = environments?.find((environment) => environment.id === environmentId); + const configured = executableEnvironments?.find( + (environment) => environment.id === environmentId, + ); if (!configured) { throw new Error(`Stateful code environment "${environmentId}" is not configured.`); } return configured; } - return environments?.find((environment) => environment.default === true); + return executableEnvironments?.find((environment) => environment.default === true); } export function resolveCodeExecutionContext(params: { diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index c7a583308f7..13cba36308b 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -81,6 +81,7 @@ describe('code environment HTTP handlers', () => { createEnvironmentId: () => 'code-generated', readSecret: jest.fn(() => 'administrator-token'), resolveTenantId: jest.fn(() => 'tenant-1'), + principalAuthEnabled: jest.fn(() => true), fetchImpl, }); const req = { @@ -132,6 +133,28 @@ describe('code environment HTTP handlers', () => { }); }); + 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() }, + principalAuthEnabled: jest.fn(() => false), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(409); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + test('registers against an operator-configured control plane and ignores client URLs', async () => { const register = jest.fn().mockResolvedValue({ resourceId: '68b2f0c498f24c1e78fa0111', @@ -263,6 +286,7 @@ describe('code environment HTTP handlers', () => { } as AppConfig), registry: { register: jest.fn(), listAccessible: jest.fn() }, readSecret: jest.fn(() => 'administrator-token'), + principalAuthEnabled: jest.fn(() => true), fetchImpl, }); const res = response(); diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 0774a25a299..39a4b8de24c 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -12,7 +12,7 @@ import type { ServerRequest } from '~/types/http'; import type { CodeBridgeFetch } from './bridge'; import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; -import { getCodeApiTenantId } from '~/auth/codeapi'; +import { getCodeApiTenantId, isCodeApiJwtAuthEnabled } from '~/auth/codeapi'; type Registry = { register: (params: { @@ -35,6 +35,7 @@ export interface CodeEnvironmentHttpDeps { createEnvironmentId?: () => string; readSecret?: (name: string) => string | undefined; resolveTenantId?: (req: ServerRequest) => string; + principalAuthEnabled?: () => boolean; fetchImpl?: CodeBridgeFetch; } @@ -102,6 +103,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) const createEnvironmentId = deps.createEnvironmentId ?? (() => `code-${nanoid(20)}`); const readSecret = deps.readSecret ?? readCodeBridgeSecret; const resolveTenantId = deps.resolveTenantId ?? getCodeApiTenantId; + const principalAuthEnabled = deps.principalAuthEnabled ?? isCodeApiJwtAuthEnabled; async function list(req: ServerRequest, res: Response): Promise { const principal = actor(req); @@ -181,7 +183,15 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (principal == null) { return res.status(401).json({ error: 'Authentication required' }); } - const body = req.body as unknown as Record; + if (!principalAuthEnabled()) { + return res.status(409).json({ + error: 'Principal code workers require Code API JWT authentication', + }); + } + const body = + typeof req.body === 'object' && req.body != null + ? (req.body as unknown as Record) + : {}; const name = typeof body.name === 'string' ? body.name.trim() : ''; const controlPlaneId = typeof body.controlPlaneId === 'string' ? body.controlPlaneId.trim() : ''; diff --git a/packages/api/src/endpoints/config/endpoints.spec.ts b/packages/api/src/endpoints/config/endpoints.spec.ts index bff43536885..c2bafd65078 100644 --- a/packages/api/src/endpoints/config/endpoints.spec.ts +++ b/packages/api/src/endpoints/config/endpoints.spec.ts @@ -225,6 +225,45 @@ describe('createEndpointsConfigService', () => { }); }); + it('does not expose a pairing-only control plane as an execution environment', async () => { + const deps = createMockDeps({ + loadDefaultEndpointsConfig: jest.fn().mockResolvedValue({ + [EModelEndpoint.agents]: { userProvide: false, order: 0 }, + }), + getAppConfig: jest.fn().mockResolvedValue( + appConfig({ + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached', + baseURL: 'https://internal-code.example.com/v1', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + }), + ), + }); + const { getEndpointsConfig } = createEndpointsConfigService(deps); + + const result = await getEndpointsConfig(fakeReq()); + + expect(result?.[EModelEndpoint.agents]?.statefulCodeSessions).toEqual({ + allowedEnvironments: ['user'], + environments: [], + }); + }); + it('merges bedrock availableRegions', async () => { const deps = createMockDeps({ loadDefaultEndpointsConfig: jest.fn().mockResolvedValue({ diff --git a/packages/api/src/endpoints/config/endpoints.ts b/packages/api/src/endpoints/config/endpoints.ts index 4375e963126..c1ce607a49e 100644 --- a/packages/api/src/endpoints/config/endpoints.ts +++ b/packages/api/src/endpoints/config/endpoints.ts @@ -75,14 +75,21 @@ export function createEndpointsConfigService(deps: EndpointsConfigDeps): { const clientStatefulCodeSessions = statefulCodeSessions ? { allowedEnvironments: statefulCodeSessions.allowedEnvironments, - environments: statefulCodeSessions.environments?.map( - ({ id, name, type, default: isDefault }) => ({ + environments: statefulCodeSessions.environments + ?.filter( + (environment) => + !( + environment.pairing?.allowPrincipalWorkers === true && + environment.pairing.workerId == null && + environment.workerId == null + ), + ) + .map(({ id, name, type, default: isDefault }) => ({ id, name, type, default: isDefault, - }), - ), + })), } : undefined; mergedConfig[EModelEndpoint.agents] = { diff --git a/packages/data-provider/specs/config-schemas.spec.ts b/packages/data-provider/specs/config-schemas.spec.ts index 057438f72c6..4e6eb140495 100644 --- a/packages/data-provider/specs/config-schemas.spec.ts +++ b/packages/data-provider/specs/config-schemas.spec.ts @@ -586,7 +586,6 @@ describe('agentsEndpointSchema', () => { name: 'Personal Code', type: 'attached', baseURL: 'https://bridge.example.com/v1', - default: true, pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', @@ -599,6 +598,53 @@ describe('agentsEndpointSchema', () => { expect(result.success).toBe(true); }); + it('rejects a pairing-only control plane as the execution default', () => { + const result = agentsEndpointSchema.safeParse({ + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'personal-code-control-plane', + name: 'Personal Code', + type: 'attached', + baseURL: 'https://bridge.example.com/v1', + default: true, + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }); + + expect(result.success).toBe(false); + }); + + it('rejects conflicting direct and pairing worker routes', () => { + const result = agentsEndpointSchema.safeParse({ + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'attached-vm', + name: 'Attached VM', + type: 'attached', + baseURL: 'https://bridge.example.com/v1', + default: true, + workerId: 'direct-worker', + pairing: { + workerId: 'paired-worker', + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }); + + expect(result.success).toBe(false); + }); + it('rejects pairing metadata on a principal-owned environment', () => { const result = agentsEndpointSchema.safeParse({ statefulCodeSessions: { diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 5cf4c5a9bc8..a1fcd77d00a 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1121,7 +1121,12 @@ export const agentsEndpointSchema = baseEndpointSchema if (!value?.environments) return; const ids = new Set(); let defaults = 0; + let executableEnvironments = 0; for (const environment of value.environments) { + const pairingOnly = + environment.pairing?.allowPrincipalWorkers === true && + environment.pairing.workerId == null && + environment.workerId == null; if (environment.pairing != null && environment.type !== 'attached') { context.addIssue({ code: z.ZodIssueCode.custom, @@ -1146,6 +1151,24 @@ export const agentsEndpointSchema = baseEndpointSchema path: ['environments', environment.id, 'baseURL'], }); } + if ( + environment.workerId != null && + environment.pairing?.workerId != null && + environment.workerId !== environment.pairing.workerId + ) { + context.addIssue({ + code: z.ZodIssueCode.custom, + message: 'Code environment workerId must match pairing.workerId', + path: ['environments', environment.id, 'workerId'], + }); + } + if (pairingOnly && environment.default === true) { + context.addIssue({ + code: z.ZodIssueCode.custom, + message: 'Pairing-only code control planes cannot be execution defaults', + path: ['environments', environment.id, 'default'], + }); + } if (ids.has(environment.id)) { context.addIssue({ code: z.ZodIssueCode.custom, @@ -1154,9 +1177,12 @@ export const agentsEndpointSchema = baseEndpointSchema }); } ids.add(environment.id); - if (environment.default === true) defaults += 1; + if (!pairingOnly) { + executableEnvironments += 1; + if (environment.default === true) defaults += 1; + } } - if (value.environments.length > 0 && defaults !== 1) { + if (executableEnvironments > 0 && defaults !== 1) { context.addIssue({ code: z.ZodIssueCode.custom, message: 'Exactly one stateful code environment must be the default', From 25238d7c71409589c40ca96b2982a66409a27627 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 14:59:18 -0400 Subject: [PATCH 04/29] fix: preserve principal worker routing --- api/server/controllers/agents/v1.js | 9 ++++- api/server/controllers/agents/v1.spec.js | 36 +++++++++++++++++++ api/server/services/Files/Code/process.js | 12 +++++-- .../services/Files/Code/process.spec.js | 7 ++++ packages/api/src/agents/handlers.spec.ts | 2 ++ packages/api/src/agents/handlers.ts | 5 +++ 6 files changed, 67 insertions(+), 4 deletions(-) diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index f412b152ff5..7cf87113a46 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -457,7 +457,14 @@ const validateStatefulCodeEnvironment = ( if (environmentId != null) { const configuredEnvironments = req.config?.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments ?? []; - if (!configuredEnvironments.some((configured) => configured.id === environmentId)) { + const configuredEnvironment = configuredEnvironments.find( + (configured) => configured.id === environmentId, + ); + const pairingOnly = + configuredEnvironment?.pairing?.allowPrincipalWorkers === true && + configuredEnvironment.pairing.workerId == null && + configuredEnvironment.workerId == null; + if (configuredEnvironment == null || pairingOnly) { res.status(400).json({ error: `Stateful code environment is not configured: ${environmentId}`, }); diff --git a/api/server/controllers/agents/v1.spec.js b/api/server/controllers/agents/v1.spec.js index ef1345586fd..8a59c0ab93f 100644 --- a/api/server/controllers/agents/v1.spec.js +++ b/api/server/controllers/agents/v1.spec.js @@ -270,6 +270,42 @@ describe('Agent Controllers - Mass Assignment Protection', () => { expect(await Agent.countDocuments()).toBe(0); }); + test('rejects a pairing-only control plane as a code environment id', async () => { + mockReq.config = { + endpoints: { + agents: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'personal-workers', + name: 'Personal workers', + type: 'attached', + baseURL: 'https://code.example.com/v1', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + }; + mockReq.body = { + name: 'Invalid Control Plane Agent', + provider: 'openai', + model: 'gpt-4', + stateful_code_sessions: true, + code_environment_id: 'personal-workers', + }; + + await createAgentHandler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(400); + expect(await Agent.countDocuments()).toBe(0); + }); + test('should block configured agent instruction content before persistence', async () => { mockReq.config = { filters: { diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 6bb352e0114..b311a7babc4 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -22,6 +22,7 @@ const { getStorageMetadata, getCodeExecutionBaseUrl, buildCodeEnvDownloadQuery, + codeExecutionHeaders, CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); const { @@ -1395,6 +1396,7 @@ async function readSandboxFile({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, req, }) { const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); @@ -1429,7 +1431,7 @@ async function readSandboxFile({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -1480,6 +1482,7 @@ async function readSandboxImage({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, maxBytes, req, }) { @@ -1543,6 +1546,7 @@ async function readSandboxImage({ session_id, runtime_session_hint, executionProfile, + bridgeWorkerId, files, req, chunkBytes, @@ -1614,6 +1618,7 @@ async function execSandboxImageChunk({ session_id, runtime_session_hint, executionProfile, + bridgeWorkerId, files, req, chunkBytes, @@ -1640,7 +1645,7 @@ async function execSandboxImageChunk({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -1708,6 +1713,7 @@ async function writeSandboxFile({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, req, }) { const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); @@ -1760,7 +1766,7 @@ async function writeSandboxFile({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 8e2b3b258a8..caf1b248f96 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -72,6 +72,10 @@ jest.mock('@librechat/api', () => { profile === 'stateful' ? 'https://code-stateful.example.com' : 'https://code-api.example.com', ), CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', + codeExecutionHeaders: ({ executionProfile, bridgeWorkerId }) => ({ + 'X-CodeAPI-Expected-Profile': executionProfile, + ...(bridgeWorkerId ? { 'X-LibreChat-Code-Worker-ID': bridgeWorkerId } : {}), + }), withTimeout: (...args) => passthroughWithTimeout(...args), hasOfficeHtmlPath: (...args) => mockHasOfficeHtmlPath(...args), /** @@ -984,6 +988,7 @@ describe('Code Process', () => { ...baseParams, codeApiBaseUrl: 'https://code-stateful.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', executionRouteKey, }); @@ -1853,12 +1858,14 @@ describe('Code Process', () => { file_path: '/mnt/data/x.txt', codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', runtime_session_hint: 'v1:user', }); const call = mockAxios.mock.calls[0][0]; expect(call.url).toBe('https://stateful-code.example.com/exec'); expect(call.headers['X-CodeAPI-Expected-Profile']).toBe('stateful'); + expect(call.headers['X-LibreChat-Code-Worker-ID']).toBe('personal-worker-1'); expect(call.data.runtime_session_hint).toBe('v1:user'); }); diff --git a/packages/api/src/agents/handlers.spec.ts b/packages/api/src/agents/handlers.spec.ts index 025f75414a5..2f31ce6592c 100644 --- a/packages/api/src/agents/handlers.spec.ts +++ b/packages/api/src/agents/handlers.spec.ts @@ -4430,6 +4430,7 @@ describe('createToolExecuteHandler', () => { baseUrl: 'https://stateful-code.example.com', codeSessionKey: 'execute_code:stateful:v1:user', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', runtimeSessionHint: 'v1:user', statefulSessions: true, }, @@ -4449,6 +4450,7 @@ describe('createToolExecuteHandler', () => { expect.objectContaining({ codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', runtime_session_hint: 'v1:user', }), ); diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 2004e1f8307..5838fa668c8 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -386,6 +386,7 @@ export interface ToolExecuteOptions { read_only?: boolean; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; }) => Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }>; @@ -447,6 +448,7 @@ export interface ToolExecuteOptions { runtime_session_hint?: string; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; req?: ServerRequest; }) => Promise<{ content: string } | null>; /** @@ -467,6 +469,7 @@ export interface ToolExecuteOptions { runtime_session_hint?: string; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; /** In-sandbox size cap; files larger than this return `tooLarge` without transferring bytes. */ maxBytes?: number; req?: ServerRequest; @@ -544,6 +547,7 @@ function getCodeExecutionContext( function codeExecutionRequestParams(context?: CodeExecutionContext): { codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; runtime_session_hint?: string; } { if (!context) { @@ -552,6 +556,7 @@ function codeExecutionRequestParams(context?: CodeExecutionContext): { return { codeApiBaseUrl: context.baseUrl, executionProfile: context.executionProfile, + ...(context.bridgeWorkerId ? { bridgeWorkerId: context.bridgeWorkerId } : {}), ...(context.runtimeSessionHint ? { runtime_session_hint: context.runtimeSessionHint } : {}), }; } From ab726fdb5c9c55677823147211e3408862613d40 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 15:39:26 -0400 Subject: [PATCH 05/29] fix: pin artifact traffic to code workers --- api/server/controllers/agents/callbacks.js | 2 ++ api/server/routes/files/files.js | 11 ++++++++++- api/server/routes/files/files.test.js | 7 ++++++- api/server/services/Files/Code/crud.js | 14 +++++++++++--- api/server/services/Files/Code/crud.spec.js | 10 +++++++++- api/server/services/Files/Code/preflight.js | 1 + api/server/services/Files/Code/process.js | 17 +++++++++++++---- api/server/services/Files/Code/process.spec.js | 5 ++++- packages/api/src/agents/harvest.ts | 2 ++ 9 files changed, 58 insertions(+), 11 deletions(-) diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 5f3579fd340..4117eb993bd 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -1048,6 +1048,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, executionProfile: metadata.codeExecutionContext?.executionProfile, executionRouteKey: metadata.codeExecutionContext?.executionRouteKey, + bridgeWorkerId: metadata.codeExecutionContext?.bridgeWorkerId, preparedBuffer, downloadFallback, }); @@ -1393,6 +1394,7 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises }) codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, executionProfile: metadata.codeExecutionContext?.executionProfile, executionRouteKey: metadata.codeExecutionContext?.executionRouteKey, + bridgeWorkerId: metadata.codeExecutionContext?.bridgeWorkerId, preparedBuffer, downloadFallback, }); diff --git a/api/server/routes/files/files.js b/api/server/routes/files/files.js index abb2713e925..7b2c3935c8e 100644 --- a/api/server/routes/files/files.js +++ b/api/server/routes/files/files.js @@ -398,7 +398,16 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => { id: req.user.id, }, req, - { baseUrl, executionProfile }, + { + baseUrl, + executionProfile, + ...((configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId) != null + ? { + bridgeWorkerId: + configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId, + } + : {}), + }, ); res.setHeader('Content-Disposition', 'attachment'); res.setHeader('Content-Type', 'application/octet-stream'); diff --git a/api/server/routes/files/files.test.js b/api/server/routes/files/files.test.js index af154d0feeb..8d80a9bbf06 100644 --- a/api/server/routes/files/files.test.js +++ b/api/server/routes/files/files.test.js @@ -1263,6 +1263,7 @@ describe('File Routes - Delete with Agent Access', () => { name: 'Managed VM', type: 'managed', baseURL: 'https://managed-code.example.com/v1', + workerId: 'personal-worker-1', default: true, owner: 'deployment', }; @@ -1290,7 +1291,11 @@ describe('File Routes - Delete with Agent Access', () => { `${sessionId}/${codeFileId}`, { kind: 'user', id: otherUserId.toString() }, expect.any(Object), - { baseUrl: environment.baseURL, executionProfile: 'stateful' }, + { + baseUrl: environment.baseURL, + executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', + }, ); }); diff --git a/api/server/services/Files/Code/crud.js b/api/server/services/Files/Code/crud.js index 27986ca2158..6999766b915 100644 --- a/api/server/services/Files/Code/crud.js +++ b/api/server/services/Files/Code/crud.js @@ -13,6 +13,7 @@ const { getCodeApiAuthHeaders, getCodeExecutionBaseUrl, createCodeExecutionRouteKey, + codeExecutionHeaders, CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); @@ -29,7 +30,7 @@ const MAX_FILE_SIZE = 150 * 1024 * 1024; * `kind: 'user', id: `; for skill/agent re-downloads pass * the kind+id (+version for skill) from the file's `metadata.codeEnvRef`. * @param {ServerRequest} req - Current authenticated request. - * @param {{baseUrl?: string, executionProfile?: 'default'|'stateful'}} [route] + * @param {{baseUrl?: string, executionProfile?: 'default'|'stateful', bridgeWorkerId?: string}} [route] * Trusted host-selected Code API route. * @returns {Promise} A promise that resolves to a readable stream of the file content. * @throws {Error} If there's an error during the download process. @@ -48,7 +49,10 @@ async function getCodeOutputDownloadStream(fileIdentifier, identity, req, route 'User-Agent': 'LibreChat/1.0', ...authHeaders, ...(route.executionProfile - ? { [CODE_API_EXPECTED_PROFILE_HEADER]: route.executionProfile } + ? codeExecutionHeaders({ + executionProfile: route.executionProfile, + bridgeWorkerId: route.bridgeWorkerId, + }) : {}), }, httpAgent: codeServerHttpAgent, @@ -127,7 +131,11 @@ async function deleteCodeEnvFile(req, file) { headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, - [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, + ...codeExecutionHeaders({ + executionProfile, + bridgeWorkerId: + configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId, + }), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, diff --git a/api/server/services/Files/Code/crud.spec.js b/api/server/services/Files/Code/crud.spec.js index c4a5dec4573..c8e69c40239 100644 --- a/api/server/services/Files/Code/crud.spec.js +++ b/api/server/services/Files/Code/crud.spec.js @@ -53,6 +53,10 @@ jest.mock('@librechat/api', () => { getCodeExecutionBaseUrl: jest.fn((profile) => profile === 'stateful' ? 'https://code-stateful.example.com' : 'https://code-api.example.com', ), + codeExecutionHeaders: jest.fn(({ executionProfile, bridgeWorkerId }) => ({ + 'X-CodeAPI-Expected-Profile': executionProfile, + ...(bridgeWorkerId ? { 'X-LibreChat-Code-Worker-ID': bridgeWorkerId } : {}), + })), CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', createAxiosInstance: jest.fn(() => mockAxios), codeServerHttpAgent: new http.Agent({ keepAlive: false }), @@ -118,12 +122,16 @@ describe('Code CRUD', () => { await getCodeOutputDownloadStream('session-1/file-1', userIdentity, undefined, { baseUrl: 'https://code-stateful.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ url: 'https://code-stateful.example.com/download/session-1/file-1?kind=user&id=user-123', - headers: expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'stateful' }), + headers: expect.objectContaining({ + 'X-CodeAPI-Expected-Profile': 'stateful', + 'X-LibreChat-Code-Worker-ID': 'personal-worker-1', + }), }), ); }); diff --git a/api/server/services/Files/Code/preflight.js b/api/server/services/Files/Code/preflight.js index 5b67c75d00d..a21c9516cb0 100644 --- a/api/server/services/Files/Code/preflight.js +++ b/api/server/services/Files/Code/preflight.js @@ -34,6 +34,7 @@ const preflightCodeOutputBatch = async ({ req, artifact, codeExecutionContext }) inspectContent, codeApiBaseUrl: codeExecutionContext?.baseUrl, executionProfile: codeExecutionContext?.executionProfile, + bridgeWorkerId: codeExecutionContext?.bridgeWorkerId, }), onInspectionUnavailable: (index) => { logger.warn( diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index b311a7babc4..d8e4a83e998 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -23,7 +23,6 @@ const { getCodeExecutionBaseUrl, buildCodeEnvDownloadQuery, codeExecutionHeaders, - CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); const { Tools, @@ -80,6 +79,7 @@ const downloadCodeOutputBuffer = async ({ maxBytes, codeApiBaseUrl, executionProfile = 'default', + bridgeWorkerId, }) => { const baseURL = codeApiBaseUrl ?? getCodeExecutionBaseUrl(executionProfile); const authHeaders = await getCodeApiAuthHeaders(req); @@ -93,7 +93,7 @@ const downloadCodeOutputBuffer = async ({ headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, - [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -134,6 +134,7 @@ const downloadCodeOutputBuffer = async ({ * @param {number} [params.maxBytes] - Remaining aggregate inspection budget. * @param {string} [params.codeApiBaseUrl] - Trusted per-agent Code API endpoint. * @param {'default'|'stateful'} [params.executionProfile] - Trusted execution profile. + * @param {string} [params.bridgeWorkerId] - Trusted worker selected for this execution. */ const prepareCodeOutputForInspection = async ({ req, @@ -144,6 +145,7 @@ const prepareCodeOutputForInspection = async ({ inspectContent = true, codeApiBaseUrl, executionProfile = 'default', + bridgeWorkerId, }) => { const { fileSizeLimit } = getCodeOutputFileSettings(req); const transportLimit = @@ -155,6 +157,7 @@ const prepareCodeOutputForInspection = async ({ maxBytes: transportLimit, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }); const safeName = sanitizeArtifactPath(name); const fallbackType = inferMimeType(name, '') || 'application/octet-stream'; @@ -524,6 +527,7 @@ const processCodeOutput = async ({ codeApiBaseUrl, executionProfile = 'default', executionRouteKey = executionProfile, + bridgeWorkerId, preparedBuffer, downloadFallback, }) => { @@ -561,6 +565,7 @@ const processCodeOutput = async ({ maxBytes: fileSizeLimit, codeApiBaseUrl, executionProfile, + bridgeWorkerId, })); // Enforce file size limit @@ -1002,7 +1007,10 @@ async function getSessionInfo(ref, req, route = {}) { 'User-Agent': 'LibreChat/1.0', ...authHeaders, ...(route.executionProfile - ? { [CODE_API_EXPECTED_PROFILE_HEADER]: route.executionProfile } + ? codeExecutionHeaders({ + executionProfile: route.executionProfile, + bridgeWorkerId: route.bridgeWorkerId, + }) : {}), }, httpAgent: codeServerHttpAgent, @@ -1120,8 +1128,9 @@ const primeFiles = async (options) => { codeApiBaseUrl, executionProfile = 'default', executionRouteKey = executionProfile, + bridgeWorkerId, } = options; - const codeApiRoute = { baseUrl: codeApiBaseUrl, executionProfile }; + const codeApiRoute = { baseUrl: codeApiBaseUrl, executionProfile, bridgeWorkerId }; const file_ids = tool_resources?.[EToolResources.execute_code]?.file_ids ?? []; const agentResourceIds = new Set(file_ids); const resourceFiles = tool_resources?.[EToolResources.execute_code]?.files ?? []; diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index caf1b248f96..2b059fca732 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -995,7 +995,10 @@ describe('Code Process', () => { expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ url: expect.stringContaining('https://code-stateful.example.com/download/'), - headers: expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'stateful' }), + headers: expect.objectContaining({ + 'X-CodeAPI-Expected-Profile': 'stateful', + 'X-LibreChat-Code-Worker-ID': 'personal-worker-1', + }), }), ); expect(result.filepath).toContain('execution_profile=stateful'); diff --git a/packages/api/src/agents/harvest.ts b/packages/api/src/agents/harvest.ts index 8fd46bef837..84f67cbfdf2 100644 --- a/packages/api/src/agents/harvest.ts +++ b/packages/api/src/agents/harvest.ts @@ -79,6 +79,7 @@ export interface CodeHarvestDeps { codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; executionRouteKey?: string; + bridgeWorkerId?: string; preparedBuffer?: Buffer; downloadFallback?: boolean; }) => Promise; @@ -284,6 +285,7 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa codeApiBaseUrl: codeExecutionContext?.baseUrl, executionProfile: codeExecutionContext?.executionProfile, executionRouteKey: codeExecutionContext?.executionRouteKey, + bridgeWorkerId: codeExecutionContext?.bridgeWorkerId, preparedBuffer, downloadFallback, }); From 0ea75c50a95678a274ba1d4f608da407a088930d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 15:58:48 -0400 Subject: [PATCH 06/29] fix: authenticate code worker selection --- api/app/clients/tools/util/handleTools.js | 2 +- api/server/services/Files/Code/crud.js | 9 ++++---- api/server/services/Files/Code/crud.spec.js | 4 ++-- api/server/services/Files/Code/process.js | 10 ++++----- .../services/Files/Code/process.spec.js | 6 +++--- api/server/services/ToolService.js | 15 ++++++++++--- packages/api/src/agents/execution.spec.ts | 10 +++++---- packages/api/src/agents/execution.ts | 6 ++++-- packages/api/src/auth/codeapi.spec.ts | 10 +++++++++ packages/api/src/auth/codeapi.ts | 21 ++++++++++++++----- 10 files changed, 64 insertions(+), 29 deletions(-) diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 13fd231c61d..0cab206daae 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -382,7 +382,7 @@ const loadTools = async ({ files, authHeaders: () => codeExecutionAuthHeaders( - () => getCodeApiAuthHeaders(options.req), + (bridgeWorkerId) => getCodeApiAuthHeaders(options.req, bridgeWorkerId), codeExecutionContext, ), ...codeExecutionContext, diff --git a/api/server/services/Files/Code/crud.js b/api/server/services/Files/Code/crud.js index 6999766b915..76945e06335 100644 --- a/api/server/services/Files/Code/crud.js +++ b/api/server/services/Files/Code/crud.js @@ -39,7 +39,7 @@ async function getCodeOutputDownloadStream(fileIdentifier, identity, req, route try { const baseURL = route.baseUrl ?? getCodeBaseURL(); const query = buildCodeEnvDownloadQuery(identity); - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, route.bridgeWorkerId); /** @type {import('axios').AxiosRequestConfig} */ const options = { method: 'get', @@ -86,7 +86,6 @@ async function deleteCodeEnvFile(req, file) { } const missingOrUnsupportedStatuses = new Set([404, 405]); - const authHeaders = await getCodeApiAuthHeaders(req); for (const [executionRouteKey, ref] of refs) { const executionProfile = ref.executionProfile ?? 'default'; const environments = @@ -126,6 +125,9 @@ async function deleteCodeEnvFile(req, file) { id: ref.id, ...(ref.kind === 'skill' ? { version: ref.version } : {}), }); + const bridgeWorkerId = + configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId; + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); const baseRequest = { method: 'delete', headers: { @@ -133,8 +135,7 @@ async function deleteCodeEnvFile(req, file) { ...authHeaders, ...codeExecutionHeaders({ executionProfile, - bridgeWorkerId: - configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.workerId, + bridgeWorkerId, }), }, httpAgent: codeServerHttpAgent, diff --git a/api/server/services/Files/Code/crud.spec.js b/api/server/services/Files/Code/crud.spec.js index c8e69c40239..7af9c52b682 100644 --- a/api/server/services/Files/Code/crud.spec.js +++ b/api/server/services/Files/Code/crud.spec.js @@ -144,7 +144,7 @@ describe('Code CRUD', () => { await getCodeOutputDownloadStream('session-1/file-1', userIdentity, req); const callConfig = mockAxios.mock.calls[0][0]; - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req, undefined); expect(callConfig.headers.Authorization).toBe('Bearer codeapi-token'); }); @@ -203,7 +203,7 @@ describe('Code CRUD', () => { await deleteCodeEnvFile(req, file); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'delete', diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index d8e4a83e998..7a679c92bd6 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -82,7 +82,7 @@ const downloadCodeOutputBuffer = async ({ bridgeWorkerId, }) => { const baseURL = codeApiBaseUrl ?? getCodeExecutionBaseUrl(executionProfile); - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); const downloadQuery = buildCodeEnvDownloadQuery({ kind: 'user', id: req.user.id }); let response; try { @@ -988,7 +988,7 @@ function checkIfActive(dateString) { async function getSessionInfo(ref, req, route = {}) { try { const baseURL = route.baseUrl ?? getCodeBaseURL(); - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, route.bridgeWorkerId); /* `/sessions/.../objects/...` is gated by codeapi's `sessionAuth` * middleware (post-Phase C). The middleware reconstructs the * sessionKey from the URL query (`kind`/`id`/`version?`) plus the @@ -1431,7 +1431,7 @@ async function readSandboxFile({ } try { - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); const response = await axios({ method: 'post', url: `${baseURL}/exec`, @@ -1645,7 +1645,7 @@ async function execSandboxImageChunk({ } try { - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); const response = await axios({ method: 'post', url: `${baseURL}/exec`, @@ -1766,7 +1766,7 @@ async function writeSandboxFile({ } try { - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); const response = await axios({ method: 'post', url: `${baseURL}/exec`, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 2b059fca732..e0984fb9bb1 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -511,7 +511,7 @@ describe('Code Process', () => { await processCodeOutput(baseParams); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'get', @@ -1412,7 +1412,7 @@ describe('Code Process', () => { mockReq, ); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'get', @@ -1934,7 +1934,7 @@ describe('Code Process', () => { await readSandboxFile({ file_path: '/mnt/data/x.txt', req: mockReq }); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'post', diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 8f1956ed60b..bf6f7dc8f1c 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -1696,7 +1696,10 @@ async function loadAgentTools({ programmaticToolsEnabled, codeExecutionEnabled, authHeaders: () => - codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), codeExecutionContext, }); @@ -2082,7 +2085,10 @@ async function loadToolsForExecution({ for (const name of ptcToolNames) { const ptcTool = createBashProgrammaticToolCallingTool({ authHeaders: () => - codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), baseUrl: codeExecutionContext.baseUrl, executionProfile: codeExecutionContext.executionProfile, runtimeSessionHint: codeExecutionContext.runtimeSessionHint, @@ -2109,7 +2115,10 @@ async function loadToolsForExecution({ try { const bashTool = createBashExecutionTool({ authHeaders: () => - codeExecutionAuthHeaders(() => getCodeApiAuthHeaders(req), codeExecutionContext), + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), ...codeExecutionContext, }); allLoadedTools.push(bashTool); diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts index 93453056963..517bc273eb5 100644 --- a/packages/api/src/agents/execution.spec.ts +++ b/packages/api/src/agents/execution.spec.ts @@ -222,13 +222,15 @@ describe('resolveCodeExecutionContext', () => { 'X-CodeAPI-Expected-Profile': 'stateful', 'X-LibreChat-Code-Worker-ID': 'opaque-worker-id', }); - await expect( - codeExecutionAuthHeaders(async () => ({ Authorization: 'Bearer user-token' }), context), - ).resolves.toEqual({ - Authorization: 'Bearer user-token', + const authHeaders = jest.fn(async (workerId?: string) => ({ + Authorization: `Bearer user-token-for-${workerId ?? 'default'}`, + })); + await expect(codeExecutionAuthHeaders(authHeaders, context)).resolves.toEqual({ + Authorization: 'Bearer user-token-for-opaque-worker-id', 'X-CodeAPI-Expected-Profile': 'stateful', 'X-LibreChat-Code-Worker-ID': 'opaque-worker-id', }); + expect(authHeaders).toHaveBeenCalledWith('opaque-worker-id'); }); it('namespaces configured deployments independently of the shared wire profile', () => { diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index 1f7e3a042d4..b4f69520656 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -180,11 +180,13 @@ export function codeExecutionHeaders( } export async function codeExecutionAuthHeaders( - authHeaders: () => Promise> | Record, + authHeaders: ( + bridgeWorkerId?: string, + ) => Promise> | Record, context: Pick, ): Promise> { return { - ...(await authHeaders()), + ...(await authHeaders(context.bridgeWorkerId)), ...codeExecutionHeaders(context), }; } diff --git a/packages/api/src/auth/codeapi.spec.ts b/packages/api/src/auth/codeapi.spec.ts index 11353b5bede..c0e344a3d21 100644 --- a/packages/api/src/auth/codeapi.spec.ts +++ b/packages/api/src/auth/codeapi.spec.ts @@ -211,6 +211,16 @@ describe('Code API JWT minting', () => { expect(claims).not.toHaveProperty('planId'); }); + it('binds cached Code API tokens to the selected code worker', async () => { + const req = baseRequest(); + const first = await mintCodeApiToken(req, 'principal-worker-a'); + const second = await mintCodeApiToken(req, 'principal-worker-b'); + + expect(decodeToken(first).claims.code_worker_id).toBe('principal-worker-a'); + expect(decodeToken(second).claims.code_worker_id).toBe('principal-worker-b'); + expect(second).not.toBe(first); + }); + it('uses the single-tenant namespace when tenant context is absent outside strict mode', async () => { mockGetTenantId.mockReturnValue(undefined); diff --git a/packages/api/src/auth/codeapi.ts b/packages/api/src/auth/codeapi.ts index a4a25edde77..d655bcee001 100644 --- a/packages/api/src/auth/codeapi.ts +++ b/packages/api/src/auth/codeapi.ts @@ -38,6 +38,7 @@ interface CodeApiClaims { service_id?: string; chc_user_id?: string; plan_id?: string; + code_worker_id?: string; auth_context_hash: string; } @@ -263,7 +264,12 @@ function canonicalContextHash(input: { return createHash('sha256').update(JSON.stringify(canonical)).digest('hex'); } -function buildClaims(req: ServerRequest, config: SigningConfig, now: number): CodeApiClaims { +function buildClaims( + req: ServerRequest, + config: SigningConfig, + now: number, + codeWorkerId?: string, +): CodeApiClaims { const user = resolveUser(req); const userId = resolveUserId(user); const tenantId = getCodeApiTenantId(req); @@ -299,6 +305,7 @@ function buildClaims(req: ServerRequest, config: SigningConfig, now: number): Co ...(serviceId ? { service_id: serviceId } : {}), ...(chcUserId ? { chc_user_id: chcUserId } : {}), ...(planId ? { plan_id: planId } : {}), + ...(codeWorkerId != null && codeWorkerId !== '' ? { code_worker_id: codeWorkerId } : {}), auth_context_hash: authContextHash, }; } @@ -330,6 +337,7 @@ function cacheKey(config: SigningConfig, claims: CodeApiClaims): string { claims.service_id ?? '', claims.chc_user_id ?? '', claims.plan_id ?? '', + claims.code_worker_id ?? '', claims.auth_context_hash, ].join(':'); } @@ -350,7 +358,7 @@ function pruneTokenCache(now: number): void { } } -export async function mintCodeApiToken(req: ServerRequest): Promise { +export async function mintCodeApiToken(req: ServerRequest, codeWorkerId?: string): Promise { if (!isCodeApiJwtAuthEnabled()) { return ''; } @@ -358,7 +366,7 @@ export async function mintCodeApiToken(req: ServerRequest): Promise { const config = getSigningConfig(); const now = Math.floor(Date.now() / 1000); pruneTokenCache(now); - const claims = buildClaims(req, config, now); + const claims = buildClaims(req, config, now, codeWorkerId); const key = cacheKey(config, claims); const cached = tokenCache.get(key); if ( @@ -381,10 +389,13 @@ export async function mintCodeApiToken(req: ServerRequest): Promise { return token; } -export async function getCodeApiAuthHeaders(req?: ServerRequest): Promise> { +export async function getCodeApiAuthHeaders( + req?: ServerRequest, + codeWorkerId?: string, +): Promise> { if (!req || !isCodeApiJwtAuthEnabled()) { return {}; } - const token = await mintCodeApiToken(req); + const token = await mintCodeApiToken(req, codeWorkerId); return token ? { Authorization: `Bearer ${token}` } : {}; } From 0712566b6e4adc9c2f7f52558d4253ae36f186e8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 17:11:20 -0400 Subject: [PATCH 07/29] test: mock worker-aware artifact headers --- .../services/Files/Code/__tests__/process-traversal.spec.js | 1 + 1 file changed, 1 insertion(+) diff --git a/api/server/services/Files/Code/__tests__/process-traversal.spec.js b/api/server/services/Files/Code/__tests__/process-traversal.spec.js index 87ab0393fa9..6866e3e309c 100644 --- a/api/server/services/Files/Code/__tests__/process-traversal.spec.js +++ b/api/server/services/Files/Code/__tests__/process-traversal.spec.js @@ -26,6 +26,7 @@ jest.mock('@librechat/api', () => { flattenArtifactPath: mockFlattenArtifactPath, createAxiosInstance: jest.fn(() => mockAxios), getCodeApiAuthHeaders: jest.fn(async () => ({})), + codeExecutionHeaders: jest.fn(() => ({})), getCodeExecutionBaseUrl: jest.fn(() => 'http://localhost:8000'), CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', classifyCodeArtifact: jest.fn(() => 'other'), From 6ea3bab7f2156b23b1c1e4a6b6b9709c8ef6fdda Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 19:45:02 -0400 Subject: [PATCH 08/29] fix: harden self-service code worker routing --- api/server/middleware/limiters/code.js | 36 ++++++++++ api/server/middleware/limiters/code.spec.js | 70 +++++++++++++++++++ api/server/routes/code-environments.js | 5 +- api/server/routes/code-environments.test.js | 9 ++- api/server/services/Files/Code/crud.js | 13 ++-- api/server/services/Files/Code/crud.spec.js | 43 +++++++++++- api/server/services/Files/Code/process.js | 4 +- .../services/Files/Code/process.spec.js | 2 + packages/api/src/agents/execution.spec.ts | 26 +++++++ packages/api/src/agents/execution.ts | 8 ++- packages/api/src/agents/handlers.ts | 1 + packages/api/src/agents/skillFiles.spec.ts | 13 ++++ packages/api/src/agents/skillFiles.ts | 5 +- packages/api/src/auth/codeapi.spec.ts | 15 ++++ packages/api/src/auth/codeapi.ts | 4 +- packages/api/src/code/http.spec.ts | 53 ++++++++++++++ packages/api/src/code/http.ts | 5 +- 17 files changed, 295 insertions(+), 17 deletions(-) create mode 100644 api/server/middleware/limiters/code.js create mode 100644 api/server/middleware/limiters/code.spec.js diff --git a/api/server/middleware/limiters/code.js b/api/server/middleware/limiters/code.js new file mode 100644 index 00000000000..6b2d081fe94 --- /dev/null +++ b/api/server/middleware/limiters/code.js @@ -0,0 +1,36 @@ +const rateLimit = require('express-rate-limit'); +const { limiterCache } = require('@librechat/api'); + +let configuredCodeEnvironmentPairingLimiter; + +const codeEnvironmentPairingLimiter = (req, res, next) => { + if (configuredCodeEnvironmentPairingLimiter == null) { + const max = Number(process.env.CODE_ENVIRONMENT_PAIRING_USER_MAX ?? 5); + const windowInMinutes = Number(process.env.CODE_ENVIRONMENT_PAIRING_USER_WINDOW ?? 60); + + configuredCodeEnvironmentPairingLimiter = rateLimit({ + windowMs: windowInMinutes * 60 * 1000, + max, + handler: (limitedReq, limitedRes) => { + const resetAt = limitedReq.rateLimit?.resetTime?.getTime?.(); + const retryAfterSeconds = Number.isFinite(resetAt) + ? Math.max(1, Math.ceil((resetAt - Date.now()) / 1000)) + : Math.max(1, Math.ceil(windowInMinutes * 60)); + limitedRes.set('Retry-After', String(retryAfterSeconds)); + return limitedRes.status(429).json({ + error: { + code: 'code_environment_pairing_rate_limited', + message: 'Code environment pairing rate limit exceeded.', + type: 'rate_limit_error', + }, + }); + }, + keyGenerator: (limitedReq) => String(limitedReq.user.id), + store: limiterCache('code_environment_pairing_user_limiter'), + }); + } + + return configuredCodeEnvironmentPairingLimiter(req, res, next); +}; + +module.exports = { codeEnvironmentPairingLimiter }; diff --git a/api/server/middleware/limiters/code.spec.js b/api/server/middleware/limiters/code.spec.js new file mode 100644 index 00000000000..bea8a40adef --- /dev/null +++ b/api/server/middleware/limiters/code.spec.js @@ -0,0 +1,70 @@ +const express = require('express'); +const request = require('supertest'); + +const mockLimiter = jest.fn((_req, _res, next) => next()); +const mockRateLimit = jest.fn(() => mockLimiter); + +jest.mock('express-rate-limit', () => mockRateLimit); +jest.mock('@librechat/api', () => ({ + limiterCache: jest.fn(() => ({})), +})); + +describe('code environment pairing rate limiter', () => { + let originalEnv; + + beforeEach(() => { + originalEnv = { ...process.env }; + jest.resetModules(); + }); + + afterEach(() => { + process.env = originalEnv; + jest.clearAllMocks(); + }); + + it('uses a configurable per-user bucket', () => { + process.env.CODE_ENVIRONMENT_PAIRING_USER_MAX = '3'; + process.env.CODE_ENVIRONMENT_PAIRING_USER_WINDOW = '15'; + const { codeEnvironmentPairingLimiter } = require('./code'); + const { limiterCache } = require('@librechat/api'); + const req = { user: { id: 'user-1' } }; + const next = jest.fn(); + + codeEnvironmentPairingLimiter(req, {}, next); + + expect(mockRateLimit).toHaveBeenLastCalledWith( + expect.objectContaining({ + max: 3, + windowMs: 900_000, + keyGenerator: expect.any(Function), + }), + ); + expect(mockRateLimit.mock.calls.at(-1)[0].keyGenerator(req)).toBe('user-1'); + expect(limiterCache).toHaveBeenCalledWith('code_environment_pairing_user_limiter'); + expect(mockLimiter).toHaveBeenCalledWith(req, {}, next); + }); + + it('returns an actionable JSON response when the limit is exceeded', async () => { + const { codeEnvironmentPairingLimiter } = require('./code'); + + codeEnvironmentPairingLimiter({ user: { id: 'user-1' } }, {}, jest.fn()); + const options = mockRateLimit.mock.calls.at(-1)[0]; + const app = express(); + app.use((req, _res, next) => { + req.rateLimit = { resetTime: new Date(Date.now() + 30_000) }; + next(); + }); + app.post('/api/code-environments/pairings', options.handler); + const response = await request(app).post('/api/code-environments/pairings'); + + expect(response.status).toBe(429); + expect(response.headers['retry-after']).toBe('30'); + expect(response.body).toEqual({ + error: { + code: 'code_environment_pairing_rate_limited', + message: 'Code environment pairing rate limit exceeded.', + type: 'rate_limit_error', + }, + }); + }); +}); diff --git a/api/server/routes/code-environments.js b/api/server/routes/code-environments.js index f0f84edaeee..9d66bdff4eb 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -2,6 +2,7 @@ const express = require('express'); const { createCodeEnvironmentHttpHandlers } = 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'); @@ -20,7 +21,9 @@ const requireCodeEnvironmentManage = requireCapability(SystemCapabilities.MANAGE router.use(requireJwtAuth); router.get('/', (req, res, next) => getHandlers().list(req, res, next)); -router.post('/pairings', (req, res, next) => getHandlers().pair(req, res, next)); +router.post('/pairings', codeEnvironmentPairingLimiter, (req, res, next) => + getHandlers().pair(req, res, next), +); router.post('/', requireCodeEnvironmentManage, (req, res, next) => getHandlers().register(req, res, next), ); diff --git a/api/server/routes/code-environments.test.js b/api/server/routes/code-environments.test.js index ed0d1f28fba..d483920040a 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -11,6 +11,10 @@ const mockRequireCapability = jest.fn((capability) => (req, _res, next) => { middlewareCalls.push(capability); next(); }); +const mockCodeEnvironmentPairingLimiter = jest.fn((_req, _res, next) => { + middlewareCalls.push('pairing-limit'); + next(); +}); const mockRegistry = {}; const mockGetCodeEnvironmentRegistry = jest.fn(() => mockRegistry); const mockHandlers = { @@ -31,6 +35,9 @@ jest.mock('@librechat/api', () => ({ jest.mock('~/server/middleware/roles/capabilities', () => ({ requireCapability: mockRequireCapability, })); +jest.mock('~/server/middleware/limiters/code', () => ({ + codeEnvironmentPairingLimiter: mockCodeEnvironmentPairingLimiter, +})); jest.mock('~/server/middleware', () => ({ requireJwtAuth: mockRequireJwtAuth })); jest.mock('~/server/services/Config', () => ({ @@ -82,7 +89,7 @@ describe('code environment routes', () => { .send({ name: 'Personal VM', controlPlaneId: 'shared-code-api' }) .expect(201); - expect(middlewareCalls).toEqual(['jwt']); + expect(middlewareCalls).toEqual(['jwt', 'pairing-limit']); expect(mockHandlers.pair).toHaveBeenCalledTimes(1); }); }); diff --git a/api/server/services/Files/Code/crud.js b/api/server/services/Files/Code/crud.js index 76945e06335..c7e63540d21 100644 --- a/api/server/services/Files/Code/crud.js +++ b/api/server/services/Files/Code/crud.js @@ -14,7 +14,6 @@ const { getCodeExecutionBaseUrl, createCodeExecutionRouteKey, codeExecutionHeaders, - CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); const axios = createAxiosInstance(); @@ -195,6 +194,7 @@ async function deleteCodeEnvFile(req, file) { * @param {number} [params.version] - Required when `kind === 'skill'`; absent otherwise. * @param {string} [params.codeApiBaseUrl] - Trusted per-agent Code API endpoint. * @param {'default'|'stateful'} [params.executionProfile] - Trusted execution profile. + * @param {string} [params.bridgeWorkerId] - Trusted worker selected for this execution. * @returns {Promise<{ storage_session_id: string; file_id: string }>} * The codeapi storage location of the uploaded file. * @throws {Error} If there's an error during the upload process. @@ -208,6 +208,7 @@ async function uploadCodeEnvFile({ version, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }) { try { const form = new FormData(); @@ -215,7 +216,7 @@ async function uploadCodeEnvFile({ appendCodeEnvFile(form, stream, filename); const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); /** @type {import('axios').AxiosRequestConfig} */ const options = { headers: { @@ -224,7 +225,7 @@ async function uploadCodeEnvFile({ 'User-Agent': 'LibreChat/1.0', 'User-Id': req.user.id, ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -276,6 +277,7 @@ async function uploadCodeEnvFile({ * `inherited: true`, never as a generated artifact. * @param {string} [params.codeApiBaseUrl] - Trusted per-agent Code API endpoint. * @param {'default'|'stateful'} [params.executionProfile] - Trusted execution profile. + * @param {string} [params.bridgeWorkerId] - Trusted worker selected for this execution. * @returns {Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }> }>} * @throws {Error} If the batch upload fails entirely. */ @@ -288,6 +290,7 @@ async function batchUploadCodeEnvFiles({ read_only = false, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }) { const form = new FormData(); appendCodeEnvFileIdentity(form, { kind, id, version }); @@ -299,7 +302,7 @@ async function batchUploadCodeEnvFiles({ } const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); /** @type {import('axios').AxiosRequestConfig} */ const options = { headers: { @@ -308,7 +311,7 @@ async function batchUploadCodeEnvFiles({ 'User-Agent': 'LibreChat/1.0', 'User-Id': req.user.id, ...authHeaders, - ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, diff --git a/api/server/services/Files/Code/crud.spec.js b/api/server/services/Files/Code/crud.spec.js index 7af9c52b682..4e5c1242dbe 100644 --- a/api/server/services/Files/Code/crud.spec.js +++ b/api/server/services/Files/Code/crud.spec.js @@ -70,7 +70,12 @@ const { getCodeApiAuthHeaders, getCodeExecutionBaseUrl, } = require('@librechat/api'); -const { deleteCodeEnvFile, getCodeOutputDownloadStream, uploadCodeEnvFile } = require('./crud'); +const { + deleteCodeEnvFile, + getCodeOutputDownloadStream, + uploadCodeEnvFile, + batchUploadCodeEnvFiles, +} = require('./crud'); describe('Code CRUD', () => { beforeEach(() => { @@ -438,7 +443,7 @@ describe('Code CRUD', () => { await uploadCodeEnvFile(baseUploadParams); const callConfig = mockAxios.post.mock.calls[0][2]; - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(baseUploadParams.req); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(baseUploadParams.req, undefined); expect(callConfig.headers.Authorization).toBe('Bearer codeapi-token'); }); @@ -455,11 +460,14 @@ describe('Code CRUD', () => { ...baseUploadParams, codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }); const [url, , callConfig] = mockAxios.post.mock.calls[0]; expect(url).toBe('https://stateful-code.example.com/upload'); expect(callConfig.headers['X-CodeAPI-Expected-Profile']).toBe('stateful'); + expect(callConfig.headers['X-LibreChat-Code-Worker-ID']).toBe('personal-worker-1'); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(baseUploadParams.req, 'personal-worker-1'); }); /* Phase C / option α (codeapi #1455): the upload wire carries the @@ -560,4 +568,35 @@ describe('Code CRUD', () => { await expect(uploadCodeEnvFile(baseUploadParams)).rejects.toThrow(); }); }); + + describe('batchUploadCodeEnvFiles', () => { + it('routes batch uploads through the selected bridge worker', async () => { + const req = { user: { id: 'user-123' } }; + mockAxios.post.mockResolvedValue({ + data: { + message: 'success', + storage_session_id: 'sess-1', + files: [{ status: 'success', fileId: 'fid-1', filename: 'data.csv' }], + succeeded: 1, + failed: 0, + }, + }); + + await batchUploadCodeEnvFiles({ + req, + files: [{ stream: Readable.from(['file-content']), filename: 'data.csv' }], + kind: 'user', + id: 'user-123', + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', + }); + + const [url, , callConfig] = mockAxios.post.mock.calls[0]; + expect(url).toBe('https://stateful-code.example.com/upload/batch'); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req, 'personal-worker-1'); + expect(callConfig.headers['X-CodeAPI-Expected-Profile']).toBe('stateful'); + expect(callConfig.headers['X-LibreChat-Code-Worker-ID']).toBe('personal-worker-1'); + }); + }); }); diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 7a679c92bd6..8da1ea2ba57 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -508,6 +508,7 @@ const runPreviewFinalize = ({ finalize, fileId, previewRevision, onResolved }) = * @param {string} [params.codeApiBaseUrl] - Trusted per-agent Code API endpoint. * @param {'default'|'stateful'} [params.executionProfile] - Trusted execution profile. * @param {string} [params.executionRouteKey] - Trusted deployment-local route identity. + * @param {string} [params.bridgeWorkerId] - Trusted bridge worker selected for this execution. * @param {Buffer} [params.preparedBuffer] - Bytes downloaded during a * no-write content inspection preflight. * @param {boolean} [params.downloadFallback] - Return the bounded download @@ -978,7 +979,7 @@ function checkIfActive(dateString) { * into codeapi storage. Carries kind/id/storage_session_id/file_id; * codeapi resolves the sessionKey from the request's auth context. * @param {ServerRequest} [req] - Current authenticated request, used to mint Code API auth. - * @param {{baseUrl?: string, executionProfile?: 'default'|'stateful'}} [route] + * @param {{baseUrl?: string, executionProfile?: 'default'|'stateful', bridgeWorkerId?: string}} [route] * Trusted host-selected Code API route. * * @returns {Promise} @@ -1256,6 +1257,7 @@ const primeFiles = async (options) => { ...(sourceRef.kind === 'skill' ? { version: sourceRef.version } : {}), codeApiBaseUrl, executionProfile, + bridgeWorkerId, }); /** diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index e0984fb9bb1..5673b5f920d 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -2411,6 +2411,7 @@ describe('Code Process', () => { agentId: 'agent-id', codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }); expect(mockAxios).not.toHaveBeenCalled(); @@ -2418,6 +2419,7 @@ describe('Code Process', () => { expect.objectContaining({ codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }), ); expect(updateFile).toHaveBeenCalledWith( diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts index 517bc273eb5..ca8c3320977 100644 --- a/packages/api/src/agents/execution.spec.ts +++ b/packages/api/src/agents/execution.spec.ts @@ -260,6 +260,32 @@ describe('resolveCodeExecutionContext', () => { expect(first.codeSessionKey).not.toBe(replacement.codeSessionKey); }); + it('namespaces replacement workers independently under a stable environment route', () => { + const environment = (workerId: string) => ({ + id: 'personal-vm', + name: 'Personal VM', + type: 'attached' as const, + baseURL: 'https://bridge.example/v1', + default: true, + owner: 'principal' as const, + workerId, + }); + const first = resolveCodeExecutionContext({ + statefulSessions: true, + environments: [environment('worker-a')], + userId: 'user-1', + }); + const replacement = resolveCodeExecutionContext({ + statefulSessions: true, + environments: [environment('worker-b')], + userId: 'user-1', + }); + + expect(first.runtimeSessionHint).toBe(replacement.runtimeSessionHint); + expect(first.executionRouteKey).not.toBe(replacement.executionRouteKey); + expect(first.codeSessionKey).not.toBe(replacement.codeSessionKey); + }); + it('uses the operator-selected default environment when the agent has no override', () => { const context = resolveCodeExecutionContext({ statefulSessions: true, diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index b4f69520656..c4744eac6fb 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -27,12 +27,16 @@ export interface CodeExecutionContext { export function createCodeExecutionRouteKey( profile: CodeExecutionProfile, - environment?: Pick, + environment?: Pick, ): string { if (profile === 'default' || environment == null) { return profile; } - const identity = JSON.stringify([environment.id, environment.baseURL.trim().replace(/\/+$/, '')]); + const identity = JSON.stringify([ + environment.id, + environment.baseURL.trim().replace(/\/+$/, ''), + environment.workerId ?? environment.pairing?.workerId ?? '', + ]); return `stateful:${createHash('sha256').update(identity).digest('hex').slice(0, 32)}`; } diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 5838fa668c8..1263df7af9a 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -398,6 +398,7 @@ export interface ToolExecuteOptions { route?: { baseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; }, ) => Promise; /** 23-hour freshness check */ diff --git a/packages/api/src/agents/skillFiles.spec.ts b/packages/api/src/agents/skillFiles.spec.ts index 429da21f329..e87a5459596 100644 --- a/packages/api/src/agents/skillFiles.spec.ts +++ b/packages/api/src/agents/skillFiles.spec.ts @@ -757,10 +757,23 @@ describe('primeSkillFiles — resource identity propagation', () => { bytes: 256, }, ], + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', + }, }); const result = await primeSkillFiles(deps); + expect(deps.batchUploadCodeEnvFiles).toHaveBeenCalledWith( + expect.objectContaining({ + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', + }), + ); + expect(result?.files).toEqual([ { id: 'file-fresh', diff --git a/packages/api/src/agents/skillFiles.ts b/packages/api/src/agents/skillFiles.ts index 5eb0645d983..b43a242897b 100644 --- a/packages/api/src/agents/skillFiles.ts +++ b/packages/api/src/agents/skillFiles.ts @@ -75,6 +75,7 @@ export interface PrimeSkillFilesParams { read_only?: boolean; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; }) => Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }>; @@ -86,12 +87,13 @@ export interface PrimeSkillFilesParams { route?: { baseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; }, ) => Promise; /** Trusted Code API route selected for the executing agent. */ codeExecutionContext?: Pick< CodeExecutionContext, - 'baseUrl' | 'executionProfile' | 'executionRouteKey' + 'baseUrl' | 'executionProfile' | 'executionRouteKey' | 'bridgeWorkerId' >; /** 23-hour freshness check */ checkIfActive?: (dateString: string) => boolean; @@ -517,6 +519,7 @@ async function executePrimeSkillFiles( read_only: true, codeApiBaseUrl: codeExecutionContext?.baseUrl, executionProfile: codeExecutionContext?.executionProfile, + bridgeWorkerId: codeExecutionContext?.bridgeWorkerId, }); return { filesToUpload, result }; }, `skill "${skill.name}"`), diff --git a/packages/api/src/auth/codeapi.spec.ts b/packages/api/src/auth/codeapi.spec.ts index c0e344a3d21..76769571543 100644 --- a/packages/api/src/auth/codeapi.spec.ts +++ b/packages/api/src/auth/codeapi.spec.ts @@ -221,6 +221,21 @@ describe('Code API JWT minting', () => { expect(second).not.toBe(first); }); + it('does not collide cache entries for colon-containing plans and worker IDs', async () => { + const first = await mintCodeApiToken(baseRequest({ subscription: { planId: 'basic:a' } }), 'b'); + const second = await mintCodeApiToken( + baseRequest({ subscription: { planId: 'basic' } }), + 'a:b', + ); + + expect(decodeToken(first).claims).toMatchObject({ plan_id: 'basic:a', code_worker_id: 'b' }); + expect(decodeToken(second).claims).toMatchObject({ + plan_id: 'basic', + code_worker_id: 'a:b', + }); + expect(second).not.toBe(first); + }); + it('uses the single-tenant namespace when tenant context is absent outside strict mode', async () => { mockGetTenantId.mockReturnValue(undefined); diff --git a/packages/api/src/auth/codeapi.ts b/packages/api/src/auth/codeapi.ts index d655bcee001..76d8664afb5 100644 --- a/packages/api/src/auth/codeapi.ts +++ b/packages/api/src/auth/codeapi.ts @@ -326,7 +326,7 @@ function signJwt(config: SigningConfig, claims: CodeApiClaims): string { } function cacheKey(config: SigningConfig, claims: CodeApiClaims): string { - return [ + return JSON.stringify([ config.alg, config.kid, claims.sub, @@ -339,7 +339,7 @@ function cacheKey(config: SigningConfig, claims: CodeApiClaims): string { claims.plan_id ?? '', claims.code_worker_id ?? '', claims.auth_context_hash, - ].join(':'); + ]); } function pruneTokenCache(now: number): void { diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 13cba36308b..007efaffffa 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -155,6 +155,59 @@ describe('code environment HTTP handlers', () => { expect(fetchImpl).not.toHaveBeenCalled(); }); + test('does not pair through a control plane removed from the caller effective config', async () => { + const fetchImpl = jest.fn(); + const baseConfig = { + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached', + baseURL: 'https://code.librechat.example', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig; + const effectiveConfig = { + endpoints: { [EModelEndpoint.agents]: { statefulCodeSessions: { environments: [] } } }, + } as AppConfig; + const getAppConfig = jest.fn(async (options) => + options.baseOnly === true ? baseConfig : effectiveConfig, + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig, + registry: { register: jest.fn(), listAccessible: jest.fn() }, + readSecret: jest.fn(() => 'administrator-token'), + principalAuthEnabled: jest.fn(() => true), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER', tenantId: 'tenant-1' }, + body: { name: 'Personal VM', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(404); + expect(getAppConfig).toHaveBeenCalledWith({ + role: 'USER', + userId: '68b2f0c498f24c1e78fa0001', + idOnTheSource: undefined, + tenantId: 'tenant-1', + }); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + test('registers against an operator-configured control plane and ignores client URLs', async () => { const register = jest.fn().mockResolvedValue({ resourceId: '68b2f0c498f24c1e78fa0111', diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 39a4b8de24c..2c9926347a1 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -13,6 +13,7 @@ import type { CodeBridgeFetch } from './bridge'; import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; import { getCodeApiTenantId, isCodeApiJwtAuthEnabled } from '~/auth/codeapi'; +import { getAppConfigOptionsFromUser } from '~/app/service'; type Registry = { register: (params: { @@ -141,7 +142,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) /** Control-plane destinations are deployment policy. Client-provided URLs * are deliberately ignored to prevent an authenticated SSRF primitive. */ - const appConfig = await deps.getAppConfig({ baseOnly: true }); + const appConfig = await deps.getAppConfig(getAppConfigOptionsFromUser(req.user)); const controlPlane = configuredControlPlane(appConfig, controlPlaneId); if (controlPlane == null) { return res.status(404).json({ error: 'Code control plane was not found' }); @@ -206,7 +207,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) }); } - const appConfig = await deps.getAppConfig({ baseOnly: true }); + const appConfig = await deps.getAppConfig(getAppConfigOptionsFromUser(req.user)); const controlPlane = configuredPrincipalControlPlane(appConfig, controlPlaneId); if (controlPlane == null) { return res.status(404).json({ error: 'Principal code control plane was not found' }); From a9e255c26011f87c66458d3d2206c4e9cc6cbe6b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 19:58:59 -0400 Subject: [PATCH 09/29] fix: fail closed for principal code pairing --- packages/api/src/app/service.spec.ts | 30 ++++++ packages/api/src/app/service.ts | 7 +- packages/api/src/auth/codeapi.spec.ts | 14 ++- packages/api/src/auth/codeapi.ts | 16 +++ packages/api/src/code/config.spec.ts | 47 +++++++++ packages/api/src/code/config.ts | 32 +++++- packages/api/src/code/http.spec.ts | 138 ++++++++++++++++++++++++++ packages/api/src/code/http.ts | 48 +++++++-- 8 files changed, 320 insertions(+), 12 deletions(-) diff --git a/packages/api/src/app/service.spec.ts b/packages/api/src/app/service.spec.ts index 4283c633c54..9443e4d8893 100644 --- a/packages/api/src/app/service.spec.ts +++ b/packages/api/src/app/service.spec.ts @@ -514,6 +514,36 @@ describe('createAppConfigService', () => { ); }); + it('propagates principal resolution failures for fail-closed callers', async () => { + const error = new Error('principal authorization unavailable'); + const deps = createDeps({ getUserPrincipals: jest.fn().mockRejectedValue(error) }); + const { getAppConfig } = createAppConfigService(deps); + + await expect(getAppConfig({ role: 'USER', userId: 'uid1', failClosed: true })).rejects.toBe( + error, + ); + }); + + it('propagates override resolution failures for fail-closed callers', async () => { + const error = new Error('override authorization unavailable'); + const deps = createDeps({ getApplicableConfigs: jest.fn().mockRejectedValue(error) }); + const { getAppConfig } = createAppConfigService(deps); + + await expect(getAppConfig({ role: 'USER', userId: 'uid1', failClosed: true })).rejects.toBe( + error, + ); + }); + + it('propagates principal augmentation failures for fail-closed callers', async () => { + const error = new Error('environment authorization unavailable'); + const deps = createDeps({ augmentConfig: jest.fn().mockRejectedValue(error) }); + const { getAppConfig } = createAppConfigService(deps); + + await expect(getAppConfig({ role: 'USER', userId: 'uid1', failClosed: true })).rejects.toBe( + error, + ); + }); + it('passes local identity through to getUserPrincipals when provided', async () => { const deps = createDeps(); const { getAppConfig } = createAppConfigService(deps); diff --git a/packages/api/src/app/service.ts b/packages/api/src/app/service.ts index 2f5d5043c87..49c59ee2561 100644 --- a/packages/api/src/app/service.ts +++ b/packages/api/src/app/service.ts @@ -78,6 +78,8 @@ export interface GetAppConfigOptions { refresh?: boolean; /** When true, return only the YAML-derived base config — no DB override queries. */ baseOnly?: boolean; + /** Propagate principal, override, and augmentation failures for security-sensitive callers. */ + failClosed?: boolean; } export interface AppConfigUserLike { @@ -214,7 +216,7 @@ export function createAppConfigService(deps: AppConfigServiceDeps): { * Use this for startup, auth strategies, and other pre-tenant code paths. */ async function getAppConfig(options: GetAppConfigOptions = {}): Promise { - const { role, userId, idOnTheSource, tenantId, refresh, baseOnly } = options; + const { role, userId, idOnTheSource, tenantId, refresh, baseOnly, failClosed } = options; const baseConfig = await ensureBaseConfig(refresh); @@ -224,6 +226,7 @@ export function createAppConfigService(deps: AppConfigServiceDeps): { const principals = await buildPrincipals(role, userId, idOnTheSource).catch( (error: unknown) => { + if (failClosed) throw error; logger.error('[getAppConfig] Error building principals, falling back to base:', error); return null; }, @@ -254,6 +257,7 @@ export function createAppConfigService(deps: AppConfigServiceDeps): { try { return await augmentConfig({ appConfig, baseConfig, principals, options }); } catch (error) { + if (failClosed) throw error; logger.error('[getAppConfig] Error augmenting principal config:', error); return appConfig; } @@ -274,6 +278,7 @@ export function createAppConfigService(deps: AppConfigServiceDeps): { merged = materializeConfigModelSpecs(mergeConfigOverrides(baseConfig, configs)); } } catch (error) { + if (failClosed) throw error; logger.error('[getAppConfig] Error resolving config overrides, falling back to base:', error); return baseConfig; } diff --git a/packages/api/src/auth/codeapi.spec.ts b/packages/api/src/auth/codeapi.spec.ts index 76769571543..2d55c6ad192 100644 --- a/packages/api/src/auth/codeapi.spec.ts +++ b/packages/api/src/auth/codeapi.spec.ts @@ -1,7 +1,7 @@ import { createHash, generateKeyPairSync, verify as cryptoVerify } from 'crypto'; import type { KeyObject } from 'crypto'; import type { ServerRequest } from '~/types'; -import { getCodeApiAuthHeaders, mintCodeApiToken } from './codeapi'; +import { assertCodeApiJwtSigningReady, getCodeApiAuthHeaders, mintCodeApiToken } from './codeapi'; jest.mock( '@librechat/data-schemas', @@ -170,6 +170,18 @@ describe('Code API JWT minting', () => { expect(decoded.claims).not.toHaveProperty('openid_token'); }); + it('validates that the configured signing key is usable by the selected algorithm', () => { + expect(() => assertCodeApiJwtSigningReady()).not.toThrow(); + + const rsaKeyPair = generateKeyPairSync('rsa', { modulusLength: 2048 }); + process.env.CODEAPI_JWT_PRIVATE_JWK_JSON = JSON.stringify( + rsaKeyPair.privateKey.export({ format: 'jwk' }), + ); + process.env.CODEAPI_JWT_ALGORITHM = 'EdDSA'; + + expect(() => assertCodeApiJwtSigningReady()).toThrow(); + }); + it('marks OpenID reuse callers without forwarding upstream credentials', async () => { process.env.OPENID_REUSE_TOKENS = 'true'; const req = baseRequest({ diff --git a/packages/api/src/auth/codeapi.ts b/packages/api/src/auth/codeapi.ts index 76d8664afb5..1d2688a1af5 100644 --- a/packages/api/src/auth/codeapi.ts +++ b/packages/api/src/auth/codeapi.ts @@ -168,6 +168,22 @@ function getSigningConfig(): SigningConfig { return signingConfigCache; } +export function assertCodeApiJwtSigningReady(): void { + const config = getSigningConfig(); + const keyType = config.key.asymmetricKeyType; + if (config.alg === 'EdDSA' && keyType !== 'ed25519' && keyType !== 'ed448') { + throw new Error(`Code API JWT algorithm EdDSA requires an EdDSA key, received ${keyType}`); + } + if (config.alg === 'RS256' && keyType !== 'rsa') { + throw new Error(`Code API JWT algorithm RS256 requires an RSA key, received ${keyType}`); + } + cryptoSign( + config.alg === 'RS256' ? 'RSA-SHA256' : null, + Buffer.from('librechat-codeapi-signing-readiness'), + config.key, + ); +} + function stringifyClaimValue(value: unknown): string | undefined { if (typeof value === 'string' && value.trim() !== '') { return value; diff --git a/packages/api/src/code/config.spec.ts b/packages/api/src/code/config.spec.ts index 6f8bf71b0a0..b0981334159 100644 --- a/packages/api/src/code/config.spec.ts +++ b/packages/api/src/code/config.spec.ts @@ -333,4 +333,51 @@ describe('mergeAccessibleCodeEnvironments', () => { deploymentEnvironment, ]); }); + + test('defaults to an executable principal environment after a pairing-only control plane', async () => { + const pairingOnly = { + id: 'self-service', + name: 'Self-service', + type: 'attached' as const, + baseURL: 'https://code.example', + owner: 'deployment' as const, + default: true, + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }; + const appConfig = { + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { environments: [pairingOnly] }, + }, + }, + } as unknown as AppConfig; + + const result = await mergeAccessibleCodeEnvironments({ + appConfig, + deploymentConfig: appConfig, + actor: { userId: '68b2f0c498f24c1e78fa0001', role: 'USER', idOnTheSource: null }, + registry: { + listRegisteredIds: jest.fn().mockResolvedValue(['personal-vm']), + listAccessibleConfigurations: jest.fn().mockResolvedValue([ + { + id: 'personal-vm', + name: 'Personal VM', + type: 'attached', + baseURL: 'https://persisted.example', + controlPlaneId: 'self-service', + owner: 'principal', + workerId: 'personal-worker', + }, + ]), + }, + }); + + const environments = result.endpoints?.agents?.statefulCodeSessions?.environments; + expect(environments?.find((environment) => environment.id === 'self-service')?.default).toBe( + false, + ); + expect(environments?.find((environment) => environment.id === 'personal-vm')?.default).toBe( + true, + ); + }); }); diff --git a/packages/api/src/code/config.ts b/packages/api/src/code/config.ts index 38a8e4088f8..3273bc19bf5 100644 --- a/packages/api/src/code/config.ts +++ b/packages/api/src/code/config.ts @@ -13,6 +13,19 @@ type ConfigurationRegistry = { listRegisteredIds: () => Promise; }; +type StatefulCodeConfig = NonNullable< + NonNullable[EModelEndpoint.agents] +>['statefulCodeSessions']; +type CodeEnvironmentConfig = NonNullable['environments']>[number]; + +function isExecutableCodeEnvironment(environment: CodeEnvironmentConfig): boolean { + return !( + environment.pairing?.allowPrincipalWorkers === true && + environment.pairing.workerId == null && + environment.workerId == null + ); +} + function retainDeploymentCodeEnvironments( appConfig: AppConfig, deploymentConfig: AppConfig, @@ -115,14 +128,27 @@ export async function mergeAccessibleCodeEnvironments({ ) { return appConfig; } - const mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments]; + let mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments].map( + (environment) => + !isExecutableCodeEnvironment(environment) && environment.default === true + ? { ...environment, default: false as const } + : environment, + ); if ( mergedEnvironments.length > 0 && !mergedEnvironments.some( - (environment) => 'default' in environment && environment.default === true, + (environment) => + isExecutableCodeEnvironment(environment) && + 'default' in environment && + environment.default === true, ) ) { - mergedEnvironments[0] = { ...mergedEnvironments[0], default: true }; + const defaultIndex = mergedEnvironments.findIndex(isExecutableCodeEnvironment); + if (defaultIndex >= 0) { + mergedEnvironments = mergedEnvironments.map((environment, index) => + index === defaultIndex ? { ...environment, default: true as const } : environment, + ); + } } return { diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 007efaffffa..2871bbadb00 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -82,6 +82,7 @@ describe('code environment HTTP handlers', () => { readSecret: jest.fn(() => 'administrator-token'), resolveTenantId: jest.fn(() => 'tenant-1'), principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), fetchImpl, }); const req = { @@ -186,6 +187,7 @@ describe('code environment HTTP handlers', () => { registry: { register: jest.fn(), listAccessible: jest.fn() }, readSecret: jest.fn(() => 'administrator-token'), principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), fetchImpl, }); const res = response(); @@ -204,10 +206,145 @@ describe('code environment HTTP handlers', () => { userId: '68b2f0c498f24c1e78fa0001', idOnTheSource: undefined, tenantId: 'tenant-1', + failClosed: true, }); expect(fetchImpl).not.toHaveBeenCalled(); }); + test('fails closed when effective pairing policy cannot be loaded', async () => { + const fetchImpl = jest.fn(); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn(async (options) => { + if (options.baseOnly === true) return {} as AppConfig; + throw new Error('authorization unavailable'); + }), + registry: { register: jest.fn(), listAccessible: jest.fn() }, + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(503); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + + test('authorizes pairing effectively but resolves destinations and secrets from deployment config', async () => { + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: jest.fn().mockResolvedValue({ + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }), + }); + const deploymentConfig = { + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached', + baseURL: 'https://code.librechat.example', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'CODE_ADMIN_TOKEN' }, + }, + ], + }, + }, + }, + } as AppConfig; + const effectiveConfig = { + endpoints: { + [EModelEndpoint.agents]: { + statefulCodeSessions: { + environments: [ + { + id: 'self-service', + name: 'Override attempt', + type: 'attached', + baseURL: 'https://attacker.example', + owner: 'deployment', + pairing: { allowPrincipalWorkers: true, tokenEnv: 'DATABASE_URL' }, + }, + ], + }, + }, + }, + } as AppConfig; + const readSecret = jest.fn((name) => + name === 'CODE_ADMIN_TOKEN' ? 'administrator-token' : 'database-secret', + ); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn(async (options) => + options.baseOnly === true ? deploymentConfig : effectiveConfig, + ), + registry: { + register: jest.fn().mockResolvedValue({ id: 'code-generated' }), + listAccessible: jest.fn(), + }, + createEnvironmentId: () => 'code-generated', + readSecret, + resolveTenantId: jest.fn(() => 'tenant-1'), + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(201); + expect(readSecret).toHaveBeenCalledWith('CODE_ADMIN_TOKEN'); + expect(readSecret).not.toHaveBeenCalledWith('DATABASE_URL'); + expect(fetchImpl).toHaveBeenCalledWith( + 'https://code.librechat.example/bridge/pairings', + expect.any(Object), + ); + }); + + test('validates JWT signing before creating upstream pairing state', async () => { + const fetchImpl = jest.fn(); + const handlers = createCodeEnvironmentHttpHandlers({ + getAppConfig: jest.fn(), + registry: { register: jest.fn(), listAccessible: jest.fn() }, + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(() => { + throw new Error('invalid signing key'); + }), + fetchImpl, + }); + const res = response(); + + await handlers.pair( + { + user: { id: '68b2f0c498f24c1e78fa0001', role: 'USER' }, + body: { name: 'Personal VM', controlPlaneId: 'self-service' }, + } as never, + res as never, + ); + + expect(res.statusCode).toBe(503); + expect(fetchImpl).not.toHaveBeenCalled(); + }); + test('registers against an operator-configured control plane and ignores client URLs', async () => { const register = jest.fn().mockResolvedValue({ resourceId: '68b2f0c498f24c1e78fa0111', @@ -340,6 +477,7 @@ describe('code environment HTTP handlers', () => { registry: { register: jest.fn(), listAccessible: jest.fn() }, readSecret: jest.fn(() => 'administrator-token'), principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), fetchImpl, }); const res = response(); diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 2c9926347a1..c341a31f173 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -10,9 +10,13 @@ import type { import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; import type { CodeBridgeFetch } from './bridge'; +import { + assertCodeApiJwtSigningReady, + getCodeApiTenantId, + isCodeApiJwtAuthEnabled, +} from '~/auth/codeapi'; import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; -import { getCodeApiTenantId, isCodeApiJwtAuthEnabled } from '~/auth/codeapi'; import { getAppConfigOptionsFromUser } from '~/app/service'; type Registry = { @@ -37,6 +41,7 @@ export interface CodeEnvironmentHttpDeps { readSecret?: (name: string) => string | undefined; resolveTenantId?: (req: ServerRequest) => string; principalAuthEnabled?: () => boolean; + principalAuthReady?: () => Promise | void; fetchImpl?: CodeBridgeFetch; } @@ -105,6 +110,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) const readSecret = deps.readSecret ?? readCodeBridgeSecret; const resolveTenantId = deps.resolveTenantId ?? getCodeApiTenantId; const principalAuthEnabled = deps.principalAuthEnabled ?? isCodeApiJwtAuthEnabled; + const principalAuthReady = deps.principalAuthReady ?? assertCodeApiJwtSigningReady; async function list(req: ServerRequest, res: Response): Promise { const principal = actor(req); @@ -142,9 +148,20 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) /** Control-plane destinations are deployment policy. Client-provided URLs * are deliberately ignored to prevent an authenticated SSRF primitive. */ - const appConfig = await deps.getAppConfig(getAppConfigOptionsFromUser(req.user)); - const controlPlane = configuredControlPlane(appConfig, controlPlaneId); - if (controlPlane == null) { + let effectiveConfig: AppConfig; + let deploymentConfig: AppConfig; + try { + [effectiveConfig, deploymentConfig] = await Promise.all([ + deps.getAppConfig({ ...getAppConfigOptionsFromUser(req.user), failClosed: true }), + deps.getAppConfig({ baseOnly: true }), + ]); + } catch (error) { + logger.error('[codeEnvironments] control-plane policy resolution failed:', error); + return res.status(503).json({ error: 'Code environment policy is unavailable' }); + } + const authorizedControlPlane = configuredControlPlane(effectiveConfig, controlPlaneId); + const controlPlane = configuredControlPlane(deploymentConfig, controlPlaneId); + if (authorizedControlPlane == null || controlPlane == null) { return res.status(404).json({ error: 'Code control plane was not found' }); } @@ -206,10 +223,27 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) error: error instanceof Error ? error.message : 'Code environment name is invalid', }); } + try { + await principalAuthReady(); + } catch (error) { + logger.error('[codeEnvironments] Code API JWT signing is unavailable:', error); + return res.status(503).json({ error: 'Principal code worker authentication is unavailable' }); + } - const appConfig = await deps.getAppConfig(getAppConfigOptionsFromUser(req.user)); - const controlPlane = configuredPrincipalControlPlane(appConfig, controlPlaneId); - if (controlPlane == null) { + let effectiveConfig: AppConfig; + let deploymentConfig: AppConfig; + try { + [effectiveConfig, deploymentConfig] = await Promise.all([ + deps.getAppConfig({ ...getAppConfigOptionsFromUser(req.user), failClosed: true }), + deps.getAppConfig({ baseOnly: true }), + ]); + } catch (error) { + logger.error('[codeEnvironments] pairing policy resolution failed:', error); + return res.status(503).json({ error: 'Code environment pairing policy is unavailable' }); + } + const authorizedControlPlane = configuredPrincipalControlPlane(effectiveConfig, controlPlaneId); + const controlPlane = configuredPrincipalControlPlane(deploymentConfig, controlPlaneId); + if (authorizedControlPlane == null || controlPlane == null) { return res.status(404).json({ error: 'Principal code control plane was not found' }); } const tokenEnv = controlPlane.pairing?.tokenEnv; From a3adb77ea3bef8b3cbb04527d8f6ec0bbe2fd176 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 14:26:24 -0400 Subject: [PATCH 10/29] feat: manage personal code environments --- api/server/controllers/UserController.js | 14 ++ api/server/controllers/UserController.spec.js | 2 + .../controllers/__tests__/deleteUser.spec.js | 25 ++- api/server/routes/code-environments.js | 1 + api/server/routes/code-environments.test.js | 8 + .../Nav/Settings/CodeEnvironments.tsx | 182 ++++++++++++++++++ .../src/components/Nav/Settings/registry.tsx | 10 + client/src/data-provider/CodeEnvironments.ts | 41 ++++ client/src/data-provider/index.ts | 1 + client/src/locales/en/translation.json | 14 ++ packages/api/src/code/bridge.ts | 43 +++++ .../src/code/environments.integration.spec.ts | 92 +++++++++ packages/api/src/code/environments.ts | 50 +++++ packages/api/src/code/http.spec.ts | 92 ++++++++- packages/api/src/code/http.ts | 127 +++++++++++- packages/api/src/code/index.ts | 1 + packages/api/src/code/lifecycle.ts | 55 ++++++ packages/data-provider/src/api-endpoints.ts | 4 + packages/data-provider/src/data-service.ts | 17 ++ packages/data-provider/src/keys.ts | 3 + packages/data-provider/src/types.ts | 27 +++ .../src/methods/codeEnvironment.ts | 16 +- .../src/schema/codeEnvironment.ts | 3 + .../data-schemas/src/types/codeEnvironment.ts | 1 + 24 files changed, 815 insertions(+), 14 deletions(-) create mode 100644 client/src/components/Nav/Settings/CodeEnvironments.tsx create mode 100644 client/src/data-provider/CodeEnvironments.ts create mode 100644 packages/api/src/code/lifecycle.ts diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index f4fe17d1e42..4a4e8417798 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, @@ -390,6 +391,19 @@ const deleteUserController = async (req, res) => { } } + const deletionAppConfig = + req.config ?? + (await getAppConfig({ + role: req.user?.role, + userId: req.user?.id, + tenantId: req.user?.tenantId, + })); + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: user.id, + appConfig: deletionAppConfig, + }); + // Block new trigger admissions across replicas while preserving the user // principal so a transient cleanup failure remains retryable. triggerDeletionFence = new Date(); diff --git a/api/server/controllers/UserController.spec.js b/api/server/controllers/UserController.spec.js index 2c6b70cba99..a7c221d40fb 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), diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 23c152b30bd..49c7e138eca 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,7 +211,13 @@ 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(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( + mockBeginAgentTriggerUserDeletion.mock.invocationCallOrder[0], + ); expect(mockBeginAgentTriggerUserDeletion.mock.invocationCallOrder[0]).toBeLessThan( mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0], ); @@ -243,6 +252,20 @@ describe('deleteUserController - 2FA enforcement', () => { ); }); + it('keeps the account when a personal code worker cannot be revoked', 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(500); + expect(mockBeginAgentTriggerUserDeletion).not.toHaveBeenCalled(); + expect(mockDeleteMessages).not.toHaveBeenCalled(); + expect(mockDeleteUserById).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/routes/code-environments.js b/api/server/routes/code-environments.js index 9d66bdff4eb..694e23581f6 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -27,5 +27,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 d483920040a..b17afb1b732 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', () => ({ @@ -92,4 +93,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 00000000000..7eeb7074ae4 --- /dev/null +++ b/client/src/components/Nav/Settings/CodeEnvironments.tsx @@ -0,0 +1,182 @@ +import { 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 command = `librechat-code pair ${pairing.endpoint} '${pairing.code}' --worker-id ${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 || ''; + + 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', + }), + }, + ); + }; + + return ( +
+
+ +

+ {localize('com_ui_code_environments_description')} +

+
+ + {controlPlanes.length > 0 ? ( +
+ setName(event.target.value)} + /> + + +
+ ) : ( +

+ {localize('com_ui_code_environment_unavailable')} +

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

{environment.name}

+

{environment.id}

+
+ + + + + + {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-white 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 2a42db87591..8da07921fa1 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 00000000000..d21fbf5fb5b --- /dev/null +++ b/client/src/data-provider/CodeEnvironments.ts @@ -0,0 +1,41 @@ +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]), + }); +} + +export function useDeleteCodeEnvironmentMutation() { + const queryClient = useQueryClient(); + return useMutation( + [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 090c8b9689c..38bb19a9589 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 791f2e30946..52cff2c39c0 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -2177,6 +2177,20 @@ "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_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/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts index 77b88d48fa8..2009c27d69c 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' | '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,39 @@ 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); + } + } 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 9031b44eb1e..4491ffed6ac 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -4,6 +4,7 @@ import { createMethods, createModels } from '@librechat/data-schemas'; import { AccessRoleIds, PrincipalType, ResourceType } from 'librechat-data-provider'; import { AccessControlService } from '~/acl/accessControlService'; import { createCodeEnvironmentRegistry } from './environments'; +import { revokeUserCodeEnvironmentWorkers } from './lifecycle'; function createSharedCache() { const values = new Map(); @@ -178,6 +179,97 @@ 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('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', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + const fetchImpl = jest.fn().mockResolvedValue({ ok: true }); + + await expect( + revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: ownerId.toString(), + appConfig: { + endpoints: { + agents: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'self-service', + name: 'Self-service', + type: 'attached', + baseURL: 'https://code.example.com/v1', + owner: 'deployment', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_ADMIN_TOKEN', + }, + }, + ], + }, + }, + }, + } 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('removes creator-owned environment records and grants when the user is deleted', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index 0625b93ce23..048362aa753 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -52,6 +52,13 @@ type CachedAccessibleCodeEnvironmentConfiguration = AccessibleCodeEnvironmentCon resourceId: string; }; +export type CodeEnvironmentLifecycleTarget = CodeEnvironmentSummary & { + baseURL: string; + workerId?: string; + controlPlaneId?: string; + workerPrincipal?: CodeEnvironmentRegistration['workerPrincipal']; +}; + 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)?$/; @@ -145,6 +152,11 @@ export function createCodeEnvironmentRegistry( ) => Promise; listRegisteredIds: () => Promise; invalidateAccessibleConfigurations: (tenantId?: string) => Promise; + remove: (params: { + actor: CodeEnvironmentPrincipalContext; + environmentId: string; + beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; + }) => Promise; } { const methods = createMethods(mongoose); const access = new AccessControlService(mongoose); @@ -201,6 +213,7 @@ export function createCodeEnvironmentRegistry( baseURL: environment.baseURL, controlPlaneId: environment.controlPlaneId, workerId: environment.workerId, + controlPlaneId: environment.controlPlaneId, workerPrincipal: environment.workerPrincipal, createdBy: new Types.ObjectId(actor.userId), }); @@ -340,11 +353,48 @@ export function createCodeEnvironmentRegistry( 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; + 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; + + await beforeDelete?.({ + ...toSummary(environment), + baseURL: environment.baseURL, + workerId: environment.workerId, + controlPlaneId: environment.controlPlaneId, + workerPrincipal: environment.workerPrincipal, + }); + const deleted = await methods.deleteCodeEnvironmentById(environment._id); + if (deleted == null) return null; + await access.removeAllPermissions({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + return toSummary(deleted); + } + return { register, listAccessible, listAccessibleConfigurations, listRegisteredIds, invalidateAccessibleConfigurations, + remove, }; } diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 2871bbadb00..ee594d4b905 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -77,7 +77,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'), @@ -122,15 +122,20 @@ describe('code environment HTTP handlers', () => { 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', 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', + }), }); }); @@ -373,7 +378,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 = { @@ -403,6 +408,7 @@ describe('code environment HTTP handlers', () => { baseURL: 'https://code.librechat.example', controlPlaneId: 'shared-code-api', workerId: 'deployment-worker', + controlPlaneId: 'shared-code-api', workerPrincipal: { type: 'deployment', id: 'shared-code-api' }, }, }); @@ -433,7 +439,7 @@ describe('code environment HTTP handlers', () => { }, }, } as AppConfig), - registry: { register, listAccessible: jest.fn() }, + registry: { register, listAccessible: jest.fn(), remove: jest.fn() }, }); const res = response(); @@ -474,7 +480,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(), @@ -494,6 +500,74 @@ 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', + workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, + }); + return { + resourceId: '68b2f0c498f24c1e78fa0111', + id: 'code-generated', + name: 'Personal VM', + type: 'attached' as const, + }; + }); + const fetchImpl = jest.fn().mockResolvedValue({ ok: 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'), @@ -525,7 +599,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 c341a31f173..f72175d5c29 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,13 +11,19 @@ import type { import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; import type { CodeBridgeFetch } from './bridge'; +import { + CodeBridgeLifecycleError, + CodeBridgePairingError, + createCodeBridgePairing, + readCodeBridgeSecret, + revokeCodeBridgeWorker, +} from './bridge'; +import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; 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 = { @@ -25,6 +32,11 @@ type Registry = { environment: CodeEnvironmentRegistration; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; + remove: (params: { + actor: CodeEnvironmentPrincipalContext; + environmentId: string; + beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; + }) => Promise; }; type StatefulCodeConfig = NonNullable< @@ -80,6 +92,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,6 +139,7 @@ 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; @@ -117,8 +152,14 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (principal == null) { return res.status(401).json({ error: 'Authentication required' }); } - const environments = await deps.registry.listAccessible(principal); - return res.status(200).json({ environments }); + const [environments, appConfig] = await Promise.all([ + deps.registry.listAccessible(principal), + deps.getAppConfig({ baseOnly: true }), + ]); + return res.status(200).json({ + environments, + controlPlanes: principalControlPlanes(appConfig), + }); } async function register(req: ServerRequest, res: Response): Promise { @@ -175,6 +216,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) baseURL: controlPlane.baseURL, controlPlaneId, workerId: controlPlane.pairing?.workerId, + controlPlaneId: controlPlane.id, workerPrincipal: { type: 'deployment', id: controlPlane.id }, }, }); @@ -278,6 +320,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) type: 'attached', baseURL: controlPlane.baseURL, workerId, + controlPlaneId: controlPlane.id, workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); @@ -287,9 +330,22 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) workerId: pairing.workerId, code: pairing.code, expiresAt: pairing.expiresAt, + endpoint: controlPlane.baseURL, }, }); } catch (error) { + 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 && @@ -306,5 +362,66 @@ 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' }); + } + const appConfig = await deps.getAppConfig({ baseOnly: true }); + try { + const environment = await deps.registry.remove({ + actor: principal, + environmentId, + beforeDelete: async (target) => { + if (target.workerPrincipal?.type !== 'user' || target.workerId == null) return; + const controlPlane = configuredPrincipalControlPlane( + appConfig, + target.controlPlaneId ?? '', + ); + if (controlPlane == null || controlPlane.baseURL !== target.baseURL) { + throw new CodeEnvironmentLifecycleHttpError( + 409, + 'Code environment control plane is no longer configured', + ); + } + const tokenEnv = controlPlane.pairing?.tokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token) { + throw new CodeEnvironmentLifecycleHttpError( + 503, + 'Code environment revocation is not configured', + ); + } + await revokeCodeBridgeWorker({ + baseURL: controlPlane.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 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 0e405e9f190..787a58b1dee 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 00000000000..66de56d256d --- /dev/null +++ b/packages/api/src/code/lifecycle.ts @@ -0,0 +1,55 @@ +import { createMethods } from '@librechat/data-schemas'; +import { EModelEndpoint } from 'librechat-data-provider'; +import type { AppConfig } from '@librechat/data-schemas'; +import type { CodeBridgeFetch } from './bridge'; +import { readCodeBridgeSecret, revokeCodeBridgeWorker } from './bridge'; + +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 ?? []; + let revoked = 0; + for (const environment of environments) { + if (environment.workerPrincipal?.type !== 'user' || environment.workerId == null) continue; + const controlPlane = controlPlanes.find( + (candidate) => + candidate.id === environment.controlPlaneId && + candidate.owner === 'deployment' && + candidate.type === 'attached' && + candidate.pairing?.allowPrincipalWorkers === true && + candidate.baseURL === environment.baseURL, + ); + if (controlPlane == null) { + throw new Error( + `Code environment control plane is unavailable for ${environment.environmentId}`, + ); + } + const tokenEnv = controlPlane.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: controlPlane.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }); + revoked += 1; + } + return revoked; +} diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index 156b6403f23..8fc14b46b1f 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 1621370b61a..955b65966b8 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -67,6 +67,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 65e26d69dee..db2752bc48b 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', } // Dynamic query keys that require parameters @@ -148,4 +149,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 8f28e0a987c..65cd0c110b7 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -534,6 +534,33 @@ export type TPublicCodeEnvironment = { pairingAvailable?: boolean; }; +export type TCodeEnvironmentSummary = { + resourceId: string; + id: string; + name: string; + type: 'managed' | 'attached'; +}; + +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/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index d9a48993034..7c2f0e11836 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -8,7 +8,7 @@ type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' > & - Pick, 'workerId' | 'workerPrincipal'>; + Pick, 'workerId' | 'controlPlaneId' | 'workerPrincipal'>; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; @@ -19,6 +19,9 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') environmentId: string, ) => Promise; listCodeEnvironmentIds: () => Promise; + findCodeEnvironmentsByCreator: ( + userId: string | Types.ObjectId, + ) => Promise; deleteCodeEnvironmentById: ( id: string | Types.ObjectId, ) => Promise; @@ -52,6 +55,16 @@ 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 { @@ -81,6 +94,7 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') findCodeEnvironmentsByIds, findCodeEnvironmentByEnvironmentId, listCodeEnvironmentIds, + findCodeEnvironmentsByCreator, deleteCodeEnvironmentById, deleteUserCodeEnvironments, }; diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index fb297e5d350..f020a014113 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -48,6 +48,9 @@ const codeEnvironmentSchema: Schema = new Schema Date: Sun, 30 Aug 2026 15:12:14 -0400 Subject: [PATCH 11/29] fix: harden personal code environment lifecycle --- api/server/controllers/UserController.js | 19 +-- .../controllers/__tests__/deleteUser.spec.js | 10 +- api/server/routes/admin/users.js | 11 +- api/server/routes/code-environments.js | 2 + api/server/routes/code-environments.test.js | 1 + .../Nav/Settings/CodeEnvironments.tsx | 149 ++++++++++-------- client/src/data-provider/CodeEnvironments.ts | 5 +- client/src/locales/en/translation.json | 1 + packages/api/src/admin/users.ts | 6 + packages/api/src/code/bridge.ts | 11 +- .../src/code/environments.integration.spec.ts | 26 ++- packages/api/src/code/environments.ts | 77 +++++++-- packages/api/src/code/http.spec.ts | 7 +- packages/api/src/code/http.ts | 59 +++++-- packages/api/src/code/lifecycle.ts | 66 ++++---- packages/data-provider/src/types.ts | 1 + .../src/methods/codeEnvironment.ts | 5 +- .../src/schema/codeEnvironment.ts | 3 + .../data-schemas/src/types/codeEnvironment.ts | 1 + 19 files changed, 300 insertions(+), 160 deletions(-) diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index 4a4e8417798..f21941320e6 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -391,19 +391,6 @@ const deleteUserController = async (req, res) => { } } - const deletionAppConfig = - req.config ?? - (await getAppConfig({ - role: req.user?.role, - userId: req.user?.id, - tenantId: req.user?.tenantId, - })); - await revokeUserCodeEnvironmentWorkers({ - mongoose, - userId: user.id, - appConfig: deletionAppConfig, - }); - // Block new trigger admissions across replicas while preserving the user // principal so a transient cleanup failure remains retryable. triggerDeletionFence = new Date(); @@ -418,6 +405,12 @@ const deleteUserController = async (req, res) => { if (triggerDeletionFence != null) { await prepareAgentTriggerUserPurge(user.id, triggerDeletionFence, user.tenantId); } + const deletionAppConfig = await getAppConfig({ baseOnly: true }); + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: user.id, + appConfig: deletionAppConfig, + }); await drainAgentTriggerDeliveriesForUser(user.id); await subagentThreadTaskStore.cancelAndDrainForOwner(user.id, user.tenantId); // Reversibly suspend the user's schedules under a per-attempt token BEFORE draining. diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 49c7e138eca..06f9a8c49a2 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -215,13 +215,13 @@ describe('deleteUserController - 2FA enforcement', () => { expect.objectContaining({ userId: 'user1' }), ); expect(mockVerifyOTPOrBackupCode).not.toHaveBeenCalled(); - expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( - mockBeginAgentTriggerUserDeletion.mock.invocationCallOrder[0], - ); expect(mockBeginAgentTriggerUserDeletion.mock.invocationCallOrder[0]).toBeLessThan( mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0], ); expect(mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0]).toBeLessThan( + mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0], + ); + expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0], ); expect(mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0]).toBeLessThan( @@ -261,7 +261,9 @@ describe('deleteUserController - 2FA enforcement', () => { await deleteUserController(req, res); expect(res.status).toHaveBeenCalledWith(500); - expect(mockBeginAgentTriggerUserDeletion).not.toHaveBeenCalled(); + expect(mockBeginAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); + expect(mockCancelAgentTriggerUserPurge).toHaveBeenCalledTimes(1); + expect(mockCancelAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); expect(mockDeleteMessages).not.toHaveBeenCalled(); expect(mockDeleteUserById).not.toHaveBeenCalled(); }); diff --git a/api/server/routes/admin/users.js b/api/server/routes/admin/users.js index ab8f5fec8b0..b07c222c779 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 694e23581f6..06da12c5dc3 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -5,6 +5,7 @@ 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(); let handlers; @@ -13,6 +14,7 @@ function getHandlers() { handlers = createCodeEnvironmentHttpHandlers({ getAppConfig, registry: getCodeEnvironmentRegistry(), + principalIsActive: db.isAgentTriggerPrincipalActive, }); } return handlers; diff --git a/api/server/routes/code-environments.test.js b/api/server/routes/code-environments.test.js index b17afb1b732..0a8466e98e1 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -45,6 +45,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')]; diff --git a/client/src/components/Nav/Settings/CodeEnvironments.tsx b/client/src/components/Nav/Settings/CodeEnvironments.tsx index 7eeb7074ae4..22b66c0182a 100644 --- a/client/src/components/Nav/Settings/CodeEnvironments.tsx +++ b/client/src/components/Nav/Settings/CodeEnvironments.tsx @@ -25,7 +25,8 @@ import { useLocalize } from '~/hooks'; function WorkerCommand({ pairing }: { pairing: CodeEnvironmentPairingResponse['pairing'] }) { const { showToast } = useToastContext(); const localize = useLocalize(); - const command = `librechat-code pair ${pairing.endpoint} '${pairing.code}' --worker-id ${pairing.workerId}`; + const shellQuote = (value: string) => `'${value.replace(/'/g, `'\\''`)}'`; + const command = `librechat-code pair ${shellQuote(pairing.endpoint)} ${shellQuote(pairing.code)} --worker-id ${shellQuote(pairing.workerId)}`; return (

@@ -88,6 +89,51 @@ export default function CodeEnvironments() { ); }; + 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 (
@@ -97,36 +143,7 @@ export default function CodeEnvironments() {

- {controlPlanes.length > 0 ? ( -
- setName(event.target.value)} - /> - - -
- ) : ( -

- {localize('com_ui_code_environment_unavailable')} -

- )} + {pairingControls} {pairing != null && } @@ -140,40 +157,42 @@ export default function CodeEnvironments() {

{environment.name}

{environment.id}

- - - - - - {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-white transition-all duration-200 hover:bg-surface-destructive-hover', - selectText: localize('com_ui_delete'), - }} - /> - + {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/data-provider/CodeEnvironments.ts b/client/src/data-provider/CodeEnvironments.ts index d21fbf5fb5b..e1e45065315 100644 --- a/client/src/data-provider/CodeEnvironments.ts +++ b/client/src/data-provider/CodeEnvironments.ts @@ -22,7 +22,10 @@ export function usePairCodeEnvironmentMutation() { Error, { name: string; controlPlaneId: string } >([MutationKeys.pairCodeEnvironment], dataService.pairCodeEnvironment, { - onSuccess: () => queryClient.invalidateQueries([QueryKeys.codeEnvironments]), + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.codeEnvironments]); + queryClient.invalidateQueries([QueryKeys.endpoints]); + }, }); } diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 52cff2c39c0..1c29a9967f7 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -2183,6 +2183,7 @@ "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", diff --git a/packages/api/src/admin/users.ts b/packages/api/src/admin/users.ts index 39e0243a792..275dc553282 100644 --- a/packages/api/src/admin/users.ts +++ b/packages/api/src/admin/users.ts @@ -37,6 +37,8 @@ export interface AdminUsersDeps { ) => Promise; cancelAgentTriggerUserPurge: (userId: string, fenceStartedAt: Date) => Promise; purgeAgentTriggerDeliveriesForUser: (userId: string) => Promise; + revokeUserCodeEnvironmentWorkers?: (userId: string) => Promise; + deleteUserCodeEnvironments?: (userId: string) => Promise; /** * Thin data-layer delete — removes the User document only. * Full cascade of user-owned resources (conversations, messages, files, tokens, etc.) @@ -72,6 +74,8 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { prepareAgentTriggerUserPurge, cancelAgentTriggerUserPurge, purgeAgentTriggerDeliveriesForUser, + revokeUserCodeEnvironmentWorkers, + deleteUserCodeEnvironments, deleteUserById, deleteUserCodeEnvironments, invalidateCodeEnvironmentConfigCache, @@ -193,6 +197,7 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { } await prepareAgentTriggerUserPurge(id, triggerDeletionFence, targetUser?.tenantId); await drainAgentTriggerDeliveriesForUser(id); + await revokeUserCodeEnvironmentWorkers?.(id); const result = await deleteUserById(id); @@ -203,6 +208,7 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { return res.status(404).json({ error: 'User not found' }); } userDeleted = true; + await deleteUserCodeEnvironments?.(id); await purgeAgentTriggerDeliveriesForUser(id); if (targetUser?.role === SystemRoles.ADMIN) { diff --git a/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts index 2009c27d69c..ff476cdee33 100644 --- a/packages/api/src/code/bridge.ts +++ b/packages/api/src/code/bridge.ts @@ -34,7 +34,7 @@ export class CodeBridgePairingError extends Error { export class CodeBridgeLifecycleError extends Error { constructor( - public readonly reason: 'rejected' | 'timeout' | 'failed', + public readonly reason: 'rejected' | 'invalid' | 'timeout' | 'failed', public readonly upstreamStatus?: number, ) { super(`Code bridge lifecycle request ${reason}`); @@ -121,6 +121,15 @@ export async function revokeCodeBridgeWorker({ 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') { diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 4491ffed6ac..211bbb07979 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -61,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( @@ -143,7 +144,10 @@ describe('code environment registry', () => { role: 'CODE_USER', idOnTheSource: null, }), - ).resolves.toEqual([roleEnvironment, groupEnvironment]); + ).resolves.toEqual([ + { ...roleEnvironment, canDelete: false }, + { ...groupEnvironment, canDelete: false }, + ]); }); test('keeps a user-bound worker private even if its ACL is granted to a role', async () => { @@ -229,10 +233,14 @@ describe('code environment registry', () => { 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 }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); await expect( revokeUserCodeEnvironmentWorkers({ @@ -243,19 +251,7 @@ describe('code environment registry', () => { agents: { statefulCodeSessions: { allowedEnvironments: ['user'], - environments: [ - { - id: 'self-service', - name: 'Self-service', - type: 'attached', - baseURL: 'https://code.example.com/v1', - owner: 'deployment', - pairing: { - allowPrincipalWorkers: true, - tokenEnv: 'CODE_ADMIN_TOKEN', - }, - }, - ], + environments: [], }, }, }, diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index 048362aa753..fa61bd0e060 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; @@ -56,6 +58,7 @@ export type CodeEnvironmentLifecycleTarget = CodeEnvironmentSummary & { baseURL: string; workerId?: string; controlPlaneId?: string; + revocationTokenEnv?: string; workerPrincipal?: CodeEnvironmentRegistration['workerPrincipal']; }; @@ -124,17 +127,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, }; } @@ -152,6 +159,7 @@ export function createCodeEnvironmentRegistry( ) => Promise; listRegisteredIds: () => Promise; invalidateAccessibleConfigurations: (tenantId?: string) => Promise; + countOwned: (actor: CodeEnvironmentPrincipalContext) => Promise; remove: (params: { actor: CodeEnvironmentPrincipalContext; environmentId: string; @@ -213,7 +221,7 @@ export function createCodeEnvironmentRegistry( baseURL: environment.baseURL, controlPlaneId: environment.controlPlaneId, workerId: environment.workerId, - controlPlaneId: environment.controlPlaneId, + revocationTokenEnv: environment.revocationTokenEnv, workerPrincipal: environment.workerPrincipal, createdBy: new Types.ObjectId(actor.userId), }); @@ -229,7 +237,7 @@ 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(); return summary; } catch (error) { @@ -272,7 +280,20 @@ export function createCodeEnvironmentRegistry( actor: CodeEnvironmentPrincipalContext, ): Promise { const environments = await findAccessible(actor); - return environments.map(toSummary); + return await Promise.all( + environments.map(async (environment) => + toSummary( + environment, + await access.checkPermission({ + userId: actor.userId.toString(), + role: actor.role, + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + requiredPermission: PermissionBits.DELETE, + }), + ), + ), + ); } async function listAccessibleConfigurations( @@ -353,6 +374,10 @@ export function createCodeEnvironmentRegistry( return configurations.map(toPublicConfiguration); } + async function countOwned(actor: CodeEnvironmentPrincipalContext): Promise { + return (await methods.findCodeEnvironmentsByCreator(actor.userId)).length; + } + async function remove({ actor, environmentId, @@ -364,6 +389,12 @@ export function createCodeEnvironmentRegistry( }): 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, @@ -372,21 +403,31 @@ export function createCodeEnvironmentRegistry( requiredPermission: PermissionBits.DELETE, }); if (!allowed) return null; + const Agent = mongoose.models.Agent; + if (Agent != null && (await Agent.exists({ code_environment_id: environmentId })) != null) { + throw new CodeEnvironmentInUseError(environmentId); + } await beforeDelete?.({ ...toSummary(environment), baseURL: environment.baseURL, workerId: environment.workerId, controlPlaneId: environment.controlPlaneId, + revocationTokenEnv: environment.revocationTokenEnv, workerPrincipal: environment.workerPrincipal, }); const deleted = await methods.deleteCodeEnvironmentById(environment._id); if (deleted == null) return null; - await access.removeAllPermissions({ - resourceType: ResourceType.CODE_ENVIRONMENT, - resourceId: environment._id, - }); - return toSummary(deleted); + try { + await access.removeAllPermissions({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + } catch (error) { + logger.warn('[code-environments] environment deleted with orphaned ACL entries', error); + } + await invalidateAccessibleConfigurations(); + return toSummary(deleted, true); } return { @@ -395,6 +436,14 @@ export function createCodeEnvironmentRegistry( listAccessibleConfigurations, listRegisteredIds, invalidateAccessibleConfigurations, + countOwned, 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 ee594d4b905..f0f9d348b69 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -126,6 +126,7 @@ describe('code environment HTTP handlers', () => { baseURL: 'https://code.librechat.example/v1', workerId: 'code-generated', controlPlaneId: 'shared-code-api', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, }, }); @@ -510,6 +511,7 @@ describe('code environment HTTP handlers', () => { baseURL: 'https://code.librechat.example/v1', workerId: 'code-generated', controlPlaneId: 'self-service', + revocationTokenEnv: 'CODE_ADMIN_TOKEN', workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, }); return { @@ -519,7 +521,10 @@ describe('code environment HTTP handlers', () => { type: 'attached' as const, }; }); - const fetchImpl = jest.fn().mockResolvedValue({ ok: true }); + const fetchImpl = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); const handlers = createCodeEnvironmentHttpHandlers({ getAppConfig: jest.fn().mockResolvedValue({ endpoints: { diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index f72175d5c29..c9f35e2d653 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -18,7 +18,11 @@ import { readCodeBridgeSecret, revokeCodeBridgeWorker, } from './bridge'; -import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; +import { + CodeEnvironmentInUseError, + CodeEnvironmentValidationError, + normalizeCodeEnvironmentName, +} from './environments'; import { assertCodeApiJwtSigningReady, getCodeApiTenantId, @@ -32,6 +36,7 @@ type Registry = { environment: CodeEnvironmentRegistration; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; + countOwned?: (actor: CodeEnvironmentPrincipalContext) => Promise; remove: (params: { actor: CodeEnvironmentPrincipalContext; environmentId: string; @@ -54,6 +59,8 @@ export interface CodeEnvironmentHttpDeps { resolveTenantId?: (req: ServerRequest) => string; principalAuthEnabled?: () => boolean; principalAuthReady?: () => Promise | void; + principalIsActive?: (userId: string) => Promise; + maxPrincipalEnvironments?: number; fetchImpl?: CodeBridgeFetch; } @@ -146,6 +153,8 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) 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); @@ -295,6 +304,15 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } const workerId = createEnvironmentId(); + if ( + deps.registry.countOwned != null && + (await deps.registry.countOwned(principal)) >= maxPrincipalEnvironments + ) { + return res.status(409).json({ error: 'Personal code environment limit reached' }); + } + if (!(await principalIsActive(principal.userId.toString()))) { + return res.status(409).json({ error: 'Account deletion is already in progress' }); + } let pairing; try { pairing = await createCodeBridgePairing({ @@ -311,6 +329,16 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) return pairingErrorResponse(error, res); } + if (!(await principalIsActive(principal.userId.toString()))) { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + return res.status(409).json({ error: 'Account deletion is already in progress' }); + } + try { const environment = await deps.registry.register({ actor: principal, @@ -321,9 +349,20 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) baseURL: controlPlane.baseURL, workerId, controlPlaneId: controlPlane.id, + revocationTokenEnv: tokenEnv, workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); + if (!(await principalIsActive(principal.userId.toString()))) { + await revokeCodeBridgeWorker({ + baseURL: controlPlane.baseURL, + token, + workerId, + fetchImpl: deps.fetchImpl, + }); + await deps.registry.remove({ actor: principal, environmentId: workerId }); + return res.status(409).json({ error: 'Account deletion is already in progress' }); + } return res.status(201).json({ environment, pairing: { @@ -372,24 +411,13 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (!environmentId) { return res.status(400).json({ error: 'Code environment id is required' }); } - const appConfig = await deps.getAppConfig({ baseOnly: true }); try { const environment = await deps.registry.remove({ actor: principal, environmentId, beforeDelete: async (target) => { if (target.workerPrincipal?.type !== 'user' || target.workerId == null) return; - const controlPlane = configuredPrincipalControlPlane( - appConfig, - target.controlPlaneId ?? '', - ); - if (controlPlane == null || controlPlane.baseURL !== target.baseURL) { - throw new CodeEnvironmentLifecycleHttpError( - 409, - 'Code environment control plane is no longer configured', - ); - } - const tokenEnv = controlPlane.pairing?.tokenEnv; + const tokenEnv = target.revocationTokenEnv; const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; if (!token) { throw new CodeEnvironmentLifecycleHttpError( @@ -398,7 +426,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) ); } await revokeCodeBridgeWorker({ - baseURL: controlPlane.baseURL, + baseURL: target.baseURL, token, workerId: target.workerId, fetchImpl: deps.fetchImpl, @@ -413,6 +441,9 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) 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', diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 66de56d256d..3376b65bb3a 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -20,36 +20,44 @@ export async function revokeUserCodeEnvironmentWorkers({ const environments = await createMethods(mongoose).findCodeEnvironmentsByCreator(userId); const controlPlanes = appConfig.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments ?? []; - let revoked = 0; - for (const environment of environments) { - if (environment.workerPrincipal?.type !== 'user' || environment.workerId == null) continue; - const controlPlane = controlPlanes.find( - (candidate) => - candidate.id === environment.controlPlaneId && - candidate.owner === 'deployment' && - candidate.type === 'attached' && - candidate.pairing?.allowPrincipalWorkers === true && - candidate.baseURL === environment.baseURL, - ); - if (controlPlane == null) { - throw new Error( - `Code environment control plane is unavailable for ${environment.environmentId}`, - ); - } - const tokenEnv = controlPlane.pairing?.tokenEnv; - const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; - if (!token) { - throw new Error( - `Code environment revocation is unavailable for ${environment.environmentId}`, + const targets = environments + .filter( + (environment) => environment.workerPrincipal?.type === 'user' && environment.workerId != null, + ) + .map((environment) => { + const fallbackControlPlane = controlPlanes.find( + (candidate) => + candidate.id === environment.controlPlaneId && + candidate.owner === 'deployment' && + candidate.type === 'attached' && + candidate.pairing?.allowPrincipalWorkers === true && + candidate.baseURL === environment.baseURL, ); - } - await revokeCodeBridgeWorker({ - baseURL: controlPlane.baseURL, - token, - workerId: environment.workerId, - fetchImpl, + 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}`, + ); + } + return { environment, token }; }); - revoked += 1; + const results = await Promise.allSettled( + targets.map(({ environment, token }) => + revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }), + ), + ); + const failures = results.filter((result) => result.status === 'rejected'); + if (failures.length > 0) { + throw new AggregateError( + failures.map((failure) => failure.reason), + 'One or more code environment workers could not be revoked', + ); } - return revoked; + return targets.length; } diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 65cd0c110b7..31f92b1e687 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -539,6 +539,7 @@ export type TCodeEnvironmentSummary = { id: string; name: string; type: 'managed' | 'attached'; + canDelete: boolean; }; export type TCodeControlPlane = { diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 7c2f0e11836..8e02295f7cf 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -8,7 +8,10 @@ type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' > & - Pick, 'workerId' | 'controlPlaneId' | 'workerPrincipal'>; + Pick< + Partial, + 'workerId' | 'controlPlaneId' | 'revocationTokenEnv' | 'workerPrincipal' + >; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index f020a014113..1164bde30cd 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -51,6 +51,9 @@ const codeEnvironmentSchema: Schema = new Schema Date: Sun, 30 Aug 2026 15:18:58 -0400 Subject: [PATCH 12/29] fix: satisfy backend lifecycle types --- packages/api/src/code/http.spec.ts | 5 +++-- packages/api/src/code/lifecycle.ts | 19 +++++++++++++------ 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index f0f9d348b69..cb2b5be7fce 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -24,7 +24,7 @@ describe('code environment HTTP handlers', () => { 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(); @@ -144,7 +144,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(() => false), fetchImpl, }); @@ -519,6 +519,7 @@ describe('code environment HTTP handlers', () => { id: 'code-generated', name: 'Personal VM', type: 'attached' as const, + canDelete: true, }; }); const fetchImpl = jest.fn().mockResolvedValue({ diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 3376b65bb3a..72fb53b2ca6 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -25,6 +25,12 @@ export async function revokeUserCodeEnvironmentWorkers({ (environment) => environment.workerPrincipal?.type === 'user' && environment.workerId != null, ) .map((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 && @@ -40,24 +46,25 @@ export async function revokeUserCodeEnvironmentWorkers({ `Code environment revocation is unavailable for ${environment.environmentId}`, ); } - return { environment, token }; + return { environment, token, workerId }; }); const results = await Promise.allSettled( - targets.map(({ environment, token }) => + targets.map(({ environment, token, workerId }) => revokeCodeBridgeWorker({ baseURL: environment.baseURL, token, - workerId: environment.workerId, + workerId, fetchImpl, }), ), ); const failures = results.filter((result) => result.status === 'rejected'); if (failures.length > 0) { - throw new AggregateError( - failures.map((failure) => failure.reason), + const error = new Error( 'One or more code environment workers could not be revoked', - ); + ) as Error & { errors: unknown[] }; + error.errors = failures.map((failure) => failure.reason); + throw error; } return targets.length; } From d3b2857f560a6d50bc52d00dcfc9cdf6d2f7775b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 16:25:33 -0400 Subject: [PATCH 13/29] fix: harden principal environment lifecycle and quotas --- api/server/controllers/UserController.js | 26 +++-- api/server/controllers/UserController.spec.js | 7 ++ .../controllers/__tests__/deleteUser.spec.js | 22 +++-- config/__tests__/delete-user.spec.js | 24 +++++ config/delete-user.js | 14 +++ packages/api/src/admin/users.spec.ts | 7 ++ packages/api/src/admin/users.ts | 9 +- .../src/code/environments.integration.spec.ts | 96 +++++++++++++++++++ packages/api/src/code/environments.ts | 51 +++++----- packages/api/src/code/http.spec.ts | 67 ++++++++++++- packages/api/src/code/http.ts | 13 ++- packages/api/src/code/lifecycle.ts | 42 ++++---- .../src/methods/codeEnvironment.ts | 32 +++++++ .../src/schema/codeEnvironment.ts | 12 +++ .../data-schemas/src/types/codeEnvironment.ts | 1 + 15 files changed, 350 insertions(+), 73 deletions(-) diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index f21941320e6..1c0ca6b5a12 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -406,11 +406,6 @@ const deleteUserController = async (req, res) => { await prepareAgentTriggerUserPurge(user.id, triggerDeletionFence, user.tenantId); } const deletionAppConfig = await getAppConfig({ baseOnly: true }); - await revokeUserCodeEnvironmentWorkers({ - mongoose, - userId: user.id, - appConfig: deletionAppConfig, - }); await drainAgentTriggerDeliveriesForUser(user.id); await subagentThreadTaskStore.cancelAndDrainForOwner(user.id, user.tenantId); // Reversibly suspend the user's schedules under a per-attempt token BEFORE draining. @@ -467,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 }); @@ -482,6 +473,23 @@ const deleteUserController = async (req, res) => { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; + try { + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: user.id, + appConfig: deletionAppConfig, + }); + } catch (error) { + logger.error('[deleteUserController] Failed to revoke code environment workers', error); + } + 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 a7c221d40fb..7d3fd9ff8e1 100644 --- a/api/server/controllers/UserController.spec.js +++ b/api/server/controllers/UserController.spec.js @@ -163,6 +163,7 @@ const { deleteConvos, acceptTerms, deleteUserById, + deleteUserCodeEnvironments, deleteMessages, beginAgentTriggerUserDeletion, cancelAgentTriggerUserDeletion, @@ -493,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 06f9a8c49a2..0324d5bc143 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -219,9 +219,6 @@ describe('deleteUserController - 2FA enforcement', () => { mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0], ); expect(mockPrepareAgentTriggerUserPurge.mock.invocationCallOrder[0]).toBeLessThan( - mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0], - ); - expect(mockRevokeUserCodeEnvironmentWorkers.mock.invocationCallOrder[0]).toBeLessThan( mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0], ); expect(mockDrainAgentTriggerDeliveriesForUser.mock.invocationCallOrder[0]).toBeLessThan( @@ -230,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(); }); @@ -252,7 +255,7 @@ describe('deleteUserController - 2FA enforcement', () => { ); }); - it('keeps the account when a personal code worker cannot be revoked', async () => { + it('finishes committed account deletion when a personal code worker cannot be revoked', async () => { const req = { user: { id: 'user1', _id: 'user1', email: 'a@b.com' }, body: {} }; const res = createRes(); mockGetUserById.mockResolvedValue({ _id: 'user1', twoFactorEnabled: false }); @@ -260,12 +263,13 @@ describe('deleteUserController - 2FA enforcement', () => { await deleteUserController(req, res); - expect(res.status).toHaveBeenCalledWith(500); + expect(res.status).toHaveBeenCalledWith(200); expect(mockBeginAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); - expect(mockCancelAgentTriggerUserPurge).toHaveBeenCalledTimes(1); - expect(mockCancelAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); - expect(mockDeleteMessages).not.toHaveBeenCalled(); - expect(mockDeleteUserById).not.toHaveBeenCalled(); + expect(mockDeleteMessages).toHaveBeenCalled(); + expect(mockDeleteUserById).toHaveBeenCalledWith('user1'); + expect(mockDeleteUserCodeEnvironments).toHaveBeenCalledWith('user1'); + expect(mockCancelAgentTriggerUserPurge).not.toHaveBeenCalled(); + expect(mockCancelAgentTriggerUserDeletion).not.toHaveBeenCalled(); }); it('proceeds with deletion when user has no 2FA record', async () => { diff --git a/config/__tests__/delete-user.spec.js b/config/__tests__/delete-user.spec.js index 282fe377678..30256c5fe7d 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,6 +49,8 @@ jest.mock('@librechat/api', () => ({ abortJob: (...args) => mockAbortJob(...args), destroy: (...args) => mockDestroy(...args), }, + getAppConfig: (...args) => mockGetAppConfig(...args), + revokeUserCodeEnvironmentWorkers: (...args) => mockRevokeUserCodeEnvironmentWorkers(...args), })); jest.mock('~/cache/getLogStores', () => jest.fn()); jest.mock('../helpers', () => ({ @@ -96,6 +101,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 +109,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 +208,22 @@ 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('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 533a239f2f1..3eedf302d6f 100644 --- a/config/delete-user.js +++ b/config/delete-user.js @@ -34,6 +34,8 @@ require('module-alias')({ base: path.resolve(__dirname, '..', 'api') }); const { GenerationJobManager, createStreamServices, + getAppConfig, + revokeUserCodeEnvironmentWorkers, waitForKeyvRedisClient, } = require('@librechat/api'); const getLogStores = require('~/cache/getLogStores'); @@ -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,16 @@ async function gracefulExit(code = 0) { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; + await revokeUserCodeEnvironmentWorkers({ + mongoose, + userId: uid, + appConfig: deletionAppConfig, + }).catch((error) => + console.error('Failed to revoke code environment workers after account deletion:', error), + ); + 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 83ff375fe2f..f2ca9f5cbde 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.' }), @@ -584,12 +585,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 275dc553282..6f1f1e42728 100644 --- a/packages/api/src/admin/users.ts +++ b/packages/api/src/admin/users.ts @@ -38,7 +38,6 @@ export interface AdminUsersDeps { cancelAgentTriggerUserPurge: (userId: string, fenceStartedAt: Date) => Promise; purgeAgentTriggerDeliveriesForUser: (userId: string) => Promise; revokeUserCodeEnvironmentWorkers?: (userId: string) => Promise; - deleteUserCodeEnvironments?: (userId: string) => Promise; /** * Thin data-layer delete — removes the User document only. * Full cascade of user-owned resources (conversations, messages, files, tokens, etc.) @@ -77,7 +76,6 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { revokeUserCodeEnvironmentWorkers, deleteUserCodeEnvironments, deleteUserById, - deleteUserCodeEnvironments, invalidateCodeEnvironmentConfigCache, deleteConfig, deleteAclEntries, @@ -197,7 +195,6 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { } await prepareAgentTriggerUserPurge(id, triggerDeletionFence, targetUser?.tenantId); await drainAgentTriggerDeliveriesForUser(id); - await revokeUserCodeEnvironmentWorkers?.(id); const result = await deleteUserById(id); @@ -208,7 +205,11 @@ export function createAdminUsersHandlers(deps: AdminUsersDeps): { return res.status(404).json({ error: 'User not found' }); } userDeleted = true; - await deleteUserCodeEnvironments?.(id); + try { + await revokeUserCodeEnvironmentWorkers?.(id); + } catch (error) { + logger.error('[adminUsers] failed to revoke code environment workers:', id, error); + } await purgeAgentTriggerDeliveriesForUser(id); if (targetUser?.role === SystemRoles.ADMIN) { diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 211bbb07979..0311391fad2 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -150,6 +150,63 @@ describe('code environment registry', () => { ]); }); + 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', + }, + }); + } + + 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.syncIndexes(); + 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', + }, + 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 () => { const registry = createCodeEnvironmentRegistry(mongoose); const access = new AccessControlService(mongoose); @@ -266,6 +323,45 @@ describe('code environment registry', () => { ); }); + 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`, + workerId: id, + revocationTokenEnv: 'CODE_ADMIN_TOKEN', + workerPrincipal: { type: 'user', id: ownerId.toString() }, + }, + }); + } + const fetchImpl = jest.fn(async (url: string) => { + 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); + }); + test('removes creator-owned environment records and grants when the user is deleted', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index fa61bd0e060..953e443ac9f 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -95,6 +95,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); @@ -152,6 +159,7 @@ export function createCodeEnvironmentRegistry( register: (params: { actor: CodeEnvironmentPrincipalContext; environment: CodeEnvironmentRegistration; + maxOwned?: number; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; listAccessibleConfigurations: ( @@ -159,7 +167,6 @@ export function createCodeEnvironmentRegistry( ) => Promise; listRegisteredIds: () => Promise; invalidateAccessibleConfigurations: (tenantId?: string) => Promise; - countOwned: (actor: CodeEnvironmentPrincipalContext) => Promise; remove: (params: { actor: CodeEnvironmentPrincipalContext; environmentId: string; @@ -209,12 +216,14 @@ 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, @@ -224,7 +233,14 @@ export function createCodeEnvironmentRegistry( 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, @@ -280,20 +296,16 @@ export function createCodeEnvironmentRegistry( actor: CodeEnvironmentPrincipalContext, ): Promise { const environments = await findAccessible(actor); - return await Promise.all( - environments.map(async (environment) => - toSummary( - environment, - await access.checkPermission({ - userId: actor.userId.toString(), - role: actor.role, - resourceType: ResourceType.CODE_ENVIRONMENT, - resourceId: environment._id, - requiredPermission: PermissionBits.DELETE, - }), - ), - ), - ); + 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 listAccessibleConfigurations( @@ -374,10 +386,6 @@ export function createCodeEnvironmentRegistry( return configurations.map(toPublicConfiguration); } - async function countOwned(actor: CodeEnvironmentPrincipalContext): Promise { - return (await methods.findCodeEnvironmentsByCreator(actor.userId)).length; - } - async function remove({ actor, environmentId, @@ -436,7 +444,6 @@ export function createCodeEnvironmentRegistry( listAccessibleConfigurations, listRegisteredIds, invalidateAccessibleConfigurations, - countOwned, remove, }; } diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index cb2b5be7fce..5d37ae8aac7 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() { @@ -119,6 +119,7 @@ describe('code environment HTTP handlers', () => { role: 'USER', idOnTheSource: null, }, + maxOwned: 5, environment: { id: 'code-generated', name: 'Personal VM', @@ -140,6 +141,70 @@ describe('code environment HTTP handlers', () => { }); }); + test('revokes an upstream pairing when the atomic owner quota is exhausted', async () => { + const fetchImpl = jest.fn(async (url: string) => ({ + ok: true, + json: async () => + url.endsWith('/bridge/pairings') + ? { + protocolVersion: 1, + workerId: 'code-generated', + code: 'a'.repeat(32), + expiresAt: new Date(Date.now() + 60_000).toISOString(), + } + : { protocolVersion: 1, revoked: true }, + })); + 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, + 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('rejects self-service pairing without principal-aware Code API auth', async () => { const fetchImpl = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index c9f35e2d653..d4f17bb67af 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -20,6 +20,7 @@ import { } from './bridge'; import { CodeEnvironmentInUseError, + CodeEnvironmentLimitError, CodeEnvironmentValidationError, normalizeCodeEnvironmentName, } from './environments'; @@ -34,9 +35,9 @@ type Registry = { register: (params: { actor: CodeEnvironmentPrincipalContext; environment: CodeEnvironmentRegistration; + maxOwned?: number; }) => Promise; listAccessible: (actor: CodeEnvironmentPrincipalContext) => Promise; - countOwned?: (actor: CodeEnvironmentPrincipalContext) => Promise; remove: (params: { actor: CodeEnvironmentPrincipalContext; environmentId: string; @@ -304,12 +305,6 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } const workerId = createEnvironmentId(); - if ( - deps.registry.countOwned != null && - (await deps.registry.countOwned(principal)) >= maxPrincipalEnvironments - ) { - return res.status(409).json({ error: 'Personal code environment limit reached' }); - } if (!(await principalIsActive(principal.userId.toString()))) { return res.status(409).json({ error: 'Account deletion is already in progress' }); } @@ -342,6 +337,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) try { const environment = await deps.registry.register({ actor: principal, + maxOwned: maxPrincipalEnvironments, environment: { id: workerId, name, @@ -393,6 +389,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 }); } diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 72fb53b2ca6..ef69907de34 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -1,4 +1,4 @@ -import { createMethods } from '@librechat/data-schemas'; +import { createMethods, logger } from '@librechat/data-schemas'; import { EModelEndpoint } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; import type { CodeBridgeFetch } from './bridge'; @@ -20,11 +20,11 @@ export async function revokeUserCodeEnvironmentWorkers({ 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, - ) - .map((environment) => { + const targets = environments.filter( + (environment) => environment.workerPrincipal?.type === 'user' && environment.workerId != null, + ); + const results = await Promise.allSettled( + targets.map(async (environment) => { const workerId = environment.workerId; if (workerId == null) { throw new Error( @@ -46,25 +46,25 @@ export async function revokeUserCodeEnvironmentWorkers({ `Code environment revocation is unavailable for ${environment.environmentId}`, ); } - return { environment, token, workerId }; - }); - const results = await Promise.allSettled( - targets.map(({ environment, token, workerId }) => - revokeCodeBridgeWorker({ + await revokeCodeBridgeWorker({ baseURL: environment.baseURL, token, workerId, fetchImpl, - }), - ), + }); + return environment.environmentId; + }), ); - const failures = results.filter((result) => result.status === 'rejected'); - if (failures.length > 0) { - const error = new Error( - 'One or more code environment workers could not be revoked', - ) as Error & { errors: unknown[] }; - error.errors = failures.map((failure) => failure.reason); - throw error; + let revoked = 0; + for (const result of results) { + if (result.status === 'fulfilled') { + revoked += 1; + continue; + } + logger.error( + `[code-environments] worker revocation failed during committed user deletion: ${userId}`, + result.reason, + ); } - return targets.length; + return revoked; } diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 8e02295f7cf..cd097d3ce2b 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -15,6 +15,10 @@ type CreateCodeEnvironmentInput = Pick< export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; + createCodeEnvironmentWithinOwnerLimit: ( + input: CreateCodeEnvironmentInput, + maxOwned: number, + ) => Promise; findCodeEnvironmentsByIds: ( ids: Array, ) => Promise; @@ -38,6 +42,33 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') return (await model().create(input)).toObject() as CodeEnvironmentDocument; } + async function createCodeEnvironmentWithinOwnerLimit( + input: CreateCodeEnvironmentInput, + maxOwned: number, + ): Promise { + for (let ownerSlot = 0; ownerSlot < maxOwned; ownerSlot++) { + try { + return ( + await model().create({ ...input, ownerSlot }) + ).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 && + 'ownerSlot' in error.keyPattern; + if (!duplicateOwnerSlot) { + throw error; + } + } + } + return null; + } + async function findCodeEnvironmentsByIds( ids: Array, ): Promise { @@ -94,6 +125,7 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') return { createCodeEnvironment, + createCodeEnvironmentWithinOwnerLimit, findCodeEnvironmentsByIds, findCodeEnvironmentByEnvironmentId, listCodeEnvironmentIds, diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index 1164bde30cd..08eb4e7bbec 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -45,6 +45,10 @@ const codeEnvironmentSchema: Schema = new Schema = new Schema Date: Sun, 30 Aug 2026 17:00:25 -0400 Subject: [PATCH 14/29] fix: fence environment lifecycle races --- .../src/code/environments.integration.spec.ts | 106 +++++++++- packages/api/src/code/environments.ts | 48 +++-- packages/api/src/code/http.spec.ts | 65 ++++-- packages/api/src/code/http.ts | 15 +- packages/api/src/code/lifecycle.ts | 33 ++- packages/data-schemas/src/methods/agent.ts | 58 +++++- .../src/methods/codeEnvironment.ts | 193 +++++++++++++++++- .../src/schema/codeEnvironment.ts | 25 ++- .../data-schemas/src/types/codeEnvironment.ts | 5 + 9 files changed, 480 insertions(+), 68 deletions(-) diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 0311391fad2..94331780ec8 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -163,6 +163,7 @@ describe('code environment registry', () => { name: id, type: 'attached', baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', }, }); } @@ -178,7 +179,10 @@ describe('code environment registry', () => { }); test('atomically limits concurrent environment registrations for one owner', async () => { - await mongoose.models.CodeEnvironment.syncIndexes(); + 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(); @@ -191,6 +195,7 @@ describe('code environment registry', () => { name: id, type: 'attached', baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', }, maxOwned: 2, } as never), @@ -219,6 +224,7 @@ describe('code environment registry', () => { name: 'Owner worker', type: 'attached', baseURL: 'https://code.example.com', + controlPlaneId: 'shared-code-api', workerId: 'owner-worker', workerPrincipal: { type: 'user', id: ownerId.toString() }, }, @@ -279,6 +285,56 @@ describe('code environment registry', () => { ).resolves.toEqual([]); }); + 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({ @@ -340,7 +396,8 @@ describe('code environment registry', () => { }, }); } - const fetchImpl = jest.fn(async (url: string) => { + const fetchImpl = jest.fn(async (input: string | URL | Request) => { + const url = String(input); if (url.includes('unreachable')) { throw new Error('control plane unavailable'); } @@ -360,6 +417,17 @@ describe('code environment registry', () => { }), ).resolves.toBe(1); expect(fetchImpl).toHaveBeenCalledTimes(2); + await expect(createMethods(mongoose).deleteUserCodeEnvironments(ownerId)).resolves.toBe(1); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-ok' }), + ).resolves.toBeNull(); + await expect( + mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-unreachable' }).lean(), + ).resolves.toMatchObject({ + revocationPendingAt: expect.any(Date), + revocationAttempts: 1, + revocationLastError: 'Code bridge lifecycle request failed', + }); }); test('removes creator-owned environment records and grants when the user is deleted', async () => { @@ -585,4 +653,38 @@ describe('code environment registry', () => { mongoose.models.AclEntry.countDocuments({ resourceType: ResourceType.CODE_ENVIRONMENT }), ).resolves.toBe(0); }); + + 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); + }); }); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index 953e443ac9f..b2b8941e45e 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -262,7 +262,7 @@ export function createCodeEnvironmentRegistry( resourceType: ResourceType.CODE_ENVIRONMENT, resourceId: created._id, }), - methods.deleteCodeEnvironmentById(created._id), + methods.discardCodeEnvironmentById(created._id), ]); for (const result of cleanup) { if (result.status === 'rejected') { @@ -411,31 +411,39 @@ export function createCodeEnvironmentRegistry( requiredPermission: PermissionBits.DELETE, }); if (!allowed) return null; - const Agent = mongoose.models.Agent; - if (Agent != null && (await Agent.exists({ code_environment_id: environmentId })) != null) { + const removal = await methods.beginCodeEnvironmentRemoval(environment._id); + if (removal == null) { throw new CodeEnvironmentInUseError(environmentId); } - - await beforeDelete?.({ - ...toSummary(environment), - baseURL: environment.baseURL, - workerId: environment.workerId, - controlPlaneId: environment.controlPlaneId, - revocationTokenEnv: environment.revocationTokenEnv, - workerPrincipal: environment.workerPrincipal, - }); - const deleted = await methods.deleteCodeEnvironmentById(environment._id); - if (deleted == null) return null; try { - await access.removeAllPermissions({ - resourceType: ResourceType.CODE_ENVIRONMENT, - resourceId: environment._id, + const Agent = mongoose.models.Agent; + if (Agent != null && (await Agent.exists({ code_environment_id: environmentId })) != null) { + throw new CodeEnvironmentInUseError(environmentId); + } + await beforeDelete?.({ + ...toSummary(environment), + baseURL: environment.baseURL, + workerId: environment.workerId, + controlPlaneId: environment.controlPlaneId, + revocationTokenEnv: environment.revocationTokenEnv, + workerPrincipal: environment.workerPrincipal, }); + const deleted = await methods.deleteCodeEnvironmentById(environment._id); + if (deleted == null) return null; + try { + await access.removeAllPermissions({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + } catch (error) { + logger.warn('[code-environments] environment deleted with orphaned ACL entries', error); + } + await invalidateAccessibleConfigurations(); + return toSummary(deleted, true); } catch (error) { - logger.warn('[code-environments] environment deleted with orphaned ACL entries', error); + await methods.cancelCodeEnvironmentRemoval(environment._id); + throw error; } - await invalidateAccessibleConfigurations(); - return toSummary(deleted, true); } return { diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 5d37ae8aac7..fed952d2734 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -20,6 +20,44 @@ 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('returns 400 when registration has no request body', async () => { const register = jest.fn(); const handlers = createCodeEnvironmentHttpHandlers({ @@ -142,18 +180,21 @@ describe('code environment HTTP handlers', () => { }); test('revokes an upstream pairing when the atomic owner quota is exhausted', async () => { - const fetchImpl = jest.fn(async (url: string) => ({ - ok: true, - json: async () => - url.endsWith('/bridge/pairings') - ? { - protocolVersion: 1, - workerId: 'code-generated', - code: 'a'.repeat(32), - expiresAt: new Date(Date.now() + 60_000).toISOString(), - } - : { protocolVersion: 1, revoked: true }, - })); + 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: { diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index d4f17bb67af..530dec90556 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -11,6 +11,12 @@ 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, @@ -18,12 +24,6 @@ import { readCodeBridgeSecret, revokeCodeBridgeWorker, } from './bridge'; -import { - CodeEnvironmentInUseError, - CodeEnvironmentLimitError, - CodeEnvironmentValidationError, - normalizeCodeEnvironmentName, -} from './environments'; import { assertCodeApiJwtSigningReady, getCodeApiTenantId, @@ -168,7 +168,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) ]); return res.status(200).json({ environments, - controlPlanes: principalControlPlanes(appConfig), + controlPlanes: principalAuthEnabled() ? principalControlPlanes(appConfig) : [], }); } @@ -224,7 +224,6 @@ 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 }, diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index ef69907de34..b767b1c33e1 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -1,5 +1,5 @@ -import { createMethods, logger } from '@librechat/data-schemas'; import { EModelEndpoint } from 'librechat-data-provider'; +import { createMethods, logger } from '@librechat/data-schemas'; import type { AppConfig } from '@librechat/data-schemas'; import type { CodeBridgeFetch } from './bridge'; import { readCodeBridgeSecret, revokeCodeBridgeWorker } from './bridge'; @@ -23,6 +23,17 @@ export async function revokeUserCodeEnvironmentWorkers({ 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; @@ -56,11 +67,29 @@ export async function revokeUserCodeEnvironmentWorkers({ }), ); let revoked = 0; - for (const result of results) { + for (const [index, result] of results.entries()) { + const target = targets[index]; if (result.status === 'fulfilled') { revoked += 1; + await CodeEnvironment?.updateOne( + { _id: target._id }, + { + $unset: { + revocationPendingAt: 1, + revocationAttempts: 1, + revocationLastError: 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, diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index 5fc0cfad811..cddcce996d2 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -9,6 +9,10 @@ import { import type { FilterQuery, Model, PipelineStage, ProjectionType, Types } from 'mongoose'; import type { AgentToolResources } from 'librechat-data-provider'; import type { IAgent, IAclEntry } from '~/types'; +import { + releaseCodeEnvironmentReference, + reserveCodeEnvironmentReference, +} from './codeEnvironment'; import { filterExistingSkillIds } from './skill'; import logger from '~/config/winston'; @@ -585,7 +589,15 @@ export function createAgentMethods( extractMCPServerNames(agentData.tools as string[] | undefined), }; - return (await Agent.create(initialAgentData)).toObject() as IAgent; + const reservation = await reserveCodeEnvironmentReference( + mongoose, + typeof agentData.code_environment_id === 'string' ? agentData.code_environment_id : undefined, + ); + try { + return (await Agent.create(initialAgentData)).toObject() as IAgent; + } finally { + await releaseCodeEnvironmentReference(mongoose, reservation); + } } /** @@ -828,11 +840,28 @@ 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 Record).code_environment_id + : undefined; + let nextEnvironmentId: string | undefined; + if (typeof directEnvironmentId === 'string') { + nextEnvironmentId = directEnvironmentId; + } else if (typeof setEnvironmentId === 'string') { + nextEnvironmentId = setEnvironmentId; + } + const reservation = await reserveCodeEnvironmentReference(mongoose, nextEnvironmentId); + let updatedAgent: IAgent | null; + try { + updatedAgent = (await Agent.findOneAndUpdate( + searchParameter, + updateData, + mongoOptions, + ).lean()) as IAgent | null; + } finally { + await releaseCodeEnvironmentReference(mongoose, reservation); + } /** `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. @@ -1251,9 +1280,20 @@ export function createAgentMethods( const revertUpdate = restoresDeploymentDefault ? { $set: revertToVersion, $unset: { code_environment_id: 1 } } : { $set: revertToVersion }; - const revertedAgent = await Agent.findOneAndUpdate(searchParameter, revertUpdate, { - new: true, - }).lean(); + const reservation = await reserveCodeEnvironmentReference( + mongoose, + typeof revertToVersion.code_environment_id === 'string' + ? revertToVersion.code_environment_id + : undefined, + ); + let revertedAgent: IAgent | null; + try { + revertedAgent = await Agent.findOneAndUpdate(searchParameter, revertUpdate, { + new: true, + }).lean(); + } finally { + await releaseCodeEnvironmentReference(mongoose, reservation); + } 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 cd097d3ce2b..56a54520ec2 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -1,9 +1,88 @@ 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'; +const CODE_ENVIRONMENT_TOMBSTONES = 'code_environment_tombstones'; + +type CodeEnvironmentTombstone = { + _id: string; + deletedAt: Date; +}; + +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); +} + +export type CodeEnvironmentReferenceReservation = { + environmentId: string; + reservationId: string; +}; + +export class CodeEnvironmentReferenceError extends Error { + 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 reserved = await CodeEnvironment.findOneAndUpdate( + { environmentId, deletionStartedAt: { $exists: false } }, + { $addToSet: { pendingAgentReferences: reservationId } }, + { new: true }, + ).lean(); + if (reserved != null) return { environmentId, reservationId }; + + const [persisted, tombstone] = await Promise.all([ + CodeEnvironment.exists({ environmentId }), + tombstones(mongoose).findOne({ _id: 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: reservation.reservationId } }, + ); +} + type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' @@ -32,6 +111,11 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') deleteCodeEnvironmentById: ( id: string | Types.ObjectId, ) => Promise; + discardCodeEnvironmentById: (id: string | Types.ObjectId) => Promise; + beginCodeEnvironmentRemoval: ( + id: string | Types.ObjectId, + ) => Promise; + cancelCodeEnvironmentRemoval: (id: string | Types.ObjectId) => Promise; deleteUserCodeEnvironments: (userId: string | Types.ObjectId) => Promise; } { const model = () => mongoose.models.CodeEnvironment as Model; @@ -39,6 +123,9 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') async function createCodeEnvironment( input: CreateCodeEnvironmentInput, ): Promise { + if ((await tombstones(mongoose).findOne({ _id: input.environmentId })) != null) { + throw new Error(`Code environment id was previously retired: ${input.environmentId}`); + } return (await model().create(input)).toObject() as CodeEnvironmentDocument; } @@ -46,10 +133,34 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') input: CreateCodeEnvironmentInput, maxOwned: number, ): Promise { + if ((await tombstones(mongoose).findOne({ _id: 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, ownerSlot }) + await model().create({ + ...input, + _id: ownerSlotId(input.createdBy, ownerSlot), + ownerSlot, + }) ).toObject() as CodeEnvironmentDocument; } catch (error) { const duplicateOwnerSlot = @@ -60,7 +171,7 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') 'keyPattern' in error && typeof error.keyPattern === 'object' && error.keyPattern != null && - 'ownerSlot' in error.keyPattern; + ('_id' in error.keyPattern || 'ownerSlot' in error.keyPattern); if (!duplicateOwnerSlot) { throw error; } @@ -102,18 +213,83 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') 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: environment.environmentId }, + { $setOnInsert: { 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 beginCodeEnvironmentRemoval( + id: string | Types.ObjectId, + ): Promise { + return await model() + .findOneAndUpdate( + { + _id: id, + deletionStartedAt: { $exists: false }, + $or: [ + { pendingAgentReferences: { $exists: false } }, + { pendingAgentReferences: { $size: 0 } }, + ], + }, + { $set: { deletionStartedAt: new Date() } }, + { new: true }, + ) + .lean(); + } + + async function cancelCodeEnvironmentRemoval(id: string | Types.ObjectId): Promise { + await model().updateOne({ _id: id }, { $unset: { deletionStartedAt: 1 } }); + } + + /** 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 }, + deletionStartedAt: { $exists: false }, + $or: [ + { pendingAgentReferences: { $exists: false } }, + { pendingAgentReferences: { $size: 0 } }, + ], + }) + .select('_id environmentId') + .lean>>(); if (environments.length === 0) return 0; - const resourceIds = environments.map((environment) => environment._id); + const Agent = mongoose.models.Agent; + const referenced = new Set( + Agent == null + ? [] + : await Agent.distinct('code_environment_id', { + code_environment_id: { $in: environments.map(({ environmentId }) => environmentId) }, + }), + ); + const deletable = environments.filter(({ environmentId }) => !referenced.has(environmentId)); + if (deletable.length === 0) return 0; + + const resourceIds = deletable.map((environment) => environment._id); + await Promise.all( + deletable.map(({ environmentId }) => + tombstones(mongoose).updateOne( + { _id: environmentId }, + { $setOnInsert: { deletedAt: new Date() } }, + { upsert: true }, + ), + ), + ); const AclEntry = mongoose.models.AclEntry as Model; await AclEntry.deleteMany({ resourceType: ResourceType.CODE_ENVIRONMENT, @@ -131,6 +307,9 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') listCodeEnvironmentIds, findCodeEnvironmentsByCreator, deleteCodeEnvironmentById, + discardCodeEnvironmentById, + beginCodeEnvironmentRemoval, + cancelCodeEnvironmentRemoval, deleteUserCodeEnvironments, }; } diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index 08eb4e7bbec..e8c5808a8f8 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -49,6 +49,23 @@ const codeEnvironmentSchema: Schema = new Schema = new Schema Date: Sun, 30 Aug 2026 20:17:29 -0400 Subject: [PATCH 15/29] fix: reconcile code environment stack --- packages/api/src/code/config.ts | 4 +++- .../api/src/code/environments.integration.spec.ts | 1 + packages/api/src/code/http.spec.ts | 11 ++++++----- packages/data-schemas/src/methods/codeEnvironment.ts | 2 +- packages/data-schemas/src/schema/codeEnvironment.ts | 3 --- packages/data-schemas/src/types/codeEnvironment.ts | 1 - 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/packages/api/src/code/config.ts b/packages/api/src/code/config.ts index 3273bc19bf5..9bc344cffc8 100644 --- a/packages/api/src/code/config.ts +++ b/packages/api/src/code/config.ts @@ -130,7 +130,9 @@ export async function mergeAccessibleCodeEnvironments({ } let mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments].map( (environment) => - !isExecutableCodeEnvironment(environment) && environment.default === true + !isExecutableCodeEnvironment(environment) && + 'default' in environment && + environment.default === true ? { ...environment, default: false as const } : environment, ); diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 94331780ec8..439a2264a0e 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -390,6 +390,7 @@ describe('code environment registry', () => { 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() }, diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index fed952d2734..185175e2dca 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -226,6 +226,7 @@ describe('code environment HTTP handlers', () => { createEnvironmentId: () => 'code-generated', readSecret: () => 'administrator-token', principalAuthEnabled: () => true, + principalAuthReady: jest.fn().mockResolvedValue(undefined), fetchImpl, }); const res = response(); @@ -290,13 +291,13 @@ describe('code environment HTTP handlers', () => { } as AppConfig; const effectiveConfig = { endpoints: { [EModelEndpoint.agents]: { statefulCodeSessions: { environments: [] } } }, - } as AppConfig; + } as unknown as AppConfig; const getAppConfig = jest.fn(async (options) => options.baseOnly === true ? baseConfig : effectiveConfig, ); 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(), @@ -330,7 +331,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, @@ -405,6 +406,7 @@ describe('code environment HTTP handlers', () => { registry: { register: jest.fn().mockResolvedValue({ id: 'code-generated' }), listAccessible: jest.fn(), + remove: jest.fn(), }, createEnvironmentId: () => 'code-generated', readSecret, @@ -436,7 +438,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'); @@ -515,7 +517,6 @@ describe('code environment HTTP handlers', () => { baseURL: 'https://code.librechat.example', controlPlaneId: 'shared-code-api', workerId: 'deployment-worker', - controlPlaneId: 'shared-code-api', workerPrincipal: { type: 'deployment', id: 'shared-code-api' }, }, }); diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 56a54520ec2..97a8c4dadf7 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -89,7 +89,7 @@ type CreateCodeEnvironmentInput = Pick< > & Pick< Partial, - 'workerId' | 'controlPlaneId' | 'revocationTokenEnv' | 'workerPrincipal' + 'workerId' | 'revocationTokenEnv' | 'workerPrincipal' >; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index e8c5808a8f8..d71f1cf60a7 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -69,9 +69,6 @@ const codeEnvironmentSchema: Schema = new Schema Date: Tue, 1 Sep 2026 13:55:22 -0400 Subject: [PATCH 16/29] style: format code execution headers --- api/server/services/Files/Code/process.js | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 61efdea64e6..479386d4590 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -1788,9 +1788,7 @@ async function execSandboxImageChunk({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, - ...(executionProfile - ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) - : {}), + ...(executionProfile ? codeExecutionHeaders({ executionProfile, bridgeWorkerId }) : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, From 1cc4c46eb3c7a8569307b6059c8bd7aacb923615 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 13:58:10 -0400 Subject: [PATCH 17/29] fix: reconcile code environment lifecycle --- api/server/controllers/UserController.js | 12 +- .../controllers/__tests__/deleteUser.spec.js | 4 +- api/server/controllers/agents/v1.js | 10 +- api/server/routes/code-environments.js | 7 +- api/server/routes/code-environments.test.js | 1 + .../Nav/Settings/CodeEnvironments.tsx | 10 +- config/__tests__/delete-user.spec.js | 8 + config/delete-user.js | 26 ++- packages/api/src/admin/users.spec.ts | 13 ++ packages/api/src/admin/users.ts | 4 +- .../src/code/environments.integration.spec.ts | 190 +++++++++++++++- packages/api/src/code/environments.ts | 126 +++++++---- packages/api/src/code/http.spec.ts | 67 ++++++ packages/api/src/code/http.ts | 47 ++-- packages/api/src/code/lifecycle.ts | 192 +++++++++++++++- .../src/methods/codeEnvironment.ts | 207 +++++++++++++----- .../src/schema/codeEnvironment.ts | 16 +- .../data-schemas/src/types/codeEnvironment.ts | 11 +- 18 files changed, 807 insertions(+), 144 deletions(-) diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index 1c0ca6b5a12..d90f7f0f5b9 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -473,6 +473,7 @@ const deleteUserController = async (req, res) => { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; + let codeEnvironmentCleanupSafe = true; try { await revokeUserCodeEnvironmentWorkers({ mongoose, @@ -480,12 +481,15 @@ const deleteUserController = async (req, res) => { appConfig: deletionAppConfig, }); } catch (error) { + codeEnvironmentCleanupSafe = false; logger.error('[deleteUserController] Failed to revoke code environment workers', error); } - try { - await db.deleteUserCodeEnvironments(user.id); - } catch (error) { - logger.error('[deleteUserController] Failed to delete code environments', 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); diff --git a/api/server/controllers/__tests__/deleteUser.spec.js b/api/server/controllers/__tests__/deleteUser.spec.js index 0324d5bc143..aabf4a63c0b 100644 --- a/api/server/controllers/__tests__/deleteUser.spec.js +++ b/api/server/controllers/__tests__/deleteUser.spec.js @@ -255,7 +255,7 @@ describe('deleteUserController - 2FA enforcement', () => { ); }); - it('finishes committed account deletion when a personal code worker cannot be revoked', async () => { + 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 }); @@ -267,7 +267,7 @@ describe('deleteUserController - 2FA enforcement', () => { expect(mockBeginAgentTriggerUserDeletion).toHaveBeenCalledTimes(1); expect(mockDeleteMessages).toHaveBeenCalled(); expect(mockDeleteUserById).toHaveBeenCalledWith('user1'); - expect(mockDeleteUserCodeEnvironments).toHaveBeenCalledWith('user1'); + expect(mockDeleteUserCodeEnvironments).not.toHaveBeenCalled(); expect(mockCancelAgentTriggerUserPurge).not.toHaveBeenCalled(); expect(mockCancelAgentTriggerUserDeletion).not.toHaveBeenCalled(); }); diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index 7cf87113a46..c0813081465 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -890,6 +890,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 }); } }; @@ -1577,7 +1580,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 }); } }; @@ -2106,6 +2111,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/code-environments.js b/api/server/routes/code-environments.js index 06da12c5dc3..4e866dbe923 100644 --- a/api/server/routes/code-environments.js +++ b/api/server/routes/code-environments.js @@ -1,5 +1,9 @@ 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'); @@ -8,6 +12,7 @@ const { requireJwtAuth } = require('~/server/middleware'); const db = require('~/models'); const router = express.Router(); +startCodeEnvironmentLifecycleReconciler({ mongoose }); let handlers; function getHandlers() { if (handlers == null) { diff --git a/api/server/routes/code-environments.test.js b/api/server/routes/code-environments.test.js index 0a8466e98e1..eda05129a1b 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -31,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', () => ({ diff --git a/client/src/components/Nav/Settings/CodeEnvironments.tsx b/client/src/components/Nav/Settings/CodeEnvironments.tsx index 22b66c0182a..a51538b96ef 100644 --- a/client/src/components/Nav/Settings/CodeEnvironments.tsx +++ b/client/src/components/Nav/Settings/CodeEnvironments.tsx @@ -1,4 +1,4 @@ -import { useState } from 'react'; +import { useEffect, useState } from 'react'; import { Button, Input, @@ -66,6 +66,14 @@ export default function CodeEnvironments() { 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; diff --git a/config/__tests__/delete-user.spec.js b/config/__tests__/delete-user.spec.js index 30256c5fe7d..8af07560157 100644 --- a/config/__tests__/delete-user.spec.js +++ b/config/__tests__/delete-user.spec.js @@ -224,6 +224,14 @@ describe('Delete user CLI', () => { ); }); + 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 3eedf302d6f..5fbd79a9c6c 100644 --- a/config/delete-user.js +++ b/config/delete-user.js @@ -238,16 +238,22 @@ async function gracefulExit(code = 0) { throw new Error('User disappeared before account deletion could commit'); } userDeleted = true; - await revokeUserCodeEnvironmentWorkers({ - mongoose, - userId: uid, - appConfig: deletionAppConfig, - }).catch((error) => - console.error('Failed to revoke code environment workers after account deletion:', error), - ); - await runAsSystem(() => methods.deleteUserCodeEnvironments(uid)).catch((error) => - console.error('Failed to delete code environment records after account deletion:', error), - ); + 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 f2ca9f5cbde..824dd725145 100644 --- a/packages/api/src/admin/users.spec.ts +++ b/packages/api/src/admin/users.spec.ts @@ -477,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) }); diff --git a/packages/api/src/admin/users.ts b/packages/api/src/admin/users.ts index 6f1f1e42728..0103c726963 100644 --- a/packages/api/src/admin/users.ts +++ b/packages/api/src/admin/users.ts @@ -205,9 +205,11 @@ 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); @@ -225,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/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 439a2264a0e..478d47955e1 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -1,10 +1,10 @@ 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'; -import { revokeUserCodeEnvironmentWorkers } from './lifecycle'; function createSharedCache() { const values = new Map(); @@ -285,6 +285,63 @@ describe('code environment registry', () => { ).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); @@ -418,10 +475,10 @@ describe('code environment registry', () => { }), ).resolves.toBe(1); expect(fetchImpl).toHaveBeenCalledTimes(2); - await expect(createMethods(mongoose).deleteUserCodeEnvironments(ownerId)).resolves.toBe(1); + await expect(createMethods(mongoose).deleteUserCodeEnvironments(ownerId)).resolves.toBe(0); await expect( - mongoose.models.CodeEnvironment.findOne({ environmentId: 'worker-ok' }), - ).resolves.toBeNull(); + 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({ @@ -429,6 +486,23 @@ describe('code environment registry', () => { revocationAttempts: 1, revocationLastError: 'Code bridge lifecycle request failed', }); + + const retryFetch = jest.fn().mockResolvedValue({ + ok: true, + json: async () => ({ protocolVersion: 1, revoked: true }), + }); + 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('removes creator-owned environment records and grants when the user is deleted', async () => { @@ -655,6 +729,49 @@ describe('code environment registry', () => { ).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', + }, + }), + ).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('preserves a creator-owned environment referenced by another surviving agent', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); @@ -688,4 +805,67 @@ describe('code environment registry', () => { 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('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', + }); + }); + }); }); diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index b2b8941e45e..f2f05cbbd0a 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -255,18 +255,24 @@ export function createCodeEnvironmentRegistry( } 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.discardCodeEnvironmentById(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; @@ -288,7 +294,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), ); } @@ -327,28 +336,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'); @@ -377,11 +399,16 @@ export function createCodeEnvironmentRegistry( .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); } @@ -415,33 +442,44 @@ export function createCodeEnvironmentRegistry( 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({ code_environment_id: environmentId })) != null) { throw new CodeEnvironmentInUseError(environmentId); } - await beforeDelete?.({ - ...toSummary(environment), - baseURL: environment.baseURL, - workerId: environment.workerId, - controlPlaneId: environment.controlPlaneId, - revocationTokenEnv: environment.revocationTokenEnv, - workerPrincipal: environment.workerPrincipal, + 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; - try { - await access.removeAllPermissions({ - resourceType: ResourceType.CODE_ENVIRONMENT, - resourceId: environment._id, - }); - } catch (error) { - logger.warn('[code-environments] environment deleted with orphaned ACL entries', error); - } await invalidateAccessibleConfigurations(); return toSummary(deleted, true); } catch (error) { - await methods.cancelCodeEnvironmentRemoval(environment._id); + if (!deletionCommitted && !externalLifecycleStarted) { + await methods.cancelCodeEnvironmentRemoval(environment._id, removalLeaseId); + } throw error; } } diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 185175e2dca..26c4fd05a32 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -247,6 +247,73 @@ describe('code environment HTTP handlers', () => { ); }); + 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({ diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 530dec90556..bf3d8a5e2f7 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -323,14 +323,29 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) return pairingErrorResponse(error, res); } - if (!(await principalIsActive(principal.userId.toString()))) { - await revokeCodeBridgeWorker({ - baseURL: controlPlane.baseURL, - token, - workerId, - fetchImpl: deps.fetchImpl, - }); - return res.status(409).json({ error: 'Account deletion is already in progress' }); + 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' }); } try { @@ -348,15 +363,15 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); - if (!(await principalIsActive(principal.userId.toString()))) { - await revokeCodeBridgeWorker({ - baseURL: controlPlane.baseURL, - token, - workerId, - fetchImpl: deps.fetchImpl, - }); + let activeAfterRegistration = false; + try { + activeAfterRegistration = await principalIsActive(principal.userId.toString()); + } catch (error) { + logger.error('[codeEnvironments] post-registration principal check failed:', error); + } + if (!activeAfterRegistration) { await deps.registry.remove({ actor: principal, environmentId: workerId }); - return res.status(409).json({ error: 'Account deletion is already in progress' }); + return res.status(409).json({ error: 'Account is unavailable for code worker pairing' }); } return res.status(201).json({ environment, diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index b767b1c33e1..38abe5090f0 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -1,9 +1,194 @@ -import { EModelEndpoint } from 'librechat-data-provider'; -import { createMethods, logger } from '@librechat/data-schemas'; -import type { AppConfig } from '@librechat/data-schemas'; +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; +let reconcileTimer: NodeJS.Timeout | undefined; +let reconcileInFlight: Promise | undefined; + +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 }, + }) + .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; + if (environment.workerPrincipal?.type !== 'user' || environment.workerId == null) { + await methods.cancelCodeEnvironmentRemoval(environment._id, leaseId); + continue; + } + const tokenEnv = environment.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (!token) continue; + try { + 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 pendingRegistrations = await CodeEnvironment.find({ + registrationPendingAt: { $lte: staleRegistration }, + }) + .limit(limit) + .lean(); + for (const environment of pendingRegistrations) { + const tokenEnv = environment.revocationTokenEnv; + const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; + if (environment.workerId != null) { + if (!token) continue; + try { + await revokeCodeBridgeWorker({ + baseURL: environment.baseURL, + token, + workerId: environment.workerId, + fetchImpl, + }); + } catch { + continue; + } + } + await AclEntry.deleteMany({ + resourceType: ResourceType.CODE_ENVIRONMENT, + resourceId: environment._id, + }); + await CodeEnvironment.deleteOne({ + _id: environment._id, + registrationPendingAt: environment.registrationPendingAt, + }); + } + + 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); + } + + const candidates = await CodeEnvironment.find({ revocationPendingAt: { $exists: true } }) + .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 }, + $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 }, + $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, @@ -74,6 +259,7 @@ export async function revokeUserCodeEnvironmentWorkers({ await CodeEnvironment?.updateOne( { _id: target._id }, { + $set: { deletionCommittedAt: new Date() }, $unset: { revocationPendingAt: 1, revocationAttempts: 1, diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 97a8c4dadf7..9b231e140ab 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -4,14 +4,30 @@ 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'; const CODE_ENVIRONMENT_TOMBSTONES = 'code_environment_tombstones'; +const REFERENCE_LEASE_MS = 2 * 60_000; +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) { @@ -34,6 +50,7 @@ export type CodeEnvironmentReferenceReservation = { }; export class CodeEnvironmentReferenceError extends Error { + readonly statusCode = 409; constructor(environmentId: string) { super(`Code environment is being removed: ${environmentId}`); this.name = 'CodeEnvironmentReferenceError'; @@ -50,16 +67,23 @@ export async function reserveCodeEnvironmentReference( | 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, deletionStartedAt: { $exists: false } }, - { $addToSet: { pendingAgentReferences: reservationId } }, + { + 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: environmentId }), + tombstones(mongoose).findOne({ _id: tombstoneId(environmentId) }), ]); if (persisted != null || tombstone != null) { throw new CodeEnvironmentReferenceError(environmentId); @@ -79,7 +103,7 @@ export async function releaseCodeEnvironmentReference( if (CodeEnvironment == null) return; await CodeEnvironment.updateOne( { environmentId: reservation.environmentId }, - { $pull: { pendingAgentReferences: reservation.reservationId } }, + { $pull: { pendingAgentReferences: { reservationId: reservation.reservationId } } }, ); } @@ -87,10 +111,7 @@ type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' > & - Pick< - Partial, - 'workerId' | 'revocationTokenEnv' | 'workerPrincipal' - >; + Pick, 'workerId' | 'revocationTokenEnv' | 'workerPrincipal'>; export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose')): { createCodeEnvironment: (input: CreateCodeEnvironmentInput) => Promise; @@ -112,10 +133,12 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') 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) => 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; @@ -123,17 +146,19 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') async function createCodeEnvironment( input: CreateCodeEnvironmentInput, ): Promise { - if ((await tombstones(mongoose).findOne({ _id: input.environmentId })) != null) { + 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)).toObject() as CodeEnvironmentDocument; + 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: input.environmentId })) != null) { + 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() @@ -160,6 +185,7 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') ...input, _id: ownerSlotId(input.createdBy, ownerSlot), ownerSlot, + registrationPendingAt: new Date(), }) ).toObject() as CodeEnvironmentDocument; } catch (error) { @@ -216,8 +242,14 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') const environment = await model().findById(id).lean(); if (environment == null) return null; await tombstones(mongoose).updateOne( - { _id: environment.environmentId }, - { $setOnInsert: { deletedAt: new Date() } }, + { _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(); @@ -227,27 +259,89 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') await model().deleteOne({ _id: id }); } + async function completeCodeEnvironmentRegistration(id: string | Types.ObjectId): Promise { + const result = await model().updateOne( + { _id: id, registrationPendingAt: { $exists: true } }, + { $unset: { registrationPendingAt: 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, - deletionStartedAt: { $exists: false }, - $or: [ - { pendingAgentReferences: { $exists: false } }, - { pendingAgentReferences: { $size: 0 } }, + 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: new Date() } }, + { + $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): Promise { - await model().updateOne({ _id: id }, { $unset: { deletionStartedAt: 1 } }); + 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. @@ -259,44 +353,47 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') .find({ createdBy: creatorId, revocationPendingAt: { $exists: false }, - deletionStartedAt: { $exists: false }, - $or: [ - { pendingAgentReferences: { $exists: false } }, - { pendingAgentReferences: { $size: 0 } }, - ], + deletionCommittedAt: { $exists: false }, }) - .select('_id environmentId') - .lean>>(); + .select('_id') + .lean>>(); if (environments.length === 0) return 0; const Agent = mongoose.models.Agent; - const referenced = new Set( - Agent == null - ? [] - : await Agent.distinct('code_environment_id', { - code_environment_id: { $in: environments.map(({ environmentId }) => environmentId) }, - }), - ); - const deletable = environments.filter(({ environmentId }) => !referenced.has(environmentId)); - if (deletable.length === 0) return 0; - - const resourceIds = deletable.map((environment) => environment._id); - await Promise.all( - deletable.map(({ environmentId }) => - tombstones(mongoose).updateOne( - { _id: environmentId }, - { $setOnInsert: { deletedAt: new Date() } }, - { upsert: true }, - ), - ), - ); 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({ code_environment_id: claimed.environmentId })) != 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 { @@ -308,8 +405,10 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') 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 d71f1cf60a7..74b807cff5e 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -50,12 +50,24 @@ const codeEnvironmentSchema: Schema = new Schema = new Schema; deletionStartedAt?: Date; + deletionLeaseId?: string; + deletionLeaseExpiresAt?: Date; + deletionCommittedAt?: Date; + registrationPendingAt?: Date; revocationPendingAt?: Date; revocationAttempts?: number; revocationLastError?: string; + revocationLeaseId?: string; + revocationLeaseExpiresAt?: Date; workerId?: string; revocationTokenEnv?: string; workerPrincipal?: { From b83b296aa73aee7051a098457ec8ceec79a81009 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 14:05:49 -0400 Subject: [PATCH 18/29] fix: close self-service pairing gaps --- api/server/middleware/limiters/code.js | 9 +++++++-- api/server/middleware/limiters/code.spec.js | 12 ++++++++++++ api/server/services/ToolService.js | 3 +++ api/server/services/__tests__/ToolService.spec.js | 2 ++ packages/api/src/code/http.spec.ts | 1 + packages/api/src/code/http.ts | 1 + 6 files changed, 26 insertions(+), 2 deletions(-) diff --git a/api/server/middleware/limiters/code.js b/api/server/middleware/limiters/code.js index 6b2d081fe94..301ae08cfcc 100644 --- a/api/server/middleware/limiters/code.js +++ b/api/server/middleware/limiters/code.js @@ -3,10 +3,15 @@ const { limiterCache } = require('@librechat/api'); let configuredCodeEnvironmentPairingLimiter; +function positiveInteger(value, fallback) { + const parsed = Number(value); + return Number.isSafeInteger(parsed) && parsed >= 1 ? parsed : fallback; +} + const codeEnvironmentPairingLimiter = (req, res, next) => { if (configuredCodeEnvironmentPairingLimiter == null) { - const max = Number(process.env.CODE_ENVIRONMENT_PAIRING_USER_MAX ?? 5); - const windowInMinutes = Number(process.env.CODE_ENVIRONMENT_PAIRING_USER_WINDOW ?? 60); + const max = positiveInteger(process.env.CODE_ENVIRONMENT_PAIRING_USER_MAX, 5); + const windowInMinutes = positiveInteger(process.env.CODE_ENVIRONMENT_PAIRING_USER_WINDOW, 60); configuredCodeEnvironmentPairingLimiter = rateLimit({ windowMs: windowInMinutes * 60 * 1000, diff --git a/api/server/middleware/limiters/code.spec.js b/api/server/middleware/limiters/code.spec.js index bea8a40adef..cfb4747c14d 100644 --- a/api/server/middleware/limiters/code.spec.js +++ b/api/server/middleware/limiters/code.spec.js @@ -44,6 +44,18 @@ describe('code environment pairing rate limiter', () => { expect(mockLimiter).toHaveBeenCalledWith(req, {}, next); }); + it('falls back to safe limits for invalid environment values', () => { + process.env.CODE_ENVIRONMENT_PAIRING_USER_MAX = 'NaN'; + process.env.CODE_ENVIRONMENT_PAIRING_USER_WINDOW = '-1'; + const { codeEnvironmentPairingLimiter } = require('./code'); + + codeEnvironmentPairingLimiter({ user: { id: 'user-1' } }, {}, jest.fn()); + + expect(mockRateLimit).toHaveBeenLastCalledWith( + expect.objectContaining({ max: 5, windowMs: 3_600_000 }), + ); + }); + it('returns an actionable JSON response when the limit is exceeded', async () => { const { codeEnvironmentPairingLimiter } = require('./code'); diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index fc57e98b370..c5d4ac0950e 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -1394,6 +1394,9 @@ async function loadToolDefinitionsWrapper({ codeApiBaseUrl: resolvedCodeExecutionContext.baseUrl, executionProfile: resolvedCodeExecutionContext.executionProfile, executionRouteKey: resolvedCodeExecutionContext.executionRouteKey, + ...(resolvedCodeExecutionContext.bridgeWorkerId + ? { bridgeWorkerId: resolvedCodeExecutionContext.bridgeWorkerId } + : {}), }); if (toolContext) { dynamicToolContextMap[Tools.execute_code] = toolContext; diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index c92edd971b6..fe7506252a7 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -955,6 +955,7 @@ describe('ToolService - Action Capability Gating', () => { executionProfile: 'stateful', runtimeSessionHint: 'v2:user:abc', statefulSessions: true, + bridgeWorkerId: 'worker-abc', }, }); @@ -965,6 +966,7 @@ describe('ToolService - Action Capability Gating', () => { agentResourceType: undefined, codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'worker-abc', }); }); diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 2871bbadb00..215700cf4ce 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -125,6 +125,7 @@ describe('code environment HTTP handlers', () => { type: 'attached', baseURL: 'https://code.librechat.example/v1', workerId: 'code-generated', + controlPlaneId: 'shared-code-api', workerPrincipal: { type: 'user', id: '68b2f0c498f24c1e78fa0001' }, }, }); diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index c341a31f173..b6c485759d8 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -278,6 +278,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) type: 'attached', baseURL: controlPlane.baseURL, workerId, + controlPlaneId: controlPlane.id, workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); From 8887182d03678ef8e70272d38f66c2d76774ef04 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 14:17:48 -0400 Subject: [PATCH 19/29] fix: close code environment reconciliation gaps --- .../src/code/environments.integration.spec.ts | 156 ++++++++++++++++++ packages/api/src/code/environments.ts | 21 ++- packages/api/src/code/http.spec.ts | 118 +++++++++++++ packages/api/src/code/http.ts | 36 +++- packages/api/src/code/lifecycle.ts | 97 +++++++++-- .../src/methods/codeEnvironment.ts | 8 +- .../src/schema/codeEnvironment.ts | 3 + .../data-schemas/src/types/codeEnvironment.ts | 3 + 8 files changed, 418 insertions(+), 24 deletions(-) diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 478d47955e1..8676de372a0 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -556,6 +556,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 }); @@ -772,6 +803,52 @@ describe('code environment registry', () => { ).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); @@ -841,6 +918,85 @@ describe('code environment registry', () => { 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 = { diff --git a/packages/api/src/code/environments.ts b/packages/api/src/code/environments.ts index f2f05cbbd0a..55f613e6d42 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -394,8 +394,25 @@ 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); } diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 26c4fd05a32..2c4f1afdd3a 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -58,6 +58,42 @@ describe('code environment HTTP handlers', () => { 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({ @@ -247,6 +283,88 @@ describe('code environment HTTP handlers', () => { ); }); + 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('revokes an issued pairing when the post-issue principal check is unavailable', async () => { const register = jest.fn(); const principalIsActive = jest diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index bf3d8a5e2f7..f7ebcf0d04c 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -162,10 +162,17 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) if (principal == null) { return res.status(401).json({ error: 'Authentication required' }); } - const [environments, appConfig] = await Promise.all([ - deps.registry.listAccessible(principal), - deps.getAppConfig({ baseOnly: true }), - ]); + 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) : [], @@ -370,7 +377,26 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) logger.error('[codeEnvironments] post-registration principal check failed:', error); } if (!activeAfterRegistration) { - await deps.registry.remove({ actor: principal, environmentId: workerId }); + 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({ diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 38abe5090f0..1b53b1e2f0f 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -7,6 +7,7 @@ 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; let reconcileTimer: NodeJS.Timeout | undefined; let reconcileInFlight: Promise | undefined; @@ -35,6 +36,7 @@ export async function reconcileCodeEnvironmentLifecycle({ deletionCommittedAt: { $exists: false }, deletionLeaseExpiresAt: { $lte: now }, }) + .sort({ _id: 1 }) .limit(limit) .select('_id') .lean>>(); @@ -42,20 +44,26 @@ export async function reconcileCodeEnvironmentLifecycle({ const environment = await methods.beginCodeEnvironmentRemoval(candidate._id); if (environment == null || environment.deletionLeaseId == null) continue; const leaseId = environment.deletionLeaseId; - if (environment.workerPrincipal?.type !== 'user' || environment.workerId == null) { + const Agent = mongoose.models.Agent; + if ( + Agent != null && + (await Agent.exists({ code_environment_id: environment.environmentId })) != null + ) { await methods.cancelCodeEnvironmentRemoval(environment._id, leaseId); continue; } - const tokenEnv = environment.revocationTokenEnv; - const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; - if (!token) continue; try { - await revokeCodeBridgeWorker({ - baseURL: environment.baseURL, - token, - workerId: environment.workerId, - fetchImpl, - }); + 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, @@ -67,16 +75,74 @@ export async function reconcileCodeEnvironmentLifecycle({ } } const staleRegistration = new Date(Date.now() - REGISTRATION_STALE_MS); - const pendingRegistrations = await CodeEnvironment.find({ + 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) - .lean(); - for (const environment of pendingRegistrations) { + .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) { - if (!token) continue; + if (!token) { + await deferRegistration(); + continue; + } try { await revokeCodeBridgeWorker({ baseURL: environment.baseURL, @@ -85,6 +151,7 @@ export async function reconcileCodeEnvironmentLifecycle({ fetchImpl, }); } catch { + await deferRegistration(); continue; } } @@ -94,7 +161,7 @@ export async function reconcileCodeEnvironmentLifecycle({ }); await CodeEnvironment.deleteOne({ _id: environment._id, - registrationPendingAt: environment.registrationPendingAt, + registrationLeaseId: leaseId, }); } diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 9b231e140ab..2fa5945c84d 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -261,8 +261,12 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') async function completeCodeEnvironmentRegistration(id: string | Types.ObjectId): Promise { const result = await model().updateOne( - { _id: id, registrationPendingAt: { $exists: true } }, - { $unset: { registrationPendingAt: 1 } }, + { + _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'); diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index 74b807cff5e..3901a28a31c 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -68,6 +68,9 @@ const codeEnvironmentSchema: Schema = new Schema Date: Tue, 1 Sep 2026 14:22:05 -0400 Subject: [PATCH 20/29] fix: close code environment UI review findings --- .../Nav/Settings/CodeEnvironments.tsx | 10 +++-- client/src/data-provider/CodeEnvironments.ts | 18 ++++---- packages/api/src/code/http.spec.ts | 45 +++++++++++++++++++ packages/api/src/code/http.ts | 9 +++- packages/data-schemas/src/methods/agent.ts | 2 +- 5 files changed, 70 insertions(+), 14 deletions(-) diff --git a/client/src/components/Nav/Settings/CodeEnvironments.tsx b/client/src/components/Nav/Settings/CodeEnvironments.tsx index a51538b96ef..6b7204f2ba1 100644 --- a/client/src/components/Nav/Settings/CodeEnvironments.tsx +++ b/client/src/components/Nav/Settings/CodeEnvironments.tsx @@ -43,9 +43,13 @@ function WorkerCommand({ pairing }: { pairing: CodeEnvironmentPairingResponse['p variant="outline" size="sm" className="self-start" - onClick={() => { - void navigator.clipboard.writeText(command); - showToast({ message: localize('com_ui_copy_to_clipboard'), status: 'success' }); + onClick={async () => { + try { + await navigator.clipboard.writeText(command); + showToast({ message: localize('com_ui_copied_to_clipboard'), status: 'success' }); + } catch { + showToast({ message: localize('com_ui_copy_failed'), status: 'error' }); + } }} > {localize('com_ui_copy_to_clipboard')} diff --git a/client/src/data-provider/CodeEnvironments.ts b/client/src/data-provider/CodeEnvironments.ts index e1e45065315..07c3ca65ea2 100644 --- a/client/src/data-provider/CodeEnvironments.ts +++ b/client/src/data-provider/CodeEnvironments.ts @@ -31,14 +31,14 @@ export function usePairCodeEnvironmentMutation() { export function useDeleteCodeEnvironmentMutation() { const queryClient = useQueryClient(); - return useMutation( - [MutationKeys.deleteCodeEnvironment], - dataService.deleteCodeEnvironment, - { - onSuccess: () => { - queryClient.invalidateQueries([QueryKeys.codeEnvironments]); - queryClient.invalidateQueries([QueryKeys.endpoints]); - }, + return useMutation< + { environment: TCodeEnvironmentsResponse['environments'][number] }, + Error, + string + >([MutationKeys.deleteCodeEnvironment], dataService.deleteCodeEnvironment, { + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.codeEnvironments]); + queryClient.invalidateQueries([QueryKeys.endpoints]); }, - ); + }); } diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 2c4f1afdd3a..4e66d1a7ccb 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -283,6 +283,51 @@ describe('code environment HTTP handlers', () => { ); }); + 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 }) => { diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index f7ebcf0d04c..67106e0fbfb 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -311,7 +311,14 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) } const workerId = createEnvironmentId(); - if (!(await principalIsActive(principal.userId.toString()))) { + 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; diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index cddcce996d2..d418add3c95 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -843,7 +843,7 @@ export function createAgentMethods( const directEnvironmentId = updateData.code_environment_id; const setEnvironmentId = typeof updateData.$set === 'object' && updateData.$set != null - ? (updateData.$set as Record).code_environment_id + ? (updateData.$set as { code_environment_id?: unknown }).code_environment_id : undefined; let nextEnvironmentId: string | undefined; if (typeof directEnvironmentId === 'string') { From 8c207b391bb9634e782f4aea907926c92581bbb5 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 14:23:38 -0400 Subject: [PATCH 21/29] fix: satisfy code environment type checks --- packages/api/src/code/config.ts | 4 +++- packages/api/src/code/environments.integration.spec.ts | 1 + packages/api/src/code/http.spec.ts | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/packages/api/src/code/config.ts b/packages/api/src/code/config.ts index 3273bc19bf5..9bc344cffc8 100644 --- a/packages/api/src/code/config.ts +++ b/packages/api/src/code/config.ts @@ -130,7 +130,9 @@ export async function mergeAccessibleCodeEnvironments({ } let mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments].map( (environment) => - !isExecutableCodeEnvironment(environment) && environment.default === true + !isExecutableCodeEnvironment(environment) && + 'default' in environment && + environment.default === true ? { ...environment, default: false as const } : environment, ); diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 9031b44eb1e..cf30112c2b8 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -158,6 +158,7 @@ describe('code environment registry', () => { type: 'attached', baseURL: 'https://code.example.com', workerId: 'owner-worker', + controlPlaneId: 'shared-code-api', workerPrincipal: { type: 'user', id: ownerId.toString() }, }, }); diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 215700cf4ce..bd3d9b332cf 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -179,7 +179,7 @@ describe('code environment HTTP handlers', () => { } as AppConfig; const effectiveConfig = { endpoints: { [EModelEndpoint.agents]: { statefulCodeSessions: { environments: [] } } }, - } as AppConfig; + } as unknown as AppConfig; const getAppConfig = jest.fn(async (options) => options.baseOnly === true ? baseConfig : effectiveConfig, ); From 3ccd22d1c5cebef396bdd252e0c4cc6dbda28bfa Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 18:55:05 -0400 Subject: [PATCH 22/29] fix: route code priming to selected worker --- api/app/clients/tools/util/handleTools.js | 1 + .../clients/tools/util/handleTools.test.js | 33 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 961ffbc4ef3..cc5ba646626 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 6ceaa578cfb..4b9f38e70bc 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}`; From 848c6e87de75acbb5974d74503cc4d66fe23ee81 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:10:06 -0400 Subject: [PATCH 23/29] fix: harden code environment lifecycle fences --- .../src/code/environments.integration.spec.ts | 46 ++++++++++ packages/api/src/code/http.spec.ts | 90 +++++++++++++++++++ packages/api/src/code/http.ts | 25 ++++++ packages/api/src/code/lifecycle.ts | 42 +++++++-- .../data-schemas/src/methods/agent.spec.ts | 49 +++++++++- packages/data-schemas/src/methods/agent.ts | 67 +++++++------- .../src/methods/codeEnvironment.ts | 54 +++++++++++ .../src/schema/codeEnvironment.ts | 1 + .../data-schemas/src/types/codeEnvironment.ts | 1 + 9 files changed, 332 insertions(+), 43 deletions(-) diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 8676de372a0..4c1abfc1c18 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -487,10 +487,56 @@ describe('code environment registry', () => { 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', diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index 4e66d1a7ccb..b3b0c04e25c 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -752,6 +752,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({ diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 67106e0fbfb..663670d1037 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -223,6 +223,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, @@ -236,6 +247,20 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) 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 = diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 1b53b1e2f0f..54517e384f4 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -8,6 +8,7 @@ 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; @@ -178,7 +179,24 @@ export async function reconcileCodeEnvironmentLifecycle({ await methods.deleteCodeEnvironmentById(environment._id); } - const candidates = await CodeEnvironment.find({ revocationPendingAt: { $exists: true } }) + 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>>(); @@ -189,9 +207,19 @@ export async function reconcileCodeEnvironmentLifecycle({ { _id: candidate._id, revocationPendingAt: { $exists: true }, - $or: [ - { revocationLeaseExpiresAt: { $exists: false } }, - { revocationLeaseExpiresAt: { $lte: now } }, + $and: [ + { + $or: [ + { revocationReconcileAfter: { $exists: false } }, + { revocationReconcileAfter: { $lte: now } }, + ], + }, + { + $or: [ + { revocationLeaseExpiresAt: { $exists: false } }, + { revocationLeaseExpiresAt: { $lte: now } }, + ], + }, ], }, { @@ -227,7 +255,10 @@ export async function reconcileCodeEnvironmentLifecycle({ await CodeEnvironment.updateOne( { _id: environment._id, revocationLeaseId: leaseId }, { - $set: { revocationLastError: message }, + $set: { + revocationLastError: message, + revocationReconcileAfter: new Date(Date.now() + REVOCATION_RETRY_MS), + }, $unset: { revocationLeaseId: 1, revocationLeaseExpiresAt: 1 }, }, ); @@ -331,6 +362,7 @@ export async function revokeUserCodeEnvironmentWorkers({ revocationPendingAt: 1, revocationAttempts: 1, revocationLastError: 1, + revocationReconcileAfter: 1, }, }, ); diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index 139fdb71c27..50349f3f006 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,52 @@ 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('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 d418add3c95..23508c44b9f 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -9,10 +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 { - releaseCodeEnvironmentReference, - reserveCodeEnvironmentReference, -} from './codeEnvironment'; +import { withCodeEnvironmentReference } from './codeEnvironment'; import { filterExistingSkillIds } from './skill'; import logger from '~/config/winston'; @@ -589,15 +586,11 @@ export function createAgentMethods( extractMCPServerNames(agentData.tools as string[] | undefined), }; - const reservation = await reserveCodeEnvironmentReference( + return await withCodeEnvironmentReference( mongoose, typeof agentData.code_environment_id === 'string' ? agentData.code_environment_id : undefined, + async () => (await Agent.create(initialAgentData)).toObject() as IAgent, ); - try { - return (await Agent.create(initialAgentData)).toObject() as IAgent; - } finally { - await releaseCodeEnvironmentReference(mongoose, reservation); - } } /** @@ -851,17 +844,16 @@ export function createAgentMethods( } else if (typeof setEnvironmentId === 'string') { nextEnvironmentId = setEnvironmentId; } - const reservation = await reserveCodeEnvironmentReference(mongoose, nextEnvironmentId); - let updatedAgent: IAgent | null; - try { - updatedAgent = (await Agent.findOneAndUpdate( - searchParameter, - updateData, - mongoOptions, - ).lean()) as IAgent | null; - } finally { - await releaseCodeEnvironmentReference(mongoose, reservation); - } + const updatedAgent = await withCodeEnvironmentReference( + mongoose, + nextEnvironmentId, + async () => + (await Agent.findOneAndUpdate( + searchParameter, + updateData, + mongoOptions, + ).lean()) as IAgent | null, + ); /** `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. @@ -1192,6 +1184,8 @@ export function createAgentMethods( if (includeSkillConfig) { projection.skills = 1; projection.skills_enabled = 1; + projection.skill_authoring_enabled = 1; + projection.skills_scope = 1; } let query = Agent.find(baseQuery, projection).sort({ updatedAt: -1, _id: 1 }); @@ -1273,27 +1267,26 @@ export function createAgentMethods( } } - const restoresDeploymentDefault = !Object.prototype.hasOwnProperty.call( - revertToVersion, - 'code_environment_id', - ); - const revertUpdate = restoresDeploymentDefault - ? { $set: revertToVersion, $unset: { code_environment_id: 1 } } - : { $set: revertToVersion }; - const reservation = await reserveCodeEnvironmentReference( + const unsetOnRestore: Record = {}; + for (const field of ['code_environment_id', 'skills_scope', 'skill_authoring_enabled']) { + if (!Object.prototype.hasOwnProperty.call(revertToVersion, field)) { + unsetOnRestore[field] = 1; + } + } + const revertUpdate = + Object.keys(unsetOnRestore).length > 0 + ? { $set: revertToVersion, $unset: unsetOnRestore } + : { $set: revertToVersion }; + const revertedAgent = await withCodeEnvironmentReference( mongoose, typeof revertToVersion.code_environment_id === 'string' ? revertToVersion.code_environment_id : undefined, + async () => + await Agent.findOneAndUpdate(searchParameter, revertUpdate, { + new: true, + }).lean(), ); - let revertedAgent: IAgent | null; - try { - revertedAgent = await Agent.findOneAndUpdate(searchParameter, revertUpdate, { - new: true, - }).lean(); - } finally { - await releaseCodeEnvironmentReference(mongoose, reservation); - } 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 2fa5945c84d..61eb6d84fb6 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -5,9 +5,11 @@ 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 = { @@ -107,6 +109,58 @@ export async function releaseCodeEnvironmentReference( ); } +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 = REFERENCE_RENEW_INTERVAL_MS, +): Promise { + const reservation = await reserveCodeEnvironmentReference(mongoose, environmentId); + if (reservation == null) return await operation(); + + let renewal = Promise.resolve(); + const timer = setInterval(() => { + renewal = renewal + .then(() => renewCodeEnvironmentReference(mongoose, reservation)) + .catch((error) => { + logger.error('[code-environments] agent reference lease renewal failed:', error); + }); + }, renewIntervalMs); + timer.unref(); + try { + return await operation(); + } finally { + clearInterval(timer); + await renewal; + await releaseCodeEnvironmentReference(mongoose, reservation); + } +} + type CreateCodeEnvironmentInput = Pick< CodeEnvironmentDocument, 'environmentId' | 'name' | 'type' | 'baseURL' | 'controlPlaneId' | 'createdBy' diff --git a/packages/data-schemas/src/schema/codeEnvironment.ts b/packages/data-schemas/src/schema/codeEnvironment.ts index 3901a28a31c..ea320b3b867 100644 --- a/packages/data-schemas/src/schema/codeEnvironment.ts +++ b/packages/data-schemas/src/schema/codeEnvironment.ts @@ -81,6 +81,7 @@ const codeEnvironmentSchema: Schema = new Schema Date: Tue, 1 Sep 2026 19:12:45 -0400 Subject: [PATCH 24/29] fix: type lease renewal interval --- packages/data-schemas/src/methods/codeEnvironment.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index 61eb6d84fb6..d0379da443d 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -138,7 +138,7 @@ export async function withCodeEnvironmentReference( mongoose: typeof import('mongoose'), environmentId: string | undefined, operation: () => Promise, - renewIntervalMs = REFERENCE_RENEW_INTERVAL_MS, + renewIntervalMs: number = REFERENCE_RENEW_INTERVAL_MS, ): Promise { const reservation = await reserveCodeEnvironmentReference(mongoose, environmentId); if (reservation == null) return await operation(); From 0cab1e5d2d2fca191c919f5f690c4d5cea690062 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:23:23 -0400 Subject: [PATCH 25/29] fix: harden environment lifecycle recovery --- config/__tests__/delete-user.spec.js | 4 +- config/delete-user.js | 2 +- .../src/code/environments.integration.spec.ts | 40 +++++++++++++++++++ packages/api/src/code/lifecycle.ts | 38 ++++++++++++++++++ .../data-schemas/src/methods/agent.spec.ts | 40 +++++++++++++++++++ packages/data-schemas/src/methods/agent.ts | 30 ++++++++++++++ .../src/methods/codeEnvironment.ts | 16 +++++++- 7 files changed, 167 insertions(+), 3 deletions(-) diff --git a/config/__tests__/delete-user.spec.js b/config/__tests__/delete-user.spec.js index 8af07560157..3ea21b63e82 100644 --- a/config/__tests__/delete-user.spec.js +++ b/config/__tests__/delete-user.spec.js @@ -49,10 +49,12 @@ jest.mock('@librechat/api', () => ({ abortJob: (...args) => mockAbortJob(...args), destroy: (...args) => mockDestroy(...args), }, - getAppConfig: (...args) => mockGetAppConfig(...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, diff --git a/config/delete-user.js b/config/delete-user.js index 5fbd79a9c6c..333d50d9a76 100644 --- a/config/delete-user.js +++ b/config/delete-user.js @@ -34,11 +34,11 @@ require('module-alias')({ base: path.resolve(__dirname, '..', 'api') }); const { GenerationJobManager, createStreamServices, - getAppConfig, 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'); diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 4c1abfc1c18..ae61e60d327 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -551,6 +551,46 @@ describe('code environment registry', () => { ).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 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(); + }); + test('removes creator-owned environment records and grants when the user is deleted', async () => { const registry = createCodeEnvironmentRegistry(mongoose); const methods = createMethods(mongoose); diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index 54517e384f4..da5ff5cc0fa 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -179,6 +179,44 @@ export async function reconcileCodeEnvironmentLifecycle({ 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 orphanedUserEnvironments = await CodeEnvironment.aggregate<{ + _id: CodeEnvironmentDocument['_id']; + }>([ + { + $match: { + 'workerPrincipal.type': 'user', + workerId: { $exists: true }, + 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 } }, + ]); + if (orphanedUserEnvironments.length > 0) { + await CodeEnvironment.updateMany( + { _id: { $in: orphanedUserEnvironments.map(({ _id }) => _id) } }, + { + $set: { revocationPendingAt: now }, + $inc: { revocationAttempts: 1 }, + $unset: { revocationLastError: 1, revocationReconcileAfter: 1 }, + }, + ); + } + const candidates = await CodeEnvironment.find({ revocationPendingAt: { $exists: true }, $and: [ diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index 50349f3f006..d1ef70c1fa1 100644 --- a/packages/data-schemas/src/methods/agent.spec.ts +++ b/packages/data-schemas/src/methods/agent.spec.ts @@ -2642,6 +2642,46 @@ describe('Agent Methods', () => { ).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 guardedWrite = withCodeEnvironmentReference( + mongoose, + environmentId, + async () => + await new Promise((resolve) => { + settleWrite = resolve; + }), + 10, + ); + 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', + }); + }); + 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 23508c44b9f..c69567f8244 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -590,6 +590,10 @@ export function createAgentMethods( 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 }); + }, ); } @@ -853,6 +857,17 @@ export function createAgentMethods( updateData, mongoOptions, ).lean()) as IAgent | null, + undefined, + async (agentAfterUpdate) => { + if (agentAfterUpdate == null || nextEnvironmentId == null) return; + const previousEnvironmentId = currentAgent?.code_environment_id; + await Agent.updateOne( + { _id: agentAfterUpdate._id, code_environment_id: nextEnvironmentId }, + previousEnvironmentId == null + ? { $unset: { code_environment_id: 1 } } + : { $set: { code_environment_id: previousEnvironmentId } }, + ); + }, ); /** `version` is a response-only field holding the count of `versions`. It is reported @@ -1286,6 +1301,21 @@ export function createAgentMethods( await Agent.findOneAndUpdate(searchParameter, revertUpdate, { new: true, }).lean(), + undefined, + async (agentAfterRevert) => { + if (agentAfterRevert == null || typeof revertToVersion.code_environment_id !== 'string') { + return; + } + await Agent.updateOne( + { + _id: agentAfterRevert._id, + code_environment_id: revertToVersion.code_environment_id, + }, + agent.code_environment_id == null + ? { $unset: { code_environment_id: 1 } } + : { $set: { code_environment_id: agent.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 d0379da443d..cb2b81e0c7a 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -139,21 +139,35 @@ export async function withCodeEnvironmentReference( 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 { - return await operation(); + const result = await operation(); + 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; From c7b2592d06e61ddfcee2d8674f5078b2e2b98a4f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:33:30 -0400 Subject: [PATCH 26/29] fix: close lifecycle recovery races --- .../__tests__/v1.duplicate-actions.spec.js | 43 +++++++++++++++ api/server/controllers/agents/v1.js | 13 ++++- .../src/code/environments.integration.spec.ts | 13 +++++ packages/api/src/code/lifecycle.ts | 21 +++++--- .../data-schemas/src/methods/agent.spec.ts | 3 ++ packages/data-schemas/src/methods/agent.ts | 52 ++++++++++++++----- .../src/methods/codeEnvironment.ts | 3 ++ 7 files changed, 127 insertions(+), 21 deletions(-) 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 cc298bd03ab..f0a4bf6095d 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 5949b7fca9a..20c1478697d 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -1571,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([ diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index ae61e60d327..1b0dbf4e29d 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -570,6 +570,16 @@ describe('code environment registry', () => { 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, @@ -589,6 +599,9 @@ describe('code environment registry', () => { 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 () => { diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index da5ff5cc0fa..daf595ebf69 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -182,13 +182,14 @@ export async function reconcileCodeEnvironmentLifecycle({ /** 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 orphanedUserEnvironments = await CodeEnvironment.aggregate<{ + const orphanedEnvironments = await CodeEnvironment.aggregate<{ _id: CodeEnvironmentDocument['_id']; + createdBy: CodeEnvironmentDocument['createdBy']; + workerId?: string; + workerPrincipal?: CodeEnvironmentDocument['workerPrincipal']; }>([ { $match: { - 'workerPrincipal.type': 'user', - workerId: { $exists: true }, revocationPendingAt: { $exists: false }, deletionCommittedAt: { $exists: false }, }, @@ -204,11 +205,14 @@ export async function reconcileCodeEnvironmentLifecycle({ { $match: { owner: { $size: 0 } } }, { $sort: { _id: 1 } }, { $limit: limit }, - { $project: { _id: 1 } }, + { $project: { _id: 1, createdBy: 1, workerId: 1, workerPrincipal: 1 } }, ]); - if (orphanedUserEnvironments.length > 0) { + const revocationTargets = orphanedEnvironments.filter( + ({ workerId, workerPrincipal }) => workerId != null && workerPrincipal?.type === 'user', + ); + if (revocationTargets.length > 0) { await CodeEnvironment.updateMany( - { _id: { $in: orphanedUserEnvironments.map(({ _id }) => _id) } }, + { _id: { $in: revocationTargets.map(({ _id }) => _id) } }, { $set: { revocationPendingAt: now }, $inc: { revocationAttempts: 1 }, @@ -216,6 +220,11 @@ export async function reconcileCodeEnvironmentLifecycle({ }, ); } + for (const creatorId of new Set( + orphanedEnvironments.map(({ createdBy }) => createdBy.toHexString()), + )) { + await methods.deleteUserCodeEnvironments(creatorId); + } const candidates = await CodeEnvironment.find({ revocationPendingAt: { $exists: true }, diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index d1ef70c1fa1..c0b78c022e2 100644 --- a/packages/data-schemas/src/methods/agent.spec.ts +++ b/packages/data-schemas/src/methods/agent.spec.ts @@ -2653,6 +2653,7 @@ describe('Agent Methods', () => { createdBy: new mongoose.Types.ObjectId(), }); let settleWrite: (() => void) | undefined; + const compensate = jest.fn().mockResolvedValue(undefined); const guardedWrite = withCodeEnvironmentReference( mongoose, environmentId, @@ -2661,6 +2662,7 @@ describe('Agent Methods', () => { settleWrite = resolve; }), 10, + compensate, ); for (let attempt = 0; attempt < 20; attempt++) { const environment = await mongoose.models.CodeEnvironment.findOne({ @@ -2680,6 +2682,7 @@ describe('Agent Methods', () => { await expect(guardedWrite).rejects.toMatchObject({ name: 'CodeEnvironmentReferenceError', }); + expect(compensate).toHaveBeenCalledWith(undefined); }); test('should prune deleted skill ids when reverting to an older version', async () => { diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index c69567f8244..f1299e05dbf 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -549,6 +549,33 @@ 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 restored = await Agent.replaceOne( + { + _id: agentAfterWrite._id, + code_environment_id: lostEnvironmentId, + updatedAt: agentAfterWrite.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. */ @@ -860,12 +887,12 @@ export function createAgentMethods( undefined, async (agentAfterUpdate) => { if (agentAfterUpdate == null || nextEnvironmentId == null) return; - const previousEnvironmentId = currentAgent?.code_environment_id; - await Agent.updateOne( - { _id: agentAfterUpdate._id, code_environment_id: nextEnvironmentId }, - previousEnvironmentId == null - ? { $unset: { code_environment_id: 1 } } - : { $set: { code_environment_id: previousEnvironmentId } }, + if (currentAgent == null) return; + await restoreAgentAfterReferenceLoss( + Agent, + agentAfterUpdate, + currentAgent.toObject() as IAgent, + nextEnvironmentId, ); }, ); @@ -1306,14 +1333,11 @@ export function createAgentMethods( if (agentAfterRevert == null || typeof revertToVersion.code_environment_id !== 'string') { return; } - await Agent.updateOne( - { - _id: agentAfterRevert._id, - code_environment_id: revertToVersion.code_environment_id, - }, - agent.code_environment_id == null - ? { $unset: { code_environment_id: 1 } } - : { $set: { code_environment_id: agent.code_environment_id } }, + await restoreAgentAfterReferenceLoss( + Agent, + agentAfterRevert, + agent.toObject() as IAgent, + revertToVersion.code_environment_id, ); }, ); diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index cb2b81e0c7a..eb51c7a428d 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -157,6 +157,9 @@ export async function withCodeEnvironmentReference( 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 { From 868d662f95fcd5da6bf8cca436ea1c10c3e3cac5 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:36:12 -0400 Subject: [PATCH 27/29] fix: type compensated agent timestamps --- packages/data-schemas/src/methods/agent.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index f1299e05dbf..2da57c55e97 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -556,11 +556,12 @@ export function createAgentMethods( 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: agentAfterWrite.updatedAt, + updatedAt, }, originalAgent, { timestamps: false }, From 03ce61730b652d27d9f3e0fe23225859eae8a6de Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:48:45 -0400 Subject: [PATCH 28/29] fix: fence lifecycle cleanup by tenant and revision --- .../src/code/environments.integration.spec.ts | 52 +++++++++++++ packages/api/src/code/environments.ts | 30 +++++++- packages/api/src/code/http.spec.ts | 73 +++++++++++++++++++ packages/api/src/code/http.ts | 14 ++++ packages/api/src/code/lifecycle.ts | 13 +++- .../data-schemas/src/methods/agent.spec.ts | 57 +++++++++++++++ packages/data-schemas/src/methods/agent.ts | 14 +++- .../src/methods/codeEnvironment.ts | 9 ++- 8 files changed, 254 insertions(+), 8 deletions(-) diff --git a/packages/api/src/code/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index 1b0dbf4e29d..47477bfe1dd 100644 --- a/packages/api/src/code/environments.integration.spec.ts +++ b/packages/api/src/code/environments.integration.spec.ts @@ -880,6 +880,8 @@ describe('code environment registry', () => { 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'); @@ -1123,4 +1125,54 @@ describe('code environment registry', () => { }); }); }); + + 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 55f613e6d42..0720644d879 100644 --- a/packages/api/src/code/environments.ts +++ b/packages/api/src/code/environments.ts @@ -62,6 +62,13 @@ export type CodeEnvironmentLifecycleTarget = CodeEnvironmentSummary & { 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)?$/; @@ -167,6 +174,7 @@ export function createCodeEnvironmentRegistry( ) => Promise; listRegisteredIds: () => Promise; invalidateAccessibleConfigurations: (tenantId?: string) => Promise; + markRevocationPending: (environmentId: string) => Promise; remove: (params: { actor: CodeEnvironmentPrincipalContext; environmentId: string; @@ -317,6 +325,22 @@ export function createCodeEnvironmentRegistry( }); } + 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( actor: CodeEnvironmentPrincipalContext, ): Promise { @@ -467,7 +491,10 @@ export function createCodeEnvironmentRegistry( let externalLifecycleStarted = false; try { const Agent = mongoose.models.Agent; - if (Agent != null && (await Agent.exists({ code_environment_id: environmentId })) != null) { + if ( + Agent != null && + (await Agent.exists(agentReferenceFilter(environmentId, environment.tenantId))) != null + ) { throw new CodeEnvironmentInUseError(environmentId); } if (beforeDelete != null) { @@ -503,6 +530,7 @@ export function createCodeEnvironmentRegistry( return { register, + markRevocationPending, listAccessible, listAccessibleConfigurations, listRegisteredIds, diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index b3b0c04e25c..ce255adc110 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -410,6 +410,79 @@ describe('code environment HTTP handlers', () => { ); }); + 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.at(-1) ?? 0, + ); + }); + test('revokes an issued pairing when the post-issue principal check is unavailable', async () => { const register = jest.fn(); const principalIsActive = jest diff --git a/packages/api/src/code/http.ts b/packages/api/src/code/http.ts index 663670d1037..60c0766a8c2 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -43,6 +43,7 @@ type Registry = { environmentId: string; beforeDelete?: (target: CodeEnvironmentLifecycleTarget) => Promise; }) => Promise; + markRevocationPending?: (environmentId: string) => Promise; }; type StatefulCodeConfig = NonNullable< @@ -387,6 +388,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) : res.status(409).json({ error: 'Account deletion is already in progress' }); } + let registrationCommitted = false; try { const environment = await deps.registry.register({ actor: principal, @@ -402,6 +404,7 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) workerPrincipal: { type: 'user', id: principal.userId.toString() }, }, }); + registrationCommitted = true; let activeAfterRegistration = false; try { activeAfterRegistration = await principalIsActive(principal.userId.toString()); @@ -441,6 +444,17 @@ export function createCodeEnvironmentHttpHandlers(deps: CodeEnvironmentHttpDeps) }, }); } 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, diff --git a/packages/api/src/code/lifecycle.ts b/packages/api/src/code/lifecycle.ts index daf595ebf69..f82de80e0bd 100644 --- a/packages/api/src/code/lifecycle.ts +++ b/packages/api/src/code/lifecycle.ts @@ -12,6 +12,13 @@ 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, @@ -48,7 +55,9 @@ export async function reconcileCodeEnvironmentLifecycle({ const Agent = mongoose.models.Agent; if ( Agent != null && - (await Agent.exists({ code_environment_id: environment.environmentId })) != null + (await Agent.exists( + agentReferenceFilter(environment.environmentId, environment.tenantId), + )) != null ) { await methods.cancelCodeEnvironmentRemoval(environment._id, leaseId); continue; @@ -139,7 +148,7 @@ export async function reconcileCodeEnvironmentLifecycle({ }; const tokenEnv = environment.revocationTokenEnv; const token = tokenEnv != null ? readSecret(tokenEnv)?.trim() : undefined; - if (environment.workerId != null) { + if (environment.workerId != null && environment.workerPrincipal?.type !== 'deployment') { if (!token) { await deferRegistration(); continue; diff --git a/packages/data-schemas/src/methods/agent.spec.ts b/packages/data-schemas/src/methods/agent.spec.ts index c0b78c022e2..017713cf248 100644 --- a/packages/data-schemas/src/methods/agent.spec.ts +++ b/packages/data-schemas/src/methods/agent.spec.ts @@ -2685,6 +2685,63 @@ describe('Agent Methods', () => { 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 2da57c55e97..e41773d676c 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -722,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; @@ -881,7 +882,9 @@ export function createAgentMethods( nextEnvironmentId, async () => (await Agent.findOneAndUpdate( - searchParameter, + currentAgent == null + ? searchParameter + : { ...searchParameter, _id: currentAgent._id, updatedAt: currentRevision }, updateData, mongoOptions, ).lean()) as IAgent | null, @@ -1288,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; @@ -1326,9 +1330,11 @@ export function createAgentMethods( ? revertToVersion.code_environment_id : undefined, async () => - await Agent.findOneAndUpdate(searchParameter, revertUpdate, { - new: true, - }).lean(), + 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') { diff --git a/packages/data-schemas/src/methods/codeEnvironment.ts b/packages/data-schemas/src/methods/codeEnvironment.ts index eb51c7a428d..cb6a563293a 100644 --- a/packages/data-schemas/src/methods/codeEnvironment.ts +++ b/packages/data-schemas/src/methods/codeEnvironment.ts @@ -46,6 +46,13 @@ function ownerSlotId(ownerId: Types.ObjectId, slot: number): Types.ObjectId { 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; @@ -442,7 +449,7 @@ export function createCodeEnvironmentMethods(mongoose: typeof import('mongoose') if (claimed == null || claimed.deletionLeaseId == null) continue; const referenced = Agent != null && - (await Agent.exists({ code_environment_id: claimed.environmentId })) != null; + (await Agent.exists(agentReferenceFilter(claimed.environmentId, claimed.tenantId))) != null; if (referenced) { await cancelCodeEnvironmentRemoval(candidate._id, claimed.deletionLeaseId); continue; From dce0d37bb531f24d9e7b827a853990f4d0b29306 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 1 Sep 2026 19:52:33 -0400 Subject: [PATCH 29/29] fix: scope environment revision fencing --- packages/api/src/code/http.spec.ts | 2 +- packages/data-schemas/src/methods/agent.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/api/src/code/http.spec.ts b/packages/api/src/code/http.spec.ts index ce255adc110..4e4efd131e6 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -479,7 +479,7 @@ describe('code environment HTTP handlers', () => { expect(res.statusCode).toBe(500); expect(markRevocationPending).toHaveBeenCalledWith('code-generated'); expect(markRevocationPending.mock.invocationCallOrder[0]).toBeLessThan( - fetchImpl.mock.invocationCallOrder.at(-1) ?? 0, + fetchImpl.mock.invocationCallOrder[fetchImpl.mock.invocationCallOrder.length - 1] ?? 0, ); }); diff --git a/packages/data-schemas/src/methods/agent.ts b/packages/data-schemas/src/methods/agent.ts index e41773d676c..8e5e1385640 100644 --- a/packages/data-schemas/src/methods/agent.ts +++ b/packages/data-schemas/src/methods/agent.ts @@ -882,7 +882,7 @@ export function createAgentMethods( nextEnvironmentId, async () => (await Agent.findOneAndUpdate( - currentAgent == null + currentAgent == null || nextEnvironmentId == null ? searchParameter : { ...searchParameter, _id: currentAgent._id, updatedAt: currentRevision }, updateData,