diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 254c917fe98..961ffbc4ef3 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( + (bridgeWorkerId) => getCodeApiAuthHeaders(options.req, bridgeWorkerId), + codeExecutionContext, + ), ...codeExecutionContext, }); }; 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/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/middleware/limiters/code.js b/api/server/middleware/limiters/code.js new file mode 100644 index 00000000000..301ae08cfcc --- /dev/null +++ b/api/server/middleware/limiters/code.js @@ -0,0 +1,41 @@ +const rateLimit = require('express-rate-limit'); +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 = 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, + 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..cfb4747c14d --- /dev/null +++ b/api/server/middleware/limiters/code.spec.js @@ -0,0 +1,82 @@ +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('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'); + + 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 316ffe1d09a..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,6 +21,9 @@ const requireCodeEnvironmentManage = requireCapability(SystemCapabilities.MANAGE router.use(requireJwtAuth); router.get('/', (req, res, next) => getHandlers().list(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 1be0bbca3a4..d483920040a 100644 --- a/api/server/routes/code-environments.test.js +++ b/api/server/routes/code-environments.test.js @@ -11,11 +11,16 @@ 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 = { 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', () => ({ @@ -30,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', () => ({ @@ -74,4 +82,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', 'pairing-limit']); + expect(mockHandlers.pair).toHaveBeenCalledTimes(1); + }); }); 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/__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'), diff --git a/api/server/services/Files/Code/crud.js b/api/server/services/Files/Code/crud.js index 27986ca2158..c7e63540d21 100644 --- a/api/server/services/Files/Code/crud.js +++ b/api/server/services/Files/Code/crud.js @@ -13,7 +13,7 @@ const { getCodeApiAuthHeaders, getCodeExecutionBaseUrl, createCodeExecutionRouteKey, - CODE_API_EXPECTED_PROFILE_HEADER, + codeExecutionHeaders, } = require('@librechat/api'); const axios = createAxiosInstance(); @@ -29,7 +29,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. @@ -38,7 +38,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', @@ -48,7 +48,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, @@ -82,7 +85,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 = @@ -122,12 +124,18 @@ 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: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, - [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, + ...codeExecutionHeaders({ + executionProfile, + bridgeWorkerId, + }), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -186,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. @@ -199,6 +208,7 @@ async function uploadCodeEnvFile({ version, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }) { try { const form = new FormData(); @@ -206,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: { @@ -215,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, @@ -267,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. */ @@ -279,6 +290,7 @@ async function batchUploadCodeEnvFiles({ read_only = false, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }) { const form = new FormData(); appendCodeEnvFileIdentity(form, { kind, id, version }); @@ -290,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: { @@ -299,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 c4a5dec4573..4e5c1242dbe 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 }), @@ -66,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(() => { @@ -118,12 +127,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', + }), }), ); }); @@ -136,7 +149,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'); }); @@ -195,7 +208,7 @@ describe('Code CRUD', () => { await deleteCodeEnvFile(req, file); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(req, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'delete', @@ -430,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'); }); @@ -447,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 @@ -552,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/preflight.js b/api/server/services/Files/Code/preflight.js index b985b039bf3..156e364bab6 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, executionRouteKey: codeExecutionContext?.executionRouteKey, }), onInspectionUnavailable: (index) => { diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index e3ad5b6c061..479386d4590 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -26,11 +26,11 @@ const { getStorageMetadata, getCodeExecutionBaseUrl, buildCodeEnvDownloadQuery, + codeExecutionHeaders, claimCodeDestination, createCodeDestinationSet, CODE_OUTPUT_PREFLIGHT_MAX_BYTES, CODE_OUTPUT_PREFLIGHT_MAX_COUNT, - CODE_API_EXPECTED_PROFILE_HEADER, sortCodeFilesByDestinationPriority, } = require('@librechat/api'); const { @@ -218,9 +218,10 @@ const downloadCodeOutputBuffer = async ({ maxBytes, codeApiBaseUrl, executionProfile = 'default', + 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 { @@ -231,7 +232,7 @@ const downloadCodeOutputBuffer = async ({ headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, - [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, + ...codeExecutionHeaders({ executionProfile, bridgeWorkerId }), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -272,6 +273,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. * @param {string} [params.executionRouteKey] - Trusted deployment-local route identity. */ const prepareCodeOutputForInspection = async ({ @@ -283,6 +285,7 @@ const prepareCodeOutputForInspection = async ({ inspectContent = true, codeApiBaseUrl, executionProfile = 'default', + bridgeWorkerId, executionRouteKey, }) => { const { fileSizeLimit } = getCodeOutputFileSettings(req); @@ -295,6 +298,7 @@ const prepareCodeOutputForInspection = async ({ maxBytes: transportLimit, codeApiBaseUrl, executionProfile, + bridgeWorkerId, }); cachePreparedCodeOutputBuffer({ req, @@ -655,6 +659,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 @@ -674,6 +679,7 @@ const processCodeOutput = async ({ codeApiBaseUrl, executionProfile = 'default', executionRouteKey = executionProfile, + bridgeWorkerId, preparedBuffer, downloadFallback, }) => { @@ -711,6 +717,7 @@ const processCodeOutput = async ({ maxBytes: fileSizeLimit, codeApiBaseUrl, executionProfile, + bridgeWorkerId, })); // Enforce file size limit @@ -1123,7 +1130,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} @@ -1133,7 +1140,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 @@ -1152,7 +1159,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, @@ -1283,8 +1293,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 ?? []; @@ -1431,6 +1442,7 @@ const primeFiles = async (options) => { ...(sourceRef.kind === 'skill' ? { version: sourceRef.version } : {}), codeApiBaseUrl, executionProfile, + bridgeWorkerId, }); /** @@ -1580,6 +1592,7 @@ async function readSandboxFile({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, req, }) { const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); @@ -1605,7 +1618,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`, @@ -1614,7 +1627,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, @@ -1663,6 +1676,7 @@ async function readSandboxFile({ * @param {number} [params.maxBytes] - In-sandbox size cap; larger files return `{ tooLarge, bytes }`. * @param {ServerRequest} [params.req] - Current authenticated request, used to mint Code API auth. * @param {string} [params.executionRouteKey] - Trusted deployment-local route identity. + * @param {string} [params.bridgeWorkerId] - Trusted bridge worker selected for this execution. * @returns {Promise<{base64: string, bytes: number} * | {tooLarge: true, reason: 'size' | 'round_trips', bytes: number} | null>} * `null` when codeapi is unavailable; throws on transport / read errors. @@ -1674,6 +1688,7 @@ async function readSandboxImage({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, executionRouteKey, maxBytes, req, @@ -1716,6 +1731,7 @@ async function readSandboxImage({ session_id, runtime_session_hint, executionProfile, + bridgeWorkerId, files, req, rateLimit, @@ -1737,6 +1753,7 @@ async function execSandboxImageChunk({ session_id, runtime_session_hint, executionProfile, + bridgeWorkerId, files, req, rateLimit, @@ -1762,7 +1779,7 @@ async function execSandboxImageChunk({ `[readSandboxImage] Rate-limited reading "${file_path}"; retrying in ${waitMs}ms`, ), attempt: async () => { - const authHeaders = await getCodeApiAuthHeaders(req); + const authHeaders = await getCodeApiAuthHeaders(req, bridgeWorkerId); return axios({ method: 'post', url: `${baseURL}/exec`, @@ -1771,7 +1788,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, @@ -1811,6 +1828,7 @@ async function writeSandboxFile({ runtime_session_hint, codeApiBaseUrl, executionProfile, + bridgeWorkerId, req, }) { const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); @@ -1854,7 +1872,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`, @@ -1863,7 +1881,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 80c7d0aadca..15b7fa9d3b7 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -88,6 +88,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), /** @@ -561,7 +565,7 @@ describe('Code Process', () => { await processCodeOutput(baseParams); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'get', @@ -1038,13 +1042,17 @@ describe('Code Process', () => { ...baseParams, codeApiBaseUrl: 'https://code-stateful.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', executionRouteKey, }); 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'); @@ -1458,7 +1466,7 @@ describe('Code Process', () => { mockReq, ); - expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq); + expect(getCodeApiAuthHeaders).toHaveBeenCalledWith(mockReq, undefined); expect(mockAxios).toHaveBeenCalledWith( expect.objectContaining({ method: 'get', @@ -1907,12 +1915,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'); }); @@ -1978,7 +1988,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', @@ -2455,6 +2465,7 @@ describe('Code Process', () => { agentId: 'agent-id', codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }); expect(mockAxios).not.toHaveBeenCalled(); @@ -2462,6 +2473,7 @@ describe('Code Process', () => { expect.objectContaining({ codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', }), ); expect(updateFile).toHaveBeenCalledWith( diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 357a3255d87..c5d4ac0950e 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, @@ -1393,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; @@ -1692,7 +1696,11 @@ async function loadAgentTools({ deferredToolsEnabled, programmaticToolsEnabled, codeExecutionEnabled, - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), codeExecutionContext, }); @@ -2077,7 +2085,11 @@ async function loadToolsForExecution({ */ for (const name of ptcToolNames) { const ptcTool = createBashProgrammaticToolCallingTool({ - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), baseUrl: codeExecutionContext.baseUrl, executionProfile: codeExecutionContext.executionProfile, runtimeSessionHint: codeExecutionContext.runtimeSessionHint, @@ -2103,7 +2115,11 @@ async function loadToolsForExecution({ if (isBashTool) { try { const bashTool = createBashExecutionTool({ - authHeaders: () => getCodeApiAuthHeaders(req), + authHeaders: () => + codeExecutionAuthHeaders( + (bridgeWorkerId) => getCodeApiAuthHeaders(req, bridgeWorkerId), + codeExecutionContext, + ), ...codeExecutionContext, }); allLoadedTools.push(bashTool); 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/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..ca8c3320977 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,108 @@ 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('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('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, + 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', + }); + 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', () => { const environment = (id: string, baseURL: string) => ({ id, @@ -159,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 9bab6cb5510..c4744eac6fb 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,16 +22,21 @@ export interface CodeExecutionContext { statefulSessions: boolean; environmentId?: string; environmentType?: CodeEnvironmentConfig['type']; + bridgeWorkerId?: string; } 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)}`; } @@ -102,14 +108,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: { @@ -152,11 +168,29 @@ export function resolveCodeExecutionContext(params: { statefulSessions: true, environmentId: configuredEnvironment?.id, environmentType: configuredEnvironment?.type, + bridgeWorkerId: configuredEnvironment?.workerId ?? configuredEnvironment?.pairing?.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: ( + bridgeWorkerId?: string, + ) => Promise> | Record, + context: Pick, +): Promise> { + return { + ...(await authHeaders(context.bridgeWorkerId)), + ...codeExecutionHeaders(context), + }; } diff --git a/packages/api/src/agents/handlers.spec.ts b/packages/api/src/agents/handlers.spec.ts index 12de41039db..32973bab5d5 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', executionRouteKey: 'stateful:deployment-a', runtimeSessionHint: 'v1:user', statefulSessions: true, @@ -4450,6 +4451,7 @@ describe('createToolExecuteHandler', () => { expect.objectContaining({ codeApiBaseUrl: 'https://stateful-code.example.com', executionProfile: 'stateful', + bridgeWorkerId: 'personal-worker-1', executionRouteKey: 'stateful:deployment-a', runtime_session_hint: 'v1:user', }), diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 23ddeacd9e4..3d27e38013a 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -389,6 +389,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 }>; @@ -400,6 +401,7 @@ export interface ToolExecuteOptions { route?: { baseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; }, ) => Promise; /** 23-hour freshness check */ @@ -450,6 +452,7 @@ export interface ToolExecuteOptions { runtime_session_hint?: string; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; executionRouteKey?: string; req?: ServerRequest; }) => Promise<{ content: string } | null>; @@ -471,6 +474,7 @@ export interface ToolExecuteOptions { runtime_session_hint?: string; codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; executionRouteKey?: string; /** In-sandbox size cap; files larger than this return `tooLarge` without transferring bytes. */ maxBytes?: number; @@ -556,6 +560,7 @@ function getCodeExecutionContext( function codeExecutionRequestParams(context?: CodeExecutionContext): { codeApiBaseUrl?: string; executionProfile?: CodeExecutionContext['executionProfile']; + bridgeWorkerId?: string; executionRouteKey?: string; runtime_session_hint?: string; } { @@ -566,6 +571,7 @@ function codeExecutionRequestParams(context?: CodeExecutionContext): { codeApiBaseUrl: context.baseUrl, executionProfile: context.executionProfile, ...(context.executionRouteKey ? { executionRouteKey: context.executionRouteKey } : {}), + ...(context.bridgeWorkerId ? { bridgeWorkerId: context.bridgeWorkerId } : {}), ...(context.runtimeSessionHint ? { runtime_session_hint: context.runtimeSessionHint } : {}), }; } 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, }); diff --git a/packages/api/src/agents/skillFiles.spec.ts b/packages/api/src/agents/skillFiles.spec.ts index 37d1f604b72..f7a666ed024 100644 --- a/packages/api/src/agents/skillFiles.spec.ts +++ b/packages/api/src/agents/skillFiles.spec.ts @@ -799,10 +799,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 8d8bda9c19b..6344c30b835 100644 --- a/packages/api/src/agents/skillFiles.ts +++ b/packages/api/src/agents/skillFiles.ts @@ -74,6 +74,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 }>; @@ -85,12 +86,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; @@ -514,6 +516,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/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 11353b5bede..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({ @@ -211,6 +223,31 @@ 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('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 24bf26bf391..1d2688a1af5 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; } @@ -167,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; @@ -218,6 +235,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'; @@ -255,13 +280,15 @@ 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 = 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); @@ -294,6 +321,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, }; } @@ -314,7 +342,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, @@ -325,8 +353,9 @@ 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(':'); + ]); } function pruneTokenCache(now: number): void { @@ -345,7 +374,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 ''; } @@ -353,7 +382,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 ( @@ -376,10 +405,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}` } : {}; } 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/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..9bc344cffc8 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,29 @@ export async function mergeAccessibleCodeEnvironments({ ) { return appConfig; } - const mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments]; + let mergedEnvironments = [...filteredEnvironments, ...effectivePrincipalEnvironments].map( + (environment) => + !isExecutableCodeEnvironment(environment) && + 'default' in 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/environments.integration.spec.ts b/packages/api/src/code/environments.integration.spec.ts index ae52b21bbad..cf30112c2b8 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', }, ]); }); @@ -143,6 +145,40 @@ 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', + controlPlaneId: 'shared-code-api', + 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 aced59bb9df..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 = { @@ -41,6 +45,7 @@ export type AccessibleCodeEnvironmentConfiguration = { baseURL: string; controlPlaneId: string; owner: 'principal'; + workerId?: string; }; type CachedAccessibleCodeEnvironmentConfiguration = AccessibleCodeEnvironmentConfiguration & { @@ -49,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'; @@ -64,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); @@ -73,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(); @@ -94,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 }; } @@ -181,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 { @@ -226,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( @@ -258,15 +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', - })); + 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 b9b4df55e51..bd3d9b332cf 100644 --- a/packages/api/src/code/http.spec.ts +++ b/packages/api/src/code/http.spec.ts @@ -37,6 +37,315 @@ 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'), + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), + 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', + 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) }), + }); + }); + + 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('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 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() }, + readSecret: jest.fn(() => 'administrator-token'), + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), + 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', + 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', @@ -95,25 +404,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', + }, }, ], }, @@ -126,20 +440,73 @@ 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'), + principalAuthEnabled: jest.fn(() => true), + principalAuthReady: jest.fn(), + 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: { @@ -159,7 +526,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(); @@ -171,7 +538,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 9cd73544889..b6c485759d8 100644 --- a/packages/api/src/code/http.ts +++ b/packages/api/src/code/http.ts @@ -9,7 +9,15 @@ import type { } from './environments'; import type { GetAppConfigOptions } from '~/app/service'; import type { ServerRequest } from '~/types/http'; -import { CodeEnvironmentValidationError } from './environments'; +import type { CodeBridgeFetch } from './bridge'; +import { + assertCodeApiJwtSigningReady, + getCodeApiTenantId, + isCodeApiJwtAuthEnabled, +} from '~/auth/codeapi'; +import { CodeBridgePairingError, createCodeBridgePairing, readCodeBridgeSecret } from './bridge'; +import { CodeEnvironmentValidationError, normalizeCodeEnvironmentName } from './environments'; +import { getAppConfigOptionsFromUser } from '~/app/service'; type Registry = { register: (params: { @@ -30,6 +38,11 @@ export interface CodeEnvironmentHttpDeps { getAppConfig: (options: GetAppConfigOptions) => Promise; registry: Registry; createEnvironmentId?: () => string; + readSecret?: (name: string) => string | undefined; + resolveTenantId?: (req: ServerRequest) => string; + principalAuthEnabled?: () => boolean; + principalAuthReady?: () => Promise | void; + fetchImpl?: CodeBridgeFetch; } function actor(req: ServerRequest): CodeEnvironmentPrincipalContext | null { @@ -50,15 +63,54 @@ function configuredControlPlane( environment.id === controlPlaneId && environment.type === 'attached' && environment.owner === 'deployment' && - environment.pairing != null, + environment.pairing?.workerId != null, ); } +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; + const principalAuthEnabled = deps.principalAuthEnabled ?? isCodeApiJwtAuthEnabled; + const principalAuthReady = deps.principalAuthReady ?? assertCodeApiJwtSigningReady; async function list(req: ServerRequest, res: Response): Promise { const principal = actor(req); @@ -86,12 +138,30 @@ 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. */ - const appConfig = await deps.getAppConfig({ baseOnly: true }); - 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' }); } @@ -105,6 +175,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 }); @@ -125,5 +196,116 @@ 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' }); + } + 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() : ''; + 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', + }); + } + 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' }); + } + + 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; + 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, + controlPlaneId: controlPlane.id, + workerPrincipal: { type: 'user', id: principal.userId.toString() }, + }, + }); + 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..c2bafd65078 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, }, ], @@ -220,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 4476ac2ce66..4e6eb140495 100644 --- a/packages/data-provider/specs/config-schemas.spec.ts +++ b/packages/data-provider/specs/config-schemas.spec.ts @@ -576,6 +576,75 @@ 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', + pairing: { + allowPrincipalWorkers: true, + tokenEnv: 'CODE_BRIDGE_ADMIN_TOKEN', + }, + }, + ], + }, + }); + + 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: { @@ -599,6 +668,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 6549df3046b..a2fdbe96df6 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(), }), ) @@ -1103,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, @@ -1128,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, @@ -1136,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', 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