From 06bf324cf02b47793bd234a1ac797663dfa28112 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 16 Aug 2026 09:42:15 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A4=EF=B8=8F=20feat:=20Per-Agent=20Cod?= =?UTF-8?q?e=20Execution=20Routing=20With=20Stateful=20Session=20Scopes=20?= =?UTF-8?q?(#14848)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: route code execution per agent profile * chore: sort execution profile imports * test: preserve stateful environment literal types * fix: isolate stateful code environments by user * fix: preserve per-agent code routing end to end * fix: route code priming by execution profile * fix: isolate code profile lifecycle state * fix: preserve mixed-profile code resources * fix: complete stateful skill routing --- .env.example | 6 +- api/app/clients/BaseClient.js | 1 + api/app/clients/tools/util/handleTools.js | 24 +- .../agents/__tests__/callbacks.spec.js | 9 +- api/server/controllers/agents/callbacks.js | 4 + api/server/controllers/agents/client.js | 12 +- api/server/controllers/agents/openai.js | 3 + api/server/controllers/agents/responses.js | 4 + api/server/routes/files/files.js | 14 ++ api/server/routes/files/files.test.js | 44 ++++ .../services/Endpoints/agents/initialize.js | 53 +++- .../Endpoints/agents/initialize.spec.js | 46 +++- .../services/Endpoints/agents/skillDeps.js | 1 + .../Code/__tests__/process-traversal.spec.js | 2 + api/server/services/Files/Code/crud.js | 83 +++++-- api/server/services/Files/Code/crud.spec.js | 99 ++++++++ api/server/services/Files/Code/process.js | 141 ++++++++--- .../services/Files/Code/process.spec.js | 190 ++++++++++++++- api/server/services/Files/process.js | 19 +- api/server/services/Files/process.spec.js | 30 +++ api/server/services/ToolService.js | 77 +++++- .../services/__tests__/ToolService.spec.js | 195 ++++++++++++++- client/src/common/agents-types.ts | 3 + .../Agents/Advanced/StatefulSessions.tsx | 106 +++++--- .../SidePanel/Agents/AgentPanel.tsx | 3 + .../SidePanel/Agents/AgentSelect.tsx | 1 + .../__tests__/AgentPanel.helpers.spec.ts | 21 ++ client/src/locales/en/translation.json | 7 +- .../src/agents/__tests__/initialize.test.ts | 43 ++++ .../agents/__tests__/run-codeTools.test.ts | 12 +- .../api/src/agents/codeFilesSession.spec.ts | 227 ++++++++++++++++++ packages/api/src/agents/codeFilesSession.ts | 125 ++++++++-- packages/api/src/agents/execution.spec.ts | 97 ++++++++ packages/api/src/agents/execution.ts | 105 ++++++++ .../src/agents/handlers.background.spec.ts | 10 +- packages/api/src/agents/handlers.spec.ts | 104 ++++++++ packages/api/src/agents/handlers.ts | 189 +++++++++++++-- packages/api/src/agents/harvest.ts | 7 + packages/api/src/agents/index.ts | 1 + packages/api/src/agents/initialize.ts | 49 +++- packages/api/src/agents/lazySubagents.ts | 3 + packages/api/src/agents/openai/service.ts | 8 +- packages/api/src/agents/prewarm.spec.ts | 118 ++++++++- packages/api/src/agents/prewarm.ts | 127 ++++++++-- packages/api/src/agents/resources.ts | 2 +- packages/api/src/agents/run.ts | 63 +---- packages/api/src/agents/skillFiles.spec.ts | 162 ++++++++++++- packages/api/src/agents/skillFiles.ts | 107 ++++++++- .../src/agents/statefulCodeSessions.spec.ts | 53 ---- packages/api/src/agents/tools.ts | 4 +- packages/api/src/agents/validation.spec.ts | 20 ++ packages/api/src/agents/validation.ts | 4 + packages/api/src/skills/deployment.ts | 20 +- packages/api/src/tools/classification.ts | 19 +- packages/data-provider/src/codeEnvRef.spec.ts | 50 +++- packages/data-provider/src/codeEnvRef.ts | 63 +++++ packages/data-provider/src/schemas.ts | 1 + .../data-provider/src/types/assistants.ts | 6 + packages/data-provider/src/types/files.ts | 3 +- .../data-schemas/src/methods/file.spec.ts | 32 +++ packages/data-schemas/src/methods/file.ts | 10 +- .../data-schemas/src/methods/skill.spec.ts | 57 ++++- packages/data-schemas/src/methods/skill.ts | 22 +- packages/data-schemas/src/schema/agent.ts | 4 + .../data-schemas/src/schema/codeEnvRef.ts | 28 +++ packages/data-schemas/src/schema/file.ts | 20 +- packages/data-schemas/src/schema/skillFile.ts | 20 +- packages/data-schemas/src/types/agent.ts | 1 + packages/data-schemas/src/types/file.ts | 3 +- packages/data-schemas/src/types/skill.ts | 3 +- 70 files changed, 2782 insertions(+), 418 deletions(-) create mode 100644 packages/api/src/agents/execution.spec.ts create mode 100644 packages/api/src/agents/execution.ts delete mode 100644 packages/api/src/agents/statefulCodeSessions.spec.ts create mode 100644 packages/data-schemas/src/schema/codeEnvRef.ts diff --git a/.env.example b/.env.example index 07ff7d0a0c..84906756e3 100644 --- a/.env.example +++ b/.env.example @@ -621,7 +621,11 @@ TTS_API_KEY= # LIBRECHAT_CODE_API_KEY= # LIBRECHAT_CODE_BASEURL= -# Prewarm stateful per-conversation sandboxes in parallel with model generation (default: true). +# Optional dedicated Code API deployment for agents with Stateful code sessions enabled. +# When configured, stateless agents continue using LIBRECHAT_CODE_BASEURL while stateful +# agents fail closed onto this endpoint. The endpoint must advertise the `stateful` profile. +# LIBRECHAT_CODE_BASEURL_STATEFUL= +# Prewarm selected stateful sandboxes in parallel with model generation (default: true). # CODE_SANDBOX_PREWARM=true # Time in milliseconds before LibreChat treats a tracked sandbox as cold (default: 2100000 / 35 minutes). # CODE_SANDBOX_COLD_AFTER_MS=2100000 diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js index 602c7a680c..ae8714bf64 100644 --- a/api/app/clients/BaseClient.js +++ b/api/app/clients/BaseClient.js @@ -1444,6 +1444,7 @@ class BaseClient { if ( file.embedded === true || file.metadata?.codeEnvRef != null || + file.metadata?.codeEnvRefs != null || file.metadata?.fileIdentifier != null ) { allFiles.push(file); diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index dce01d274b..5a85c0d224 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -20,6 +20,7 @@ const { ASK_USER_QUESTION_TOOL_NAME, resolveWebSearchSSRFAgents, buildWebSearchDynamicContext, + resolveCodeExecutionContext, } = require('@librechat/api'); const { Tools, @@ -331,9 +332,23 @@ const loadTools = async ({ for (const tool of tools) { if (tool === Tools.execute_code) { requestedTools[tool] = async () => { + const statefulSessions = + agent?.stateful_code_sessions === true && + (await checkCapability(options.req, AgentCapabilities.stateful_code_sessions)); + const codeExecutionContext = + options.codeExecutionContext ?? + resolveCodeExecutionContext({ + statefulSessions, + environment: agent?.stateful_code_environment, + userId: user, + agentId: agent?.id, + conversationId: options.req?.body?.conversationId, + }); const { files, toolContext } = await primeCodeFiles({ ...options, agentId: agent?.id, + codeApiBaseUrl: codeExecutionContext.baseUrl, + executionProfile: codeExecutionContext.executionProfile, }); if (toolContext) { dynamicToolContextMap[tool] = toolContext; @@ -341,18 +356,11 @@ const loadTools = async ({ if (files?.length) { primedCodeFiles = files; } - /* Hedge the execute_code description toward persistence only when the - * admin `stateful_code_sessions` capability is on AND the agent opted - * in via the builder (off by default); the matching wire hint is set - * in the run config. Older @librechat/agents ignore the param. */ - const statefulSessions = - agent?.stateful_code_sessions === true && - (await checkCapability(options.req, AgentCapabilities.stateful_code_sessions)); return createCodeExecutionTool({ user_id: user, files, authHeaders: () => getCodeApiAuthHeaders(options.req), - statefulSessions, + ...codeExecutionContext, }); }; continue; diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index a183a212f5..d37f393367 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -444,6 +444,7 @@ describe('createToolEndCallback', () => { name, toolName = 'execute_code', hostFileAuthoring = false, + codeExecutionContext, }) { return { output: { @@ -455,7 +456,7 @@ describe('createToolEndCallback', () => { files: [{ id: fileId, name, session_id: 'sess-1' }], }, }, - metadata: { run_id: runId, thread_id: threadId }, + metadata: { run_id: runId, thread_id: threadId, codeExecutionContext }, }; } @@ -679,6 +680,10 @@ describe('createToolEndCallback', () => { name: 'created.txt', toolName: 'create_file', hostFileAuthoring: true, + codeExecutionContext: { + baseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', + }, }); await toolEndCallback({ output: event.output }, event.metadata); await Promise.all(artifactPromises); @@ -690,6 +695,8 @@ describe('createToolEndCallback', () => { messageId: 'run-create', toolCallId: 'tool-create', conversationId: 'thread789', + codeApiBaseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', }), ); expect(res.write).toHaveBeenCalledTimes(1); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index f72c0d61fa..a43b387d73 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -964,6 +964,8 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo * ids. */ session_id: file.storage_session_id ?? output.artifact.session_id, + codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, + executionProfile: metadata.codeExecutionContext?.executionProfile, }); const fileMetadata = result?.file ?? null; const finalize = result?.finalize; @@ -1286,6 +1288,8 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises }) * ids. */ session_id: file.storage_session_id ?? output.artifact.session_id, + codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl, + executionProfile: metadata.codeExecutionContext?.executionProfile, }); const fileMetadata = result?.file ?? null; const finalize = result?.finalize; diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index a44bddf3b1..8c5d4883d7 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -1415,7 +1415,7 @@ class AgentClient extends BaseClient { this.contextHandlers?.processFile(file); continue; } - if (file.metadata?.codeEnvRef) { + if (file.metadata?.codeEnvRef || file.metadata?.codeEnvRefs) { continue; } } @@ -2585,7 +2585,7 @@ class AgentClient extends BaseClient { abortController = new AbortController(); } - /** Fire-and-forget: boot the per-conversation stateful sandbox in + /** Fire-and-forget: boot each selected stateful environment in * parallel with generation so the first execute_code/bash call lands * on a warm VM. No-op unless a reachable agent resolved * `statefulCodeSessions`. */ @@ -2631,13 +2631,7 @@ class AgentClient extends BaseClient { ? await this.options.primeInvokedSkills(payload) : undefined; - /** - * Seed `Graph.sessions` with code-env files primed across every - * reachable agent (primary, handoff/addedConvo, and nested - * subagents) plus skill-priming output. The merge logic and its - * run-wide semantics live in `buildInitialToolSessions`; see that - * helper's doc for why this is intentionally NOT per-agent. - */ + /** Seed each reachable agent's trusted code-session partition. */ const initialSessions = buildInitialToolSessions({ skillSessions: skillPrimeResult?.initialSessions, agents: [this.options.agent, ...(this.agentConfigs ? this.agentConfigs.values() : [])], diff --git a/api/server/controllers/agents/openai.js b/api/server/controllers/agents/openai.js index ed647e1e62..e8a66e4368 100644 --- a/api/server/controllers/agents/openai.js +++ b/api/server/controllers/agents/openai.js @@ -87,6 +87,7 @@ function createToolLoader(signal, definitionsOnly = true) { provider, tool_options, tool_resources, + codeExecutionContext, accessibleMcpServerNames, }) { const agent = { id: agentId, tools, provider, model, tool_options }; @@ -97,6 +98,7 @@ function createToolLoader(signal, definitionsOnly = true) { agent, signal, tool_resources, + codeExecutionContext, agentResourceType: ResourceType.REMOTE_AGENT, definitionsOnly, accessibleMcpServerNames, @@ -508,6 +510,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => { req, res, agentResourceType: ResourceType.REMOTE_AGENT, + conversationId, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index b104800949..a8e6301c39 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -101,6 +101,7 @@ function createToolLoader(signal, definitionsOnly = true) { provider, tool_options, tool_resources, + codeExecutionContext, accessibleMcpServerNames, }) { const agent = { id: agentId, tools, provider, model, tool_options }; @@ -111,6 +112,7 @@ function createToolLoader(signal, definitionsOnly = true) { agent, signal, tool_resources, + codeExecutionContext, agentResourceType: ResourceType.REMOTE_AGENT, definitionsOnly, accessibleMcpServerNames, @@ -735,6 +737,7 @@ const executeResponse = async (envelope, { req, res }) => { req, res, agentResourceType: ResourceType.REMOTE_AGENT, + conversationId, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, @@ -919,6 +922,7 @@ const executeResponse = async (envelope, { req, res }) => { req, res, agentResourceType: ResourceType.REMOTE_AGENT, + conversationId, toolNames, agent: ctx.agent ?? agent, signal: abortController.signal, diff --git a/api/server/routes/files/files.js b/api/server/routes/files/files.js index 458d9ec1d0..72295c4a9b 100644 --- a/api/server/routes/files/files.js +++ b/api/server/routes/files/files.js @@ -10,6 +10,7 @@ const { startUploadSseStream, resolveUploadErrorMessage, verifyAgentUploadPermission, + getCodeExecutionBaseUrl, } = require('@librechat/api'); const { Time, @@ -329,6 +330,18 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => { return res.status(400).send('Bad request'); } + const requestedProfile = req.query.execution_profile; + if ( + requestedProfile != null && + requestedProfile !== 'default' && + requestedProfile !== 'stateful' + ) { + logger.debug(`${logPrefix} invalid execution_profile`); + return res.status(400).send('Bad request'); + } + const executionProfile = requestedProfile ?? 'default'; + const baseUrl = getCodeExecutionBaseUrl(executionProfile); + const { getDownloadStream } = getStrategyFunctions(FileSources.execute_code); if (!getDownloadStream) { logger.warn( @@ -352,6 +365,7 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => { id: req.user.id, }, req, + { baseUrl, executionProfile }, ); res.set(response.headers); response.data.pipe(res); diff --git a/api/server/routes/files/files.test.js b/api/server/routes/files/files.test.js index d087fc508d..4fce9cb30e 100644 --- a/api/server/routes/files/files.test.js +++ b/api/server/routes/files/files.test.js @@ -45,6 +45,11 @@ jest.mock('sharp', () => jest.mock('@librechat/api', () => ({ ...jest.requireActual('@librechat/api'), refreshS3FileUrls: jest.fn(), + getCodeExecutionBaseUrl: jest.fn((profile) => + profile === 'stateful' + ? process.env.LIBRECHAT_CODE_BASEURL_STATEFUL + : 'https://code-default.example.com/v1', + ), })); jest.mock('~/cache', () => ({ @@ -1088,4 +1093,43 @@ describe('File Routes - Delete with Agent Access', () => { expect(response.status).toBe(401); }); }); + + describe('GET /files/code/download/:session_id/:fileId', () => { + it('routes a persisted stateful fallback through the stateful Code API', async () => { + const getDownloadStream = jest.fn().mockResolvedValue({ + headers: { 'content-type': 'text/plain' }, + data: Readable.from(['stateful output']), + }); + getStrategyFunctions.mockReturnValue({ getDownloadStream }); + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'https://code-stateful.example.com/v1'; + + try { + const sessionId = 's'.repeat(21); + const codeFileId = 'f'.repeat(21); + const response = await request(app).get( + `/files/code/download/${sessionId}/${codeFileId}?execution_profile=stateful`, + ); + + expect(response.status).toBe(200); + expect(response.text).toBe('stateful output'); + expect(getDownloadStream).toHaveBeenCalledWith( + `${sessionId}/${codeFileId}`, + { kind: 'user', id: otherUserId.toString() }, + expect.any(Object), + { baseUrl: 'https://code-stateful.example.com/v1', executionProfile: 'stateful' }, + ); + } finally { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + } + }); + + it('rejects an unknown execution profile', async () => { + const response = await request(app).get( + `/files/code/download/${'s'.repeat(21)}/${'f'.repeat(21)}?execution_profile=attacker`, + ); + + expect(response.status).toBe(400); + expect(getStrategyFunctions).not.toHaveBeenCalled(); + }); + }); }); diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 00e43956de..51d9bcf6eb 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -1,11 +1,11 @@ const { logger } = require('@librechat/data-schemas'); -const { createContentAggregator } = require('@librechat/agents'); +const { createContentAggregator, GraphNodeKeys } = require('@librechat/agents'); const { checkAccess, loadSkillStates, initializeAgent, isMemoryEnabled, - primeInvokedSkills, + primeInvokedSkillsForProfiles, validateAgentModel, extractManualSkills, GenerationJobManager, @@ -17,6 +17,7 @@ const { resolveModelSpecSkillIds, getAgentStartupTelemetry, buildAgentContextAttachmentsByAgentId, + collectCodeExecutionProfileRoutes, getLazySubagentConfigId, } = require('@librechat/api'); const { @@ -97,6 +98,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false, jobC provider, tool_options, tool_resources, + codeExecutionContext, accessibleMcpServerNames, }) { const agent = { id: agentId, tools, provider, model, tool_options }; @@ -109,6 +111,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false, jobC streamId, jobCreatedAt, tool_resources, + codeExecutionContext, definitionsOnly, accessibleMcpServerNames, }); @@ -167,7 +170,7 @@ const initializeClient = async ({ /** @type {Map} */ const toolInputValidationErrors = new Map(); const { contentParts, aggregateContent, stepMap } = createContentAggregator(); - const toolEndCallback = createToolEndCallback({ + const artifactToolEndCallback = createToolEndCallback({ req, res, artifactPromises, @@ -279,6 +282,30 @@ const initializeClient = async ({ * }>} */ const agentToolContexts = new Map(); + /** Attach only the host-resolved route for the actually executing agent. + * Runnable metadata is transport data and may contain caller-controlled + * keys, so discard any incoming route context before resolving from the + * server-owned per-agent map. This covers both traditional TOOL_END events + * and event-driven ON_TOOL_EXECUTE callbacks. */ + const toolEndCallback = async (data, metadata = {}) => { + const node = typeof metadata.langgraph_node === 'string' ? metadata.langgraph_node : ''; + const nodeAgentId = node.startsWith(GraphNodeKeys.TOOLS) + ? node.slice(GraphNodeKeys.TOOLS.length) + : undefined; + const executingAgentId = + metadata.executingAgentId ?? metadata.agentId ?? metadata.agent_id ?? nodeAgentId; + const soleContext = + agentToolContexts.size === 1 ? agentToolContexts.values().next().value : null; + const trustedContext = + (typeof executingAgentId === 'string' ? agentToolContexts.get(executingAgentId) : null) ?? + soleContext; + const callbackMetadata = { ...metadata }; + delete callbackMetadata.codeExecutionContext; + if (trustedContext?.codeExecutionContext) { + callbackMetadata.codeExecutionContext = trustedContext.codeExecutionContext; + } + return artifactToolEndCallback(data, callbackMetadata); + }; /** @type {Map} */ const endpointTokenConfigByAgentId = new Map(); @@ -293,6 +320,7 @@ const initializeClient = async ({ res, signal, streamId, + conversationId, toolNames, agent: ctx.agent, toolRegistry: ctx.toolRegistry, @@ -796,6 +824,7 @@ const initializeClient = async ({ codeEnvAvailable === true && agent.stateful_code_sessions === true && agent.tools?.includes(Tools.execute_code) === true, + statefulCodeEnvironment: agent.stateful_code_environment, includeReasoningHistory: getIncludeReasoningHistory(agent), }); @@ -977,6 +1006,7 @@ const initializeClient = async ({ configId: metadata.configId, codeEnvAvailable: metadata.codeEnvAvailable, statefulCodeSessions: metadata.statefulCodeSessions, + statefulCodeEnvironment: metadata.statefulCodeEnvironment, includeReasoningHistory: metadata.includeReasoningHistory, lazySubagentConfigs: lazyChildren, subagentAgentConfigs: eagerChildren, @@ -1057,16 +1087,23 @@ const initializeClient = async ({ /** History priming uses the user's full ACL-accessible skill set (not * per-agent scoped) because prior turns may reference skills no longer - * in any active agent's scope; the ACL check is the security gate. - * `codeEnvAvailable` comes from `primaryConfig` — @see - * `InitializedAgent.codeEnvAvailable` for the per-agent narrowing. */ + * in any active agent's scope; the ACL check is the security gate. Each + * selected Code API deployment receives its own upload, and only session + * partitions routed to that deployment receive those storage pointers. */ + const codeExecutionProfiles = collectCodeExecutionProfileRoutes( + [primaryConfig, ...agentConfigs.values()], + { + userId: req.user.id, + conversationId, + }, + ); const handlePrimeInvokedSkills = skillsCapabilityEnabled ? (payload) => - primeInvokedSkills({ + primeInvokedSkillsForProfiles({ req, payload, accessibleSkillIds, - codeEnvAvailable: primaryConfig.codeEnvAvailable === true, + executionProfiles: codeExecutionProfiles, ...getSkillToolDeps(), }) : undefined; diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 9d846be38e..a8ff3840f1 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -45,8 +45,9 @@ jest.mock('@librechat/api', () => ({ * the tool context (agent, tool_resources, skill ACLs) was preserved. */ let capturedToolExecuteOptions; let capturedDefaultHandlerOptions; +const mockArtifactToolEndCallback = jest.fn(); jest.mock('~/server/controllers/agents/callbacks', () => ({ - createToolEndCallback: jest.fn(() => jest.fn()), + createToolEndCallback: jest.fn(() => mockArtifactToolEndCallback), createAttachmentEmitter: jest.fn(() => jest.fn()), createBackgroundCodeResultHandler: jest.fn(() => jest.fn()), getDefaultHandlers: jest.fn((opts) => { @@ -158,6 +159,42 @@ describe('initializeClient — processAgent ACL gate', () => { tool_resources: {}, resendFiles: true, maxContextTokens: 4096, + codeExecutionContext: { + baseUrl: 'https://code-default.example.com', + codeSessionKey: 'execute_code', + executionProfile: 'default', + statefulSessions: false, + }, + }); + + it('replaces untrusted artifact route metadata with the executing agent context', async () => { + mockInitializeAgent.mockResolvedValue(makePrimaryConfig([])); + + await initializeClient({ + req: makeReq(), + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + + const data = { output: { name: 'execute_code', artifact: { files: [] } } }; + await capturedDefaultHandlerOptions.toolEndCallback(data, { + langgraph_node: `tools=${PRIMARY_ID}`, + codeExecutionContext: { + baseUrl: 'https://attacker.invalid', + executionProfile: 'stateful', + }, + }); + + expect(mockArtifactToolEndCallback).toHaveBeenLastCalledWith( + data, + expect.objectContaining({ + codeExecutionContext: expect.objectContaining({ + baseUrl: 'https://code-default.example.com', + executionProfile: 'default', + }), + }), + ); }); it('threads the owning job epoch into resumable event handlers', async () => { @@ -766,6 +803,7 @@ describe('initializeClient — subagent loading', () => { model: 'gpt-4', author: new mongoose.Types.ObjectId(), tools: ['web'], + stateful_code_environment: 'agent-user', }); await grantView(subAgent); @@ -795,7 +833,11 @@ describe('initializeClient — subagent loading', () => { expect(mockInitializeAgent).toHaveBeenCalledTimes(1); expect(agentClientArgs.agent.lazySubagentConfigs).toHaveLength(1); expect(agentClientArgs.agent.lazySubagentConfigs[0]).toEqual( - expect.objectContaining({ id: SUBAGENT_ID, configId: expect.any(String) }), + expect.objectContaining({ + id: SUBAGENT_ID, + configId: expect.any(String), + statefulCodeEnvironment: 'agent-user', + }), ); expect(agentClientArgs.agent.lazySubagentConfigs[0]).not.toHaveProperty('tools'); expect(agentClientArgs.agent.lazySubagentConfigs[0]).not.toHaveProperty('tool_resources'); diff --git a/api/server/services/Endpoints/agents/skillDeps.js b/api/server/services/Endpoints/agents/skillDeps.js index d5542ab31d..2d68113892 100644 --- a/api/server/services/Endpoints/agents/skillDeps.js +++ b/api/server/services/Endpoints/agents/skillDeps.js @@ -290,6 +290,7 @@ function buildAgentToolContext({ agent, config }) { accessibleSkillIds: config.accessibleSkillIds, activeSkillNames: config.activeSkillNames, codeEnvAvailable: config.codeEnvAvailable, + codeExecutionContext: config.codeExecutionContext, skillAuthoringAvailable: config.skillAuthoringAvailable, fileAuthoringToolNames: config.fileAuthoringToolNames, skillPrimedIdsByName: 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 57609c545a..b0abc0a5c5 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,8 @@ jest.mock('@librechat/api', () => { flattenArtifactPath: mockFlattenArtifactPath, createAxiosInstance: jest.fn(() => mockAxios), getCodeApiAuthHeaders: jest.fn(async () => ({})), + getCodeExecutionBaseUrl: jest.fn(() => 'http://localhost:8000'), + CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', classifyCodeArtifact: jest.fn(() => 'other'), extractCodeArtifactText: jest.fn(async () => null), /* `processCodeOutput` calls this to derive the trust flag persisted diff --git a/api/server/services/Files/Code/crud.js b/api/server/services/Files/Code/crud.js index 553599d335..f04e5ff682 100644 --- a/api/server/services/Files/Code/crud.js +++ b/api/server/services/Files/Code/crud.js @@ -1,6 +1,7 @@ const FormData = require('form-data'); const { logger } = require('@librechat/data-schemas'); const { getCodeBaseURL } = require('@librechat/agents'); +const { getCodeEnvRefs } = require('librechat-data-provider'); const { logAxiosError, appendCodeEnvFile, @@ -10,6 +11,8 @@ const { appendCodeEnvFileIdentity, buildCodeEnvDownloadQuery, getCodeApiAuthHeaders, + getCodeExecutionBaseUrl, + CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); const axios = createAxiosInstance(); @@ -24,12 +27,15 @@ const MAX_FILE_SIZE = 150 * 1024 * 1024; * matching sessionKey. For code-output downloads this is always * `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] + * 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. */ -async function getCodeOutputDownloadStream(fileIdentifier, identity, req) { +async function getCodeOutputDownloadStream(fileIdentifier, identity, req, route = {}) { try { - const baseURL = getCodeBaseURL(); + const baseURL = route.baseUrl ?? getCodeBaseURL(); const query = buildCodeEnvDownloadQuery(identity); const authHeaders = await getCodeApiAuthHeaders(req); /** @type {import('axios').AxiosRequestConfig} */ @@ -40,6 +46,9 @@ async function getCodeOutputDownloadStream(fileIdentifier, identity, req) { headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...(route.executionProfile + ? { [CODE_API_EXPECTED_PROFILE_HEADER]: route.executionProfile } + : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -66,26 +75,26 @@ async function getCodeOutputDownloadStream(fileIdentifier, identity, req) { * @returns {Promise} */ async function deleteCodeEnvFile(req, file) { - const ref = file?.metadata?.codeEnvRef; - if (!ref) { + const refs = getCodeEnvRefs(file?.metadata); + if (refs.length === 0) { return; } - let lastError; const missingOrUnsupportedStatuses = new Set([404, 405]); - try { - const baseURL = getCodeBaseURL(); + const authHeaders = await getCodeApiAuthHeaders(req); + for (const [executionProfile, ref] of refs) { + const baseURL = getCodeExecutionBaseUrl(executionProfile); const query = buildCodeEnvDownloadQuery({ kind: ref.kind, id: ref.id, ...(ref.kind === 'skill' ? { version: ref.version } : {}), }); - const authHeaders = await getCodeApiAuthHeaders(req); const baseRequest = { method: 'delete', headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, + [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -96,10 +105,13 @@ async function deleteCodeEnvFile(req, file) { `${baseURL}/files/${ref.storage_session_id}/${ref.file_id}${query}`, ]; + let lastError; + let deleted = false; for (const url of urls) { try { await axios({ ...baseRequest, url }); - return; + deleted = true; + break; } catch (error) { lastError = error; if (!missingOrUnsupportedStatuses.has(error.response?.status)) { @@ -107,19 +119,16 @@ async function deleteCodeEnvFile(req, file) { } } } - } catch (error) { - lastError = error; - } - - if (lastError) { - logAxiosError({ - error: lastError, - message: `Error deleting code environment file: ${lastError.message}`, - }); - if (lastError.response?.status === 404) { - return; + if (!deleted && lastError) { + logAxiosError({ + error: lastError, + message: `Error deleting code environment file: ${lastError.message}`, + }); + if (lastError.response?.status === 404) { + continue; + } + throw new Error(lastError.message || 'An error occurred during file deletion.'); } - throw new Error(lastError.message || 'An error occurred during file deletion.'); } } @@ -142,17 +151,28 @@ async function deleteCodeEnvFile(req, file) { * ignores this for `kind: 'user'` (auth context provides userId), but it's * sent uniformly for shape symmetry with the discriminated union. * @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. * @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. */ -async function uploadCodeEnvFile({ req, stream, filename, kind, id, version }) { +async function uploadCodeEnvFile({ + req, + stream, + filename, + kind, + id, + version, + codeApiBaseUrl, + executionProfile, +}) { try { const form = new FormData(); appendCodeEnvFileIdentity(form, { kind, id, version }); appendCodeEnvFile(form, stream, filename); - const baseURL = getCodeBaseURL(); + const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); const authHeaders = await getCodeApiAuthHeaders(req); /** @type {import('axios').AxiosRequestConfig} */ const options = { @@ -162,6 +182,7 @@ async function uploadCodeEnvFile({ req, stream, filename, kind, id, version }) { 'User-Agent': 'LibreChat/1.0', 'User-Id': req.user.id, ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -211,10 +232,21 @@ async function uploadCodeEnvFile({ req, stream, filename, kind, id, version }) { * through subsequent download/walk passes — sandboxed-code modifications * are dropped on the floor and the original ref is echoed back as * `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. * @returns {Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }> }>} * @throws {Error} If the batch upload fails entirely. */ -async function batchUploadCodeEnvFiles({ req, files, kind, id, version, read_only = false }) { +async function batchUploadCodeEnvFiles({ + req, + files, + kind, + id, + version, + read_only = false, + codeApiBaseUrl, + executionProfile, +}) { const form = new FormData(); appendCodeEnvFileIdentity(form, { kind, id, version }); if (read_only) { @@ -224,7 +256,7 @@ async function batchUploadCodeEnvFiles({ req, files, kind, id, version, read_onl appendCodeEnvFile(form, file.stream, file.filename); } - const baseURL = getCodeBaseURL(); + const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); const authHeaders = await getCodeApiAuthHeaders(req); /** @type {import('axios').AxiosRequestConfig} */ const options = { @@ -234,6 +266,7 @@ async function batchUploadCodeEnvFiles({ req, files, kind, id, version, read_onl 'User-Agent': 'LibreChat/1.0', 'User-Id': req.user.id, ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), }, 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 4e9c755ae1..7cc9e743fd 100644 --- a/api/server/services/Files/Code/crud.spec.js +++ b/api/server/services/Files/Code/crud.spec.js @@ -50,6 +50,10 @@ jest.mock('@librechat/api', () => { }), logAxiosError: jest.fn(({ message }) => message), getCodeApiAuthHeaders: jest.fn(async () => ({})), + getCodeExecutionBaseUrl: jest.fn((profile) => + profile === 'stateful' ? 'https://code-stateful.example.com' : 'https://code-api.example.com', + ), + CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', createAxiosInstance: jest.fn(() => mockAxios), codeServerHttpAgent: new http.Agent({ keepAlive: false }), codeServerHttpsAgent: new https.Agent({ keepAlive: false }), @@ -107,6 +111,22 @@ describe('Code CRUD', () => { expect(callConfig.timeout).toBe(15000); }); + it('uses the trusted stateful route and fail-closed profile header', async () => { + mockAxios.mockResolvedValue({ data: Readable.from(['chunk']) }); + + await getCodeOutputDownloadStream('session-1/file-1', userIdentity, undefined, { + baseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', + }); + + 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' }), + }), + ); + }); + it('forwards Code API auth headers when a request is provided', async () => { const req = { user: { id: 'user-123' } }; getCodeApiAuthHeaders.mockResolvedValue({ Authorization: 'Bearer codeapi-token' }); @@ -190,6 +210,65 @@ describe('Code CRUD', () => { ); }); + it('deletes a stateful artifact from its originating profile', async () => { + mockAxios.mockResolvedValue({ status: 204 }); + const statefulFile = { + metadata: { + codeEnvRef: { + ...file.metadata.codeEnvRef, + executionProfile: 'stateful', + }, + }, + }; + + await deleteCodeEnvFile(req, statefulFile); + + expect(mockAxios).toHaveBeenCalledWith( + expect.objectContaining({ + url: 'https://code-stateful.example.com/sessions/session-1/objects/file-1?kind=agent&id=agent-abc', + headers: expect.objectContaining({ + 'X-CodeAPI-Expected-Profile': 'stateful', + }), + }), + ); + }); + + it('deletes every profile-local object retained for a shared file record', async () => { + mockAxios.mockResolvedValue({ status: 204 }); + const dualProfileFile = { + metadata: { + codeEnvRef: file.metadata.codeEnvRef, + codeEnvRefs: { + default: file.metadata.codeEnvRef, + stateful: { + ...file.metadata.codeEnvRef, + storage_session_id: 'stateful-session', + file_id: 'stateful-file', + executionProfile: 'stateful', + }, + }, + }, + }; + + await deleteCodeEnvFile(req, dualProfileFile); + + expect(mockAxios).toHaveBeenCalledTimes(2); + expect(mockAxios).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + url: expect.stringContaining('/sessions/session-1/objects/file-1'), + headers: expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'default' }), + }), + ); + expect(mockAxios).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + url: expect.stringContaining('/sessions/stateful-session/objects/stateful-file'), + headers: expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'stateful' }), + }), + ); + }); + it.each([404, 405])( 'falls back to the legacy code environment delete route after a %s', async (status) => { @@ -317,6 +396,26 @@ describe('Code CRUD', () => { expect(callConfig.headers.Authorization).toBe('Bearer codeapi-token'); }); + it('routes uploads through the trusted stateful endpoint and profile header', async () => { + mockAxios.post.mockResolvedValue({ + data: { + message: 'success', + storage_session_id: 'sess-1', + files: [{ fileId: 'fid-1', filename: 'data.csv' }], + }, + }); + + await uploadCodeEnvFile({ + ...baseUploadParams, + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }); + + 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'); + }); + /* Phase C / option α (codeapi #1455): the upload wire carries the * resource identity codeapi uses for sessionKey derivation. Without * these on the form, codeapi falls back to user bucketing for every diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index 33b28a2d1d..053136a262 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -17,7 +17,9 @@ const { extractCodeArtifactText, getExtractedTextFormat, getStorageMetadata, + getCodeExecutionBaseUrl, buildCodeEnvDownloadQuery, + CODE_API_EXPECTED_PROFILE_HEADER, } = require('@librechat/api'); const { Tools, @@ -31,6 +33,9 @@ const { EModelEndpoint, ErrorTypes, mergeFileConfig, + getCodeEnvRefs, + mergeCodeEnvRef, + getCodeEnvRefForProfile, getEndpointFileConfig, } = require('librechat-data-provider'); const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions'); @@ -53,6 +58,7 @@ const axios = createAxiosInstance(); * @param {string} params.toolCallId - The tool call ID that generated the file. * @param {string} params.messageId - The current message ID. * @param {number} params.expiresAt - Expiration timestamp (24 hours from creation). + * @param {'default'|'stateful'} [params.executionProfile] - Code API route for later fallback download. * @returns {Object} Fallback response with download URL. */ const createDownloadFallback = ({ @@ -64,11 +70,13 @@ const createDownloadFallback = ({ session_id, toolCallId, conversationId, + executionProfile, }) => { const basePath = getBasePath(); + const profileQuery = executionProfile === 'stateful' ? '?execution_profile=stateful' : ''; return { filename: name, - filepath: `${basePath}/api/files/code/download/${session_id}/${id}`, + filepath: `${basePath}/api/files/code/download/${session_id}/${id}${profileQuery}`, expiresAt, conversationId, toolCallId, @@ -314,6 +322,8 @@ const runPreviewFinalize = ({ finalize, fileId, previewRevision, onResolved }) = * @param {string} params.session_id - The code execution session ID. * @param {string} params.conversationId - The current conversation ID. * @param {string} params.messageId - The current message ID. + * @param {string} [params.codeApiBaseUrl] - Trusted per-agent Code API endpoint. + * @param {'default'|'stateful'} [params.executionProfile] - Trusted execution profile. * @returns {Promise<{ file: MongoFile & { messageId: string, toolCallId: string }, finalize?: () => Promise }>} */ const processCodeOutput = async ({ @@ -326,10 +336,12 @@ const processCodeOutput = async ({ session_id, agentId, freshClaimAfter, + codeApiBaseUrl, + executionProfile = 'default', }) => { const appConfig = req.config; const currentDate = new Date(); - const baseURL = getCodeBaseURL(); + const baseURL = codeApiBaseUrl ?? getCodeExecutionBaseUrl(executionProfile); const fileExt = path.extname(name).toLowerCase(); const isImage = fileExt && imageExtRegex.test(name); @@ -356,6 +368,7 @@ const processCodeOutput = async ({ headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, + [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile, }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -378,6 +391,7 @@ const processCodeOutput = async ({ toolCallId, session_id, conversationId, + executionProfile, expiresAt: currentDate.getTime() + 86400000, }), }; @@ -391,6 +405,7 @@ const processCodeOutput = async ({ id: req.user.id, storage_session_id: session_id, file_id: id, + executionProfile, }; /* `safeName` keeps the directory structure (`a/b/file.txt` -> `a/b/file.txt`) @@ -504,6 +519,14 @@ const processCodeOutput = async ({ * silently excludes the file from priming on subsequent turns. */ const persistedMessageId = isUpdate ? (claimed.messageId ?? messageId) : messageId; + /* A generated-output write replaces the file's bytes, so pointers to + * earlier content in another profile must not survive as reusable refs. */ + const codeEnvReferenceSet = mergeCodeEnvRef(undefined, codeEnvRef); + const codeEnvMetadata = { + ...claimed.metadata, + ...codeEnvReferenceSet, + sourceDispatchedAt, + }; if (isImage) { const usage = isUpdate ? (claimed.usage ?? 0) + 1 : 1; @@ -524,6 +547,7 @@ const processCodeOutput = async ({ usage, filename: safeName, conversationId, + executionProfile, user: req.user.id, tenantId: req.user.tenantId, type: `image/${appConfig.imageOutputType}`, @@ -531,7 +555,7 @@ const processCodeOutput = async ({ updatedAt: formattedDate, source: appConfig.fileStrategy, context: FileContext.execute_code, - metadata: { codeEnvRef, sourceDispatchedAt }, + metadata: codeEnvMetadata, ...(await getRetentionExpiry(req)), }; if (!(await commitCodeFile(file))) { @@ -554,6 +578,7 @@ const processCodeOutput = async ({ toolCallId, session_id, conversationId, + executionProfile, expiresAt: currentDate.getTime() + 86400000, }), }; @@ -633,7 +658,7 @@ const processCodeOutput = async ({ tenantId: req.user.tenantId, bytes: buffer.length, updatedAt: formattedDate, - metadata: { codeEnvRef, sourceDispatchedAt }, + metadata: codeEnvMetadata, source: appConfig.fileStrategy, context: FileContext.execute_code, usage: isUpdate ? (claimed.usage ?? 0) + 1 : 1, @@ -731,6 +756,7 @@ const processCodeOutput = async ({ toolCallId, session_id, conversationId, + executionProfile, expiresAt: currentDate.getTime() + 86400000, }), }; @@ -752,14 +778,16 @@ 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] + * Trusted host-selected Code API route. * * @returns {Promise} * A promise that resolves to the `lastModified` time string of the file if successful, or null if there is an * error in initialization or fetching the info. */ -async function getSessionInfo(ref, req) { +async function getSessionInfo(ref, req, route = {}) { try { - const baseURL = getCodeBaseURL(); + const baseURL = route.baseUrl ?? getCodeBaseURL(); const authHeaders = await getCodeApiAuthHeaders(req); /* `/sessions/.../objects/...` is gated by codeapi's `sessionAuth` * middleware (post-Phase C). The middleware reconstructs the @@ -778,6 +806,9 @@ async function getSessionInfo(ref, req) { headers: { 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...(route.executionProfile + ? { [CODE_API_EXPECTED_PROFILE_HEADER]: route.executionProfile } + : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -886,7 +917,15 @@ const getReuploadFailureCategory = (error) => { * }>} */ const primeFiles = async (options) => { - const { tool_resources, req, agentId, agentResourceType } = options; + const { + tool_resources, + req, + agentId, + agentResourceType, + codeApiBaseUrl, + executionProfile = 'default', + } = options; + const codeApiRoute = { baseUrl: codeApiBaseUrl, executionProfile }; const file_ids = tool_resources?.[EToolResources.execute_code]?.file_ids ?? []; const agentResourceIds = new Set(file_ids); const resourceFiles = tool_resources?.[EToolResources.execute_code]?.files ?? []; @@ -943,15 +982,16 @@ const primeFiles = async (options) => { continue; } - const ref = file.metadata?.codeEnvRef; - if (!ref) { + const ref = getCodeEnvRefForProfile(file.metadata, executionProfile); + const sourceRef = ref ?? getCodeEnvRefs(file.metadata)[0]?.[1]; + if (!sourceRef) { skippedNoRef += 1; logger.debug(`[primeCodeFiles] file=${file.file_id} path=skip reason=no-codeenvref`); continue; } requiredCodeFiles += 1; - const session_id = ref.storage_session_id; - const id = ref.file_id; + const session_id = sourceRef.storage_session_id; + const id = sourceRef.file_id; /** * `pushFile` accepts optional overrides so the reupload path can @@ -982,22 +1022,14 @@ const primeFiles = async (options) => { * we still send it for shape uniformity with shared kinds. */ files.push({ id: overrideId ?? id, - resource_id: ref.id, + resource_id: sourceRef.id, storage_session_id: overrideSessionId ?? session_id, name: file.filename, - kind: ref.kind, - ...(ref.kind === 'skill' ? { version: ref.version } : {}), + kind: sourceRef.kind, + ...(sourceRef.kind === 'skill' ? { version: sourceRef.version } : {}), }); }; - if (sessions.has(session_id)) { - logger.debug( - `[primeCodeFiles] file=${file.file_id} path=cache-hit-by-session storage_session_id=${session_id}`, - ); - pushFile(); - continue; - } - const reuploadFile = async () => { try { const { getDownloadStream } = getStrategyFunctions(file.source); @@ -1014,9 +1046,11 @@ const primeFiles = async (options) => { req: options.req, stream, filename: file.filename, - kind: ref.kind, - id: ref.id, - ...(ref.kind === 'skill' ? { version: ref.version } : {}), + kind: sourceRef.kind, + id: sourceRef.id, + ...(sourceRef.kind === 'skill' ? { version: sourceRef.version } : {}), + codeApiBaseUrl, + executionProfile, }); /** @@ -1033,21 +1067,20 @@ const primeFiles = async (options) => { * pointer changes. */ const newRef = { - kind: ref.kind, - id: ref.id, + kind: sourceRef.kind, + id: sourceRef.id, storage_session_id: uploaded.storage_session_id, file_id: uploaded.file_id, - ...(ref.kind === 'skill' ? { version: ref.version } : {}), + executionProfile, + ...(sourceRef.kind === 'skill' ? { version: sourceRef.version } : {}), }; - const updatedMetadata = { - ...file.metadata, - codeEnvRef: newRef, - }; + const updatedRefs = mergeCodeEnvRef(file.metadata, newRef); await updateFile({ file_id: file.file_id, - metadata: updatedMetadata, + 'metadata.codeEnvRef': updatedRefs.codeEnvRef, + [`metadata.codeEnvRefs.${executionProfile}`]: newRef, }); sessions.set(newRef.storage_session_id, true); pushFile(newRef.storage_session_id, newRef.file_id); @@ -1066,7 +1099,22 @@ const primeFiles = async (options) => { ); } }; - const uploadTime = await getSessionInfo(ref, req); + if (!ref) { + logger.debug( + `[primeCodeFiles] file=${file.file_id} path=reupload reason=profile-missing ` + + `requestedProfile=${executionProfile}`, + ); + await reuploadFile(); + continue; + } + if (sessions.has(session_id)) { + logger.debug( + `[primeCodeFiles] file=${file.file_id} path=cache-hit-by-session storage_session_id=${session_id}`, + ); + pushFile(); + continue; + } + const uploadTime = await getSessionInfo(ref, req, codeApiRoute); if (!uploadTime) { logger.debug( `[primeCodeFiles] file=${file.file_id} path=reupload reason=no-uploadtime ` + @@ -1144,8 +1192,16 @@ const primeFiles = async (options) => { * @param {ServerRequest} [params.req] - Current authenticated request, used to mint Code API auth. * @returns {Promise<{content: string} | null>} */ -async function readSandboxFile({ file_path, session_id, files, runtime_session_hint, req }) { - const baseURL = getCodeBaseURL(); +async function readSandboxFile({ + file_path, + session_id, + files, + runtime_session_hint, + codeApiBaseUrl, + executionProfile, + req, +}) { + const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); if (!baseURL) { return null; } @@ -1177,6 +1233,7 @@ async function readSandboxFile({ file_path, session_id, files, runtime_session_h 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -1225,10 +1282,12 @@ async function readSandboxImage({ session_id, files, runtime_session_hint, + codeApiBaseUrl, + executionProfile, maxBytes, req, }) { - const baseURL = getCodeBaseURL(); + const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); if (!baseURL) { return null; } @@ -1287,6 +1346,7 @@ async function readSandboxImage({ file_path, session_id, runtime_session_hint, + executionProfile, files, req, chunkBytes, @@ -1357,6 +1417,7 @@ async function execSandboxImageChunk({ file_path, session_id, runtime_session_hint, + executionProfile, files, req, chunkBytes, @@ -1383,6 +1444,7 @@ async function execSandboxImageChunk({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), }, httpAgent: codeServerHttpAgent, httpsAgent: codeServerHttpsAgent, @@ -1448,9 +1510,11 @@ async function writeSandboxFile({ session_id, files, runtime_session_hint, + codeApiBaseUrl, + executionProfile, req, }) { - const baseURL = getCodeBaseURL(); + const baseURL = codeApiBaseUrl ?? getCodeBaseURL(); if (!baseURL) { return null; } @@ -1500,6 +1564,7 @@ async function writeSandboxFile({ 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...(executionProfile ? { [CODE_API_EXPECTED_PROFILE_HEADER]: executionProfile } : {}), }, 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 99e4f7a491..7f6c149fad 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -63,6 +63,10 @@ jest.mock('@librechat/api', () => { flattenArtifactPath: jest.fn((name) => name.replace(/\//g, '__')), createAxiosInstance: jest.fn(() => mockAxios), getCodeApiAuthHeaders: jest.fn(async () => ({})), + getCodeExecutionBaseUrl: jest.fn((profile) => + profile === 'stateful' ? 'https://code-stateful.example.com' : 'https://code-api.example.com', + ), + CODE_API_EXPECTED_PROFILE_HEADER: 'X-CodeAPI-Expected-Profile', withTimeout: (...args) => passthroughWithTimeout(...args), hasOfficeHtmlPath: (...args) => mockHasOfficeHtmlPath(...args), /** @@ -821,6 +825,24 @@ describe('Code Process', () => { }); describe('fallback behavior', () => { + it('preserves the stateful route in generated downloads and fallbacks', async () => { + mockAxios.mockRejectedValue(new Error('Network error')); + + const { file: result } = await processCodeOutput({ + ...baseParams, + codeApiBaseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', + }); + + expect(mockAxios).toHaveBeenCalledWith( + expect.objectContaining({ + url: expect.stringContaining('https://code-stateful.example.com/download/'), + headers: expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'stateful' }), + }), + ); + expect(result.filepath).toContain('execution_profile=stateful'); + }); + it('should fallback to download URL when saveBuffer is not available', async () => { const smallBuffer = Buffer.alloc(100); mockAxios.mockResolvedValue({ data: smallBuffer }); @@ -901,11 +923,33 @@ describe('Code Process', () => { id: 'user-123', storage_session_id: 'session-123', file_id: 'file-id-123', + executionProfile: 'default', + }, + codeEnvRefs: { + default: { + kind: 'user', + id: 'user-123', + storage_session_id: 'session-123', + file_id: 'file-id-123', + executionProfile: 'default', + }, }, sourceDispatchedAt: expect.any(Number), }); }); + it('persists the originating profile on a stateful artifact ref', async () => { + mockAxios.mockResolvedValue({ data: Buffer.alloc(100) }); + + const { file: result } = await processCodeOutput({ + ...baseParams, + codeApiBaseUrl: 'https://code-stateful.example.com', + executionProfile: 'stateful', + }); + + expect(result.metadata.codeEnvRef.executionProfile).toBe('stateful'); + }); + /* Phase C lock-in: outputs are ALWAYS user-scoped, never skill-scoped. * Even when an execution turn invoked a skill (so input files were * `kind: 'skill'` shared cross-user), the resulting output bucket @@ -935,12 +979,14 @@ describe('Code Process', () => { id: 'user-A', storage_session_id: 'session-123', file_id: 'file-id-123', + executionProfile: 'default', }); expect(outputB.metadata.codeEnvRef).toEqual({ kind: 'user', id: 'user-B', storage_session_id: 'session-123', file_id: 'file-id-123', + executionProfile: 'default', }); // No skill identity leaks into the output ref under any property. @@ -1208,6 +1254,35 @@ describe('Code Process', () => { }), ); }); + + it('checks freshness against the trusted stateful endpoint and profile', async () => { + mockAxios.mockResolvedValue({ + data: { lastModified: '2026-08-15T00:00:00Z' }, + }); + + await getSessionInfo( + { + kind: 'user', + id: 'user-123', + storage_session_id: 'session-123', + file_id: 'file-123', + }, + mockReq, + { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }, + ); + + expect(mockAxios).toHaveBeenCalledWith( + expect.objectContaining({ + url: expect.stringMatching(/^https:\/\/stateful-code\.example\.com\/sessions\//), + headers: expect.objectContaining({ + 'X-CodeAPI-Expected-Profile': 'stateful', + }), + }), + ); + }); }); describe('deferred-preview flow (office-bucket files)', () => { @@ -1610,6 +1685,22 @@ describe('Code Process', () => { expect(call.data.lang).toBe('bash'); }); + it('routes to the selected profile endpoint and asserts the expected profile', async () => { + mockAxios.mockResolvedValueOnce({ data: { stdout: 'ok', stderr: '' } }); + + await readSandboxFile({ + file_path: '/mnt/data/x.txt', + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + 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.data.runtime_session_hint).toBe('v1:user'); + }); + it('omits session_id and files when not provided', async () => { mockAxios.mockResolvedValueOnce({ data: { stdout: '', stderr: '' } }); @@ -2118,6 +2209,83 @@ describe('Code Process', () => { expect(uploadArgs.version).toBe(4); }); + it('reuploads instead of reusing a ref from the other execution profile', async () => { + const dbFile = { + file_id: 'librechat-file-id', + filename: 'sentinel.txt', + filepath: '/uploads/sentinel.txt', + source: 'local', + context: 'execute_code', + metadata: { + codeEnvRef: { + kind: 'user', + id: 'user-123', + storage_session_id: 'DEFAULT_SESSION', + file_id: 'DEFAULT_ID', + executionProfile: 'default', + }, + }, + }; + getFiles.mockResolvedValue([dbFile]); + const { handleFileUpload } = setupReuploadMocks({ + storage_session_id: 'STATEFUL_SESSION', + file_id: 'STATEFUL_ID', + }); + + await primeFiles({ + req: { user: { id: 'user-123', role: 'USER' } }, + tool_resources: { + execute_code: { file_ids: ['librechat-file-id'], files: [] }, + }, + agentId: 'agent-id', + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }); + + expect(mockAxios).not.toHaveBeenCalled(); + expect(handleFileUpload).toHaveBeenCalledWith( + expect.objectContaining({ + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }), + ); + expect(updateFile).toHaveBeenCalledWith( + expect.objectContaining({ + 'metadata.codeEnvRef': expect.objectContaining({ executionProfile: 'default' }), + 'metadata.codeEnvRefs.stateful': expect.objectContaining({ + executionProfile: 'stateful', + }), + }), + ); + + const persistedMetadata = { + ...dbFile.metadata, + codeEnvRef: updateFile.mock.calls[0][0]['metadata.codeEnvRef'], + codeEnvRefs: { + default: dbFile.metadata.codeEnvRef, + stateful: updateFile.mock.calls[0][0]['metadata.codeEnvRefs.stateful'], + }, + }; + getFiles.mockResolvedValue([{ ...dbFile, metadata: persistedMetadata }]); + mockAxios.mockResolvedValue({ data: { lastModified: new Date().toISOString() } }); + + await primeFiles({ + req: { user: { id: 'user-123', role: 'USER' } }, + tool_resources: { + execute_code: { file_ids: ['librechat-file-id'], files: [] }, + }, + agentId: 'agent-id', + executionProfile: 'default', + }); + + expect(handleFileUpload).toHaveBeenCalledTimes(1); + expect(mockAxios).toHaveBeenCalledWith( + expect.objectContaining({ + url: expect.stringContaining('/sessions/DEFAULT_SESSION/objects/DEFAULT_ID'), + }), + ); + }); + it('persists fresh codeEnvRef (kind/id preserved) on the DB record after reupload', async () => { const dbFile = { file_id: 'librechat-file-id', @@ -2149,14 +2317,20 @@ describe('Code Process', () => { expect(updateFile).toHaveBeenCalledWith( expect.objectContaining({ file_id: 'librechat-file-id', - metadata: expect.objectContaining({ - codeEnvRef: { - kind: 'user', - id: 'user-123', - storage_session_id: 'NEW_SESSION', - file_id: 'NEW_ID', - }, - }), + 'metadata.codeEnvRef': { + kind: 'user', + id: 'user-123', + storage_session_id: 'NEW_SESSION', + file_id: 'NEW_ID', + executionProfile: 'default', + }, + 'metadata.codeEnvRefs.default': { + kind: 'user', + id: 'user-123', + storage_session_id: 'NEW_SESSION', + file_id: 'NEW_ID', + executionProfile: 'default', + }, }), ); }); diff --git a/api/server/services/Files/process.js b/api/server/services/Files/process.js index 527e7b0e31..71777891d7 100644 --- a/api/server/services/Files/process.js +++ b/api/server/services/Files/process.js @@ -10,6 +10,7 @@ const { imageExtRegex, EModelEndpoint, EToolResources, + mergeCodeEnvRef, mergeFileConfig, AgentCapabilities, checkOpenAIStorage, @@ -69,7 +70,8 @@ const createSanitizedUploadWrapper = (uploadFunction) => { }; }; -const hasCodeEnvRef = (file) => file?.metadata?.codeEnvRef != null; +const hasCodeEnvRef = (file) => + file?.metadata?.codeEnvRef != null || file?.metadata?.codeEnvRefs != null; const isMissingStorageError = (err) => { const code = err?.code ?? err?.status ?? err?.statusCode ?? err?.response?.status; @@ -727,14 +729,13 @@ const processAgentFileUpload = async ({ req, res, metadata, sseStream }) => { * `fileIdentifier` key would be silently dropped by mongoose strict * mode and the file would lose its sandbox reference on subsequent * priming turns. */ - fileInfoMetadata = { - codeEnvRef: { - kind: codeKind, - id: codeId, - storage_session_id: uploaded.storage_session_id, - file_id: uploaded.file_id, - }, - }; + fileInfoMetadata = mergeCodeEnvRef(undefined, { + kind: codeKind, + id: codeId, + storage_session_id: uploaded.storage_session_id, + file_id: uploaded.file_id, + executionProfile: 'default', + }); } else if (tool_resource === EToolResources.file_search) { const isFileSearchEnabled = await checkCapability(req, AgentCapabilities.file_search); if (!isFileSearchEnabled) { diff --git a/api/server/services/Files/process.spec.js b/api/server/services/Files/process.spec.js index 48c37166df..9f8ac36b66 100644 --- a/api/server/services/Files/process.spec.js +++ b/api/server/services/Files/process.spec.js @@ -791,6 +791,16 @@ describe('processAgentFileUpload', () => { id: 'user-123', storage_session_id: 'sess-1', file_id: 'fid-1', + executionProfile: 'default', + }, + codeEnvRefs: { + default: { + kind: 'user', + id: 'user-123', + storage_session_id: 'sess-1', + file_id: 'fid-1', + executionProfile: 'default', + }, }, }, }), @@ -819,6 +829,16 @@ describe('processAgentFileUpload', () => { id: 'agent-abc', storage_session_id: 'sess-2', file_id: 'fid-2', + executionProfile: 'default', + }, + codeEnvRefs: { + default: { + kind: 'agent', + id: 'agent-abc', + storage_session_id: 'sess-2', + file_id: 'fid-2', + executionProfile: 'default', + }, }, }, }), @@ -879,6 +899,16 @@ describe('processAgentFileUpload', () => { id: 'agent-abc', storage_session_id: 'sess-5', file_id: 'fid-5', + executionProfile: 'default', + }, + codeEnvRefs: { + default: { + kind: 'agent', + id: 'agent-abc', + storage_session_id: 'sess-5', + file_id: 'fid-5', + executionProfile: 'default', + }, }, }, }), diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index c6654c0daf..f03b6e8933 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -35,6 +35,7 @@ const { isNormalizationSensitiveName, AGENT_EXPECTED_MCP_TOOLS_UNAVAILABLE, isFatalAgentInitializationError, + resolveCodeExecutionContext, } = require('@librechat/api'); const { Time, @@ -583,6 +584,7 @@ async function loadToolDefinitionsWrapper({ streamId = null, jobCreatedAt, tool_resources, + codeExecutionContext, accessibleMcpServerNames, }) { if (!agent.tools || agent.tools.length === 0) { @@ -608,6 +610,18 @@ async function loadToolDefinitionsWrapper({ const codeExecutionEnabled = agent.tools?.includes(Tools.execute_code) === true && enabledCapabilities.has(AgentCapabilities.execute_code); + const resolvedCodeExecutionContext = + codeExecutionContext ?? + resolveCodeExecutionContext({ + statefulSessions: + codeExecutionEnabled && + enabledCapabilities.has(AgentCapabilities.stateful_code_sessions) && + agent.stateful_code_sessions === true, + environment: agent.stateful_code_environment, + userId: req.user.id, + agentId: agent.id, + conversationId: req.body?.conversationId, + }); const hasMCPTools = agent.tools?.some((tool) => tool?.includes(Constants.mcp_delimiter)); const mcpPermissionContext = createMCPPermissionContext(req); const canUseMCP = hasMCPTools ? await mcpPermissionContext.canUseServers(req.user) : true; @@ -1158,6 +1172,8 @@ async function loadToolDefinitionsWrapper({ tool_resources, agentId: agent.id, agentResourceType, + codeApiBaseUrl: resolvedCodeExecutionContext.baseUrl, + executionProfile: resolvedCodeExecutionContext.executionProfile, }); if (toolContext) { dynamicToolContextMap[Tools.execute_code] = toolContext; @@ -1258,6 +1274,7 @@ async function loadAgentTools({ streamId = null, jobCreatedAt, definitionsOnly = true, + codeExecutionContext: providedCodeExecutionContext, accessibleMcpServerNames, }) { if (definitionsOnly) { @@ -1270,6 +1287,7 @@ async function loadAgentTools({ streamId, jobCreatedAt, tool_resources, + codeExecutionContext: providedCodeExecutionContext, accessibleMcpServerNames, }); } catch (error) { @@ -1368,6 +1386,23 @@ async function loadAgentTools({ }); } + const codeExecutionEnabled = + agent.tools?.includes(Tools.execute_code) === true && + enabledCapabilities.has(AgentCapabilities.execute_code); + const statefulCodeSessions = + codeExecutionEnabled && + enabledCapabilities.has(AgentCapabilities.stateful_code_sessions) && + agent.stateful_code_sessions === true; + const codeExecutionContext = + providedCodeExecutionContext ?? + resolveCodeExecutionContext({ + statefulSessions: statefulCodeSessions, + environment: agent.stateful_code_environment, + userId: req.user.id, + agentId: agent.id, + conversationId: req.body?.conversationId, + }); + const { loadedTools, toolContextMap, dynamicToolContextMap, primedCodeFiles } = await loadTools({ agent, signal, @@ -1388,6 +1423,7 @@ async function loadAgentTools({ returnMetadata: true, mcpPermissionContext, requestScopedConnections: getMCPRequestContext(req, res), + codeExecutionContext, [Tools.web_search]: webSearchCallbacks, }, webSearch: appConfig.webSearch, @@ -1398,9 +1434,6 @@ async function loadAgentTools({ /** Build tool registry from MCP tools and create PTC/tool search tools if configured */ const deferredToolsEnabled = checkCapability(AgentCapabilities.deferred_tools); const programmaticToolsEnabled = enabledCapabilities.has(AgentCapabilities.programmatic_tools); - const codeExecutionEnabled = - agent.tools?.includes(Tools.execute_code) === true && - enabledCapabilities.has(AgentCapabilities.execute_code); const { toolRegistry, toolDefinitions, additionalTools, hasDeferredTools } = await buildToolClassification({ loadedTools, @@ -1412,6 +1445,7 @@ async function loadAgentTools({ programmaticToolsEnabled, codeExecutionEnabled, authHeaders: () => getCodeApiAuthHeaders(req), + codeExecutionContext, }); const agentTools = []; @@ -1646,6 +1680,7 @@ async function loadAgentTools({ * @param {Object} [params.tool_resources] - Tool resources * @param {string|null} [params.streamId] - Stream ID for web search callbacks * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events + * @param {string} [params.conversationId] - Resolved conversation identity for this request * @param {boolean} [params.actionsEnabled] - Whether the actions capability is enabled * @param {readonly string[]} [params.accessibleMcpServerNames] - COMPLETE accessible-server audit resolved at initialization * @returns {Promise<{ loadedTools: Array, configurable: Object }>} @@ -1666,6 +1701,7 @@ async function loadToolsForExecution({ tool_resources, streamId = null, jobCreatedAt, + conversationId, actionsEnabled, accessibleMcpServerNames, }) { @@ -1695,9 +1731,19 @@ async function loadToolsForExecution({ const isBashToolRequested = toolNames.includes(AgentConstants.BASH_TOOL); const isLegacyExecuteCodeRequested = toolNames.includes(Tools.execute_code); const isCodeExecutionToolRequested = isBashToolRequested || isLegacyExecuteCodeRequested; + const isSkillToolRequested = toolNames.includes(AgentConstants.SKILL_TOOL); + const isSandboxFileToolRequested = toolNames.some((name) => + [AgentConstants.READ_FILE, AgentConstants.CREATE_FILE, AgentConstants.EDIT_FILE].includes(name), + ); let enabledCapabilities; - if (actionsEnabled === undefined || isPTCRequested || isCodeExecutionToolRequested) { + if ( + actionsEnabled === undefined || + isPTCRequested || + isCodeExecutionToolRequested || + isSkillToolRequested || + isSandboxFileToolRequested + ) { enabledCapabilities = await resolveAgentCapabilities(req, appConfig, agent?.id); } if (actionsEnabled === undefined) { @@ -1707,17 +1753,21 @@ async function loadToolsForExecution({ enabledCapabilities?.has(AgentCapabilities.execute_code) === true && agent?.tools?.includes(Tools.execute_code) === true; - /** - * Opt bash_tool into the hedged stateful-session description. Gated on code - * execution being enabled AND the admin `stateful_code_sessions` capability - * AND the agent's own builder opt-in; off by default. Sets prompt text only - * (the wire hint is set at run config). PTC keeps its stateless prompt in - * v1. Older @librechat/agents ignore the param. - */ + /** Resolve the trusted endpoint/profile from the actually executing agent. + * This stays per-agent across handoffs and subagents; no graph-global stateful + * flag or model-supplied value is consulted. */ const statefulCodeSessions = codeExecutionEnabled && enabledCapabilities?.has(AgentCapabilities.stateful_code_sessions) === true && agent?.stateful_code_sessions === true; + const codeExecutionContext = resolveCodeExecutionContext({ + statefulSessions: statefulCodeSessions, + environment: agent?.stateful_code_environment, + userId: req.user.id, + agentId: agent?.id, + conversationId: conversationId ?? req.body?.conversationId, + }); + configurable.codeExecutionContext = codeExecutionContext; const isPTC = isPTCRequested && @@ -1747,6 +1797,9 @@ async function loadToolsForExecution({ for (const name of ptcToolNames) { const ptcTool = createBashProgrammaticToolCallingTool({ authHeaders: () => getCodeApiAuthHeaders(req), + baseUrl: codeExecutionContext.baseUrl, + executionProfile: codeExecutionContext.executionProfile, + runtimeSessionHint: codeExecutionContext.runtimeSessionHint, }); ptcTool.name = name; allLoadedTools.push(ptcTool); @@ -1770,7 +1823,7 @@ async function loadToolsForExecution({ try { const bashTool = createBashExecutionTool({ authHeaders: () => getCodeApiAuthHeaders(req), - statefulSessions: statefulCodeSessions, + ...codeExecutionContext, }); allLoadedTools.push(bashTool); } catch (error) { diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index 67a25dc7bf..7b2cb74542 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -1,3 +1,4 @@ +const { createHash } = require('node:crypto'); const { Constants: AgentConstants } = require('@librechat/agents'); const { Tools, @@ -16,6 +17,40 @@ const mockGetMCPServerTools = jest.fn(); const mockGetCachedTools = jest.fn(); const mockSendEvent = jest.fn(); const mockEmitChunk = jest.fn(); +const mockResolveCodeExecutionContext = jest.fn( + ({ statefulSessions, environment, userId, agentId, conversationId }) => { + if (!statefulSessions) { + return { + baseUrl: (process.env.LIBRECHAT_CODE_BASEURL ?? 'https://api.librechat.ai').replace( + /\/$/, + '', + ), + codeSessionKey: 'execute_code', + executionProfile: 'default', + statefulSessions: false, + }; + } + const baseUrl = process.env.LIBRECHAT_CODE_BASEURL_STATEFUL?.replace(/\/$/, ''); + if (!baseUrl) { + throw new Error('LIBRECHAT_CODE_BASEURL_STATEFUL is not configured'); + } + const fingerprint = (...parts) => + createHash('sha256').update(JSON.stringify(parts)).digest('hex').slice(0, 32); + let runtimeSessionHint = `v2:user:${fingerprint(userId)}`; + if (environment === 'agent-user') { + runtimeSessionHint = `v2:agent-user:${fingerprint(userId, agentId)}`; + } else if (environment === 'conversation') { + runtimeSessionHint = `v2:conversation:${fingerprint(userId, conversationId)}`; + } + return { + baseUrl, + codeSessionKey: `execute_code:stateful:${runtimeSessionHint}`, + executionProfile: 'stateful', + runtimeSessionHint, + statefulSessions: true, + }; + }, +); jest.mock('~/server/services/Config', () => ({ getEndpointsConfig: (...args) => mockGetEndpointsConfig(...args), getMCPServerTools: (...args) => mockGetMCPServerTools(...args), @@ -35,6 +70,7 @@ jest.mock('@librechat/api', () => ({ GenerationJobManager: { emitChunk: (...args) => mockEmitChunk(...args), }, + resolveCodeExecutionContext: (...args) => mockResolveCodeExecutionContext(...args), })); const mockLoadToolsUtil = jest.fn(); @@ -267,7 +303,43 @@ describe('ToolService - Action Capability Gating', () => { agentResourceType: ResourceType.REMOTE_AGENT, }; expect(primeSearchFiles).toHaveBeenCalledWith(expectedParams); - expect(primeCodeFiles).toHaveBeenCalledWith(expectedParams); + expect(primeCodeFiles).toHaveBeenCalledWith({ + ...expectedParams, + codeApiBaseUrl: 'https://api.librechat.ai', + executionProfile: 'default', + }); + }); + + it('primes code files through the initializer-selected stateful route', async () => { + const capabilities = [AgentCapabilities.tools, AgentCapabilities.execute_code]; + const req = createMockReq(capabilities); + const tool_resources = { execute_code: { file_ids: ['stateful-file'] } }; + const { primeFiles: primeCodeFiles } = require('~/server/services/Files/Code/process'); + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + + await loadAgentTools({ + req, + res: {}, + agent: { id: 'stateful-agent', tools: [Tools.execute_code] }, + tool_resources, + definitionsOnly: true, + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + codeSessionKey: 'execute_code:stateful:v2:user:abc', + executionProfile: 'stateful', + runtimeSessionHint: 'v2:user:abc', + statefulSessions: true, + }, + }); + + expect(primeCodeFiles).toHaveBeenCalledWith({ + req, + tool_resources, + agentId: 'stateful-agent', + agentResourceType: undefined, + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }); }); it('propagates a typed CodeAPI resource recovery failure before model invocation', async () => { @@ -1402,6 +1474,127 @@ describe('ToolService - Action Capability Gating', () => { expect(mockLoadToolsUtil).not.toHaveBeenCalled(); }); + it('keeps stateless and stateful agents on isolated execution profiles in one run', async () => { + const capabilities = [ + AgentCapabilities.tools, + AgentCapabilities.execute_code, + AgentCapabilities.stateful_code_sessions, + ]; + const req = createMockReq(capabilities); + req.body = { conversationId: 'conversation-1' }; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + process.env.LIBRECHAT_CODE_BASEURL = 'http://code-default.test/v1'; + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + + try { + const stateless = await loadToolsForExecution({ + req, + res: {}, + agent: { id: 'stateless-agent', tools: [Tools.execute_code] }, + toolNames: [], + }); + const stateful = await loadToolsForExecution({ + req, + res: {}, + agent: { + id: 'stateful-agent', + tools: [Tools.execute_code], + stateful_code_sessions: true, + stateful_code_environment: 'agent-user', + }, + toolNames: [], + }); + + expect(stateless.configurable.codeExecutionContext).toEqual({ + baseUrl: 'http://code-default.test/v1', + codeSessionKey: 'execute_code', + executionProfile: 'default', + statefulSessions: false, + }); + expect(stateful.configurable.codeExecutionContext).toEqual({ + baseUrl: 'http://code-stateful.test/v1', + codeSessionKey: 'execute_code:stateful:v2:agent-user:7c684f0773d9642c122f67aa30e9e0f4', + executionProfile: 'stateful', + runtimeSessionHint: 'v2:agent-user:7c684f0773d9642c122f67aa30e9e0f4', + statefulSessions: true, + }); + } finally { + delete process.env.LIBRECHAT_CODE_BASEURL; + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + } + }); + + it('resolves stateful routing for host file tools with the controller conversation ID', async () => { + const capabilities = [ + AgentCapabilities.tools, + AgentCapabilities.execute_code, + AgentCapabilities.stateful_code_sessions, + ]; + const req = createMockReq(capabilities); + req.body = {}; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + + try { + const result = await loadToolsForExecution({ + req, + res: {}, + conversationId: 'resolved-api-conversation', + agent: { + id: 'stateful-agent', + tools: [Tools.execute_code], + stateful_code_sessions: true, + stateful_code_environment: 'conversation', + }, + toolNames: [AgentConstants.READ_FILE], + actionsEnabled: false, + }); + + expect(result.configurable.codeExecutionContext.executionProfile).toBe('stateful'); + expect(mockResolveCodeExecutionContext).toHaveBeenLastCalledWith( + expect.objectContaining({ + statefulSessions: true, + conversationId: 'resolved-api-conversation', + }), + ); + } finally { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + } + }); + + it('resolves stateful routing when handle_skill is the only requested tool', async () => { + const capabilities = [ + AgentCapabilities.tools, + AgentCapabilities.execute_code, + AgentCapabilities.stateful_code_sessions, + ]; + const req = createMockReq(capabilities); + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + + try { + const result = await loadToolsForExecution({ + req, + res: {}, + agent: { + id: 'stateful-agent', + tools: [Tools.execute_code], + stateful_code_sessions: true, + stateful_code_environment: 'agent-user', + }, + toolNames: [AgentConstants.SKILL_TOOL], + actionsEnabled: false, + }); + + expect(result.configurable.codeExecutionContext.executionProfile).toBe('stateful'); + expect(mockResolveCodeExecutionContext).toHaveBeenLastCalledWith( + expect.objectContaining({ statefulSessions: true, environment: 'agent-user' }), + ); + } finally { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + } + }); + it('loads bash PTC under the legacy programmatic tool name when code capabilities are enabled', async () => { const capabilities = [ AgentCapabilities.tools, diff --git a/client/src/common/agents-types.ts b/client/src/common/agents-types.ts index b519de7eef..c5ac569ac2 100644 --- a/client/src/common/agents-types.ts +++ b/client/src/common/agents-types.ts @@ -6,6 +6,7 @@ import type { SupportContact, AgentProvider, MemoryScope, + StatefulCodeEnvironment, GraphEdge, Agent, } from 'librechat-data-provider'; @@ -44,6 +45,8 @@ export type AgentForm = { skills_enabled?: boolean; /** Memory partition: 'agent' isolates memories per (user, agent); default shared pool */ memory_scope?: MemoryScope; + /** Sharing scope for stateful Code API workspaces. */ + stateful_code_environment?: StatefulCodeEnvironment; provider?: AgentProvider | OptionWithIcon; /** @deprecated Use edges instead */ agent_ids?: string[]; diff --git a/client/src/components/SidePanel/Agents/Advanced/StatefulSessions.tsx b/client/src/components/SidePanel/Agents/Advanced/StatefulSessions.tsx index d4dee3d9fe..92b6ea7a29 100644 --- a/client/src/components/SidePanel/Agents/Advanced/StatefulSessions.tsx +++ b/client/src/components/SidePanel/Agents/Advanced/StatefulSessions.tsx @@ -2,12 +2,18 @@ import { useFormContext } from 'react-hook-form'; import { AgentCapabilities } from 'librechat-data-provider'; import { Switch, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, HoverCard, HoverCardPortal, HoverCardContent, HoverCardTrigger, CircleHelpIcon, } from '@librechat/client'; +import type { StatefulCodeEnvironment } from 'librechat-data-provider'; import type { AgentForm } from '~/common'; import { useLocalize } from '~/hooks'; import { ESide } from '~/common'; @@ -19,41 +25,83 @@ export default function StatefulSessions() { const enabled = watch(AgentCapabilities.stateful_code_sessions) ?? false; const codeEnabled = watch(AgentCapabilities.execute_code); + const environment = watch('stateful_code_environment') ?? 'user'; const handleChange = (value: boolean) => { setValue(AgentCapabilities.stateful_code_sessions, value, { shouldDirty: true }); + if (value && !watch('stateful_code_environment')) { + setValue('stateful_code_environment', 'user', { shouldDirty: true }); + } }; return ( - -
-
-
- {localize('com_ui_stateful_sessions')} -
- - - -
- - -
-

- {localize('com_nav_info_stateful_sessions')} -

+
+ +
+
+
+ {localize('com_ui_stateful_sessions')}
- - - -
- + + + +
+ + +
+

+ {localize('com_nav_info_stateful_sessions')} +

+
+
+
+ +
+ + {enabled && codeEnabled === true && ( +
+ + +

+ {localize('com_nav_info_stateful_code_environment')} +

+
+ )} +
); } diff --git a/client/src/components/SidePanel/Agents/AgentPanel.tsx b/client/src/components/SidePanel/Agents/AgentPanel.tsx index f36aa8ddad..1763eb99f7 100644 --- a/client/src/components/SidePanel/Agents/AgentPanel.tsx +++ b/client/src/components/SidePanel/Agents/AgentPanel.tsx @@ -75,6 +75,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n end_after_tools, hide_sequential_outputs, stateful_code_sessions, + stateful_code_environment, recursion_limit, category, support_contact, @@ -89,6 +90,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n * execute_code is disabled so a stale opt-in can't silently reactivate later. */ const normalizedStatefulCodeSessions = data.execute_code === true ? stateful_code_sessions : false; + const normalizedStatefulCodeEnvironment = stateful_code_environment ?? 'user'; const shouldResetAvatar = avatarActionState === 'reset' && Boolean(agent_id) && !isEphemeralAgent(agent_id); @@ -111,6 +113,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n end_after_tools, hide_sequential_outputs, stateful_code_sessions: normalizedStatefulCodeSessions, + stateful_code_environment: normalizedStatefulCodeEnvironment, recursion_limit, category, support_contact, diff --git a/client/src/components/SidePanel/Agents/AgentSelect.tsx b/client/src/components/SidePanel/Agents/AgentSelect.tsx index d9a3775b50..50bd290f62 100644 --- a/client/src/components/SidePanel/Agents/AgentSelect.tsx +++ b/client/src/components/SidePanel/Agents/AgentSelect.tsx @@ -86,6 +86,7 @@ function AgentSelect({ avatar_file: null, avatar_preview: fullAgent.avatar?.filepath ?? '', avatar_action: null, + stateful_code_environment: fullAgent.stateful_code_environment ?? 'user', }; Object.entries(fullAgent).forEach(([name, value]) => { diff --git a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts index 599e2e216e..da3471824e 100644 --- a/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts +++ b/client/src/components/SidePanel/Agents/__tests__/AgentPanel.helpers.spec.ts @@ -94,6 +94,27 @@ describe('composeAgentUpdatePayload', () => { expect(payload.stateful_code_sessions).toBe(true); }); + + it('defaults stateful environments to the scalable user scope', () => { + const form = createForm(); + form.execute_code = true; + form.stateful_code_sessions = true; + + const { payload } = composeAgentUpdatePayload(form, 'agent_123'); + + expect(payload.stateful_code_environment).toBe('user'); + }); + + it('preserves an explicit stateful environment scope', () => { + const form = createForm(); + form.execute_code = true; + form.stateful_code_sessions = true; + form.stateful_code_environment = 'agent-user'; + + const { payload } = composeAgentUpdatePayload(form, 'agent_123'); + + expect(payload.stateful_code_environment).toBe('agent-user'); + }); }); describe('persistAvatarChanges', () => { diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 83d705302f..788c84ad2c 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -535,7 +535,8 @@ "com_nav_info_save_draft": "When enabled, the text and attachments you enter in the chat form will be automatically saved locally as drafts. These drafts will be available even if you reload the page or switch to a different conversation. Drafts are stored locally on your device and are deleted once the message is sent.", "com_nav_info_show_thinking": "When enabled, the chat will display the thinking dropdowns open by default, allowing you to view the AI's reasoning in real-time. When disabled, the thinking dropdowns will remain closed by default for a cleaner and more streamlined interface", "com_nav_info_smooth_streaming": "When enabled, newly streamed words fade in smoothly for the latest response. This is purely visual — it does not delay token delivery — and is disabled automatically when your device prefers reduced motion.", - "com_nav_info_stateful_sessions": "When enabled, this agent's code executions reuse one persistent sandbox workspace per conversation: files, installed packages, and working state usually carry over between runs. The workspace may occasionally reset, so anything important should be saved under /mnt/data. Requires Code Interpreter and the app-level stateful sessions capability.", + "com_nav_info_stateful_sessions": "When enabled, this agent uses the dedicated stateful Code API instead of the default stateless service. Files, installed packages, and working state usually carry over between runs. The workspace may occasionally reset, so save anything important under /mnt/data. Requires Code Interpreter and the app-level stateful sessions capability.", + "com_nav_info_stateful_code_environment": "Choose who shares this agent's stateful workspace. This does not share live files with stateless code sessions.", "com_nav_info_user_name_display": "When enabled, the username of the sender will be shown above each message you send. When disabled, you will only see \"You\" above your messages.", "com_nav_keep_screen_awake": "Keep screen awake during response generation", "com_nav_lang_arabic": "العربية", @@ -2031,6 +2032,10 @@ "com_ui_stack_trace": "Stack Trace", "com_ui_standard": "Standard", "com_ui_stateful_sessions": "Stateful code sessions", + "com_ui_stateful_code_environment": "Stateful environment", + "com_ui_stateful_code_environment_user": "User workspace (recommended)", + "com_ui_stateful_code_environment_agent_user": "Agent + user workspace", + "com_ui_stateful_code_environment_conversation": "Conversation workspace", "com_ui_status_prefix": "Status:", "com_ui_steer": "Steer", "com_ui_steer_already_applied": "That steering message already reached the agent, so it was left in the response", diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index dfbed9aacd..fb0ff6de28 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -1739,6 +1739,49 @@ describe('initializeAgent — execute_code capability expansion', () => { expect(result.fileAuthoringToolNames).toEqual(new Set(['create_file', 'edit_file'])); }); + it('routes code-file priming through the stateful profile before tools load', async () => { + const { agent, req, res, loadTools, db } = createMocks(); + agent.tools = ['execute_code']; + agent.stateful_code_sessions = true; + agent.stateful_code_environment = 'agent-user'; + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'https://stateful-code.example.com/v1/'; + + try { + const result = await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + codeEnvAvailable: true, + statefulSessionsAvailable: true, + }, + db, + ); + + expect(loadTools).toHaveBeenCalledWith( + expect.objectContaining({ + codeExecutionContext: expect.objectContaining({ + baseUrl: 'https://stateful-code.example.com/v1', + executionProfile: 'stateful', + statefulSessions: true, + }), + }), + ); + expect(result.codeExecutionContext).toEqual( + expect.objectContaining({ + baseUrl: 'https://stateful-code.example.com/v1', + executionProfile: 'stateful', + }), + ); + } finally { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + } + }); + it('upgrades read_file to the skill-aware description when active skills are in scope', async () => { const { agent, req, res, loadTools, db } = createMocks(); agent.tools = ['execute_code']; diff --git a/packages/api/src/agents/__tests__/run-codeTools.test.ts b/packages/api/src/agents/__tests__/run-codeTools.test.ts index 5a027bf828..f977e465dd 100644 --- a/packages/api/src/agents/__tests__/run-codeTools.test.ts +++ b/packages/api/src/agents/__tests__/run-codeTools.test.ts @@ -72,9 +72,9 @@ function makeAgent(overrides?: Record) { }; } -async function captureRunConfig(): Promise> { +async function captureRunConfig(agent = makeAgent()): Promise> { await createRun({ - agents: [makeAgent()] as never, + agents: [agent] as never, signal: new AbortController().signal, streaming: true, streamUsage: true, @@ -105,4 +105,12 @@ describe('createRun code-tool eager/session wiring', () => { expect.arrayContaining(['create_file', 'edit_file', 'read_file']), ); }); + + it('passes the trusted per-agent code-session partition to the SDK', async () => { + const codeSessionKey = 'execute_code:stateful:v1:user'; + const runConfig = await captureRunConfig(makeAgent({ codeSessionKey })); + const [agentInput] = (runConfig.graphConfig as { agents: Array> }) + .agents; + expect(agentInput.codeSessionKey).toBe(codeSessionKey); + }); }); diff --git a/packages/api/src/agents/codeFilesSession.spec.ts b/packages/api/src/agents/codeFilesSession.spec.ts index a5ac577703..d1c63981ae 100644 --- a/packages/api/src/agents/codeFilesSession.spec.ts +++ b/packages/api/src/agents/codeFilesSession.spec.ts @@ -1,7 +1,9 @@ import { Constants } from '@librechat/agents'; import type { CodeEnvFile, CodeSessionContext, ToolSessionMap } from '@librechat/agents'; import { + buildAgentInitialToolSessions, buildInitialToolSessions, + collectCodeExecutionProfileRoutes, seedCodeFilesIntoSessions, type CodeFilesAgent, } from './codeFilesSession'; @@ -144,6 +146,65 @@ describe('seedCodeFilesIntoSessions', () => { expect(entry.files).toHaveLength(2); expect(entry.files!.map((f) => f.storage_session_id).sort()).toEqual(['sess-A', 'sess-B']); }); + + it('seeds only the requested code-session partition', () => { + const statefulKey = 'execute_code:stateful:v1:user'; + const existing: ToolSessionMap = new Map(); + existing.set(Constants.EXECUTE_CODE, { + session_id: 'stateless-session', + files: [file('s1', 'stateless-session', 'stateless.txt')], + lastUpdated: 1, + } satisfies CodeSessionContext); + + const result = seedCodeFilesIntoSessions( + [file('w1', 'stateful-session', 'stateful.txt')], + existing, + statefulKey, + ); + + expect(result!.get(Constants.EXECUTE_CODE)?.files?.map((f) => f.id)).toEqual(['s1']); + expect(result!.get(statefulKey)?.files?.map((f) => f.id)).toEqual(['w1']); + }); +}); + +describe('buildAgentInitialToolSessions', () => { + it('clones only the agent partition and merges files resolved after the run seed', () => { + const statefulKey = 'execute_code:stateful:v2:user:user-1'; + const statelessFile = file('stateless', 'stateless-session', 'stateless.txt'); + const skillFile = file('skill', 'stateful-skill-session', 'skills/tool.py'); + const lazyAttachment = file('attachment', 'stateful-user-session', 'input.csv'); + const runSessions: ToolSessionMap = new Map([ + [ + Constants.EXECUTE_CODE, + { + session_id: statelessFile.storage_session_id, + files: [statelessFile], + lastUpdated: 1, + } satisfies CodeSessionContext, + ], + [ + statefulKey, + { + session_id: skillFile.storage_session_id, + files: [skillFile], + lastUpdated: 2, + } satisfies CodeSessionContext, + ], + ]); + + const result = buildAgentInitialToolSessions( + { codeSessionKey: statefulKey, primedCodeFiles: [lazyAttachment] }, + runSessions, + ); + + expect(result?.has(Constants.EXECUTE_CODE)).toBe(false); + expect(result?.get(statefulKey)?.files?.map((entry) => entry.id)).toEqual([ + 'skill', + 'attachment', + ]); + expect(result).not.toBe(runSessions); + expect(runSessions.get(statefulKey)?.files?.map((entry) => entry.id)).toEqual(['skill']); + }); }); describe('buildInitialToolSessions', () => { @@ -151,8 +212,10 @@ describe('buildInitialToolSessions', () => { name: string, primedCodeFiles?: CodeEnvFile[], subagents?: CodeFilesAgent[], + codeSessionKey?: string, ): CodeFilesAgent & { __label: string } => ({ __label: name, + codeSessionKey, primedCodeFiles, subagentAgentConfigs: subagents, }); @@ -327,4 +390,168 @@ describe('buildInitialToolSessions', () => { expect(entry.files).toHaveLength(2); expect(entry.files!.map((f) => f.name).sort()).toEqual(['shared.csv', 'top.csv']); }); + + it('keeps stateless and stateful agent files in separate partitions', () => { + const statefulKey = 'execute_code:stateful:v1:user'; + const skillSessions: ToolSessionMap = new Map(); + skillSessions.set(Constants.EXECUTE_CODE, { + session_id: 'skill-sess', + files: [file('skill-1', 'skill-sess', 'skill.py')], + lastUpdated: 1, + } satisfies CodeSessionContext); + skillSessions.set(statefulKey, { + session_id: 'stateful-skill-sess', + files: [file('stateful-skill-1', 'stateful-skill-sess', 'skill.py')], + lastUpdated: 1, + } satisfies CodeSessionContext); + + const result = buildInitialToolSessions({ + skillSessions, + agents: [ + agent('stateless', [file('s1', 'stateless-sess', 'stateless.txt')]), + agent('stateful', [file('w1', 'stateful-sess', 'stateful.txt')], undefined, statefulKey), + ], + }); + + expect(result!.get(Constants.EXECUTE_CODE)?.files?.map((f) => f.id)).toEqual(['skill-1', 's1']); + expect(result!.get(statefulKey)?.files?.map((f) => f.id)).toEqual(['stateful-skill-1', 'w1']); + }); + + it('shares user-scoped stateful files but isolates agent-user scopes', () => { + const userKey = 'execute_code:stateful:v1:user'; + const firstAgentKey = 'execute_code:stateful:v1:agent-user:agent-a'; + const secondAgentKey = 'execute_code:stateful:v1:agent-user:agent-b'; + + const result = buildInitialToolSessions({ + agents: [ + agent('user-a', [file('u1', 'user-a-sess', 'a.txt')], undefined, userKey), + agent('user-b', [file('u2', 'user-b-sess', 'b.txt')], undefined, userKey), + agent('agent-a', [file('a1', 'agent-a-sess', 'private-a.txt')], undefined, firstAgentKey), + agent('agent-b', [file('b1', 'agent-b-sess', 'private-b.txt')], undefined, secondAgentKey), + ], + }); + + expect(result!.get(userKey)?.files?.map((f) => f.id)).toEqual(['u1', 'u2']); + expect(result!.get(firstAgentKey)?.files?.map((f) => f.id)).toEqual(['a1']); + expect(result!.get(secondAgentKey)?.files?.map((f) => f.id)).toEqual(['b1']); + expect(result!.has(Constants.EXECUTE_CODE)).toBe(false); + }); + + it('preserves a profile-local stateful skill seed without agent files', () => { + const statefulKey = 'execute_code:stateful:v1:user'; + const skillSessions: ToolSessionMap = new Map(); + skillSessions.set(statefulKey, { + session_id: 'skill-sess', + files: [file('skill-1', 'skill-sess', 'skill.py')], + lastUpdated: 1, + } satisfies CodeSessionContext); + + const result = buildInitialToolSessions({ + skillSessions, + agents: [agent('stateful', undefined, undefined, statefulKey)], + }); + + expect(result!.get(statefulKey)?.files?.map((f) => f.id)).toEqual(['skill-1']); + }); + + it('never copies a default-profile skill pointer into a stateful partition', () => { + const statefulKey = 'execute_code:stateful:v1:user'; + const skillSessions: ToolSessionMap = new Map(); + skillSessions.set(Constants.EXECUTE_CODE, { + session_id: 'default-skill-sess', + files: [file('skill-1', 'default-skill-sess', 'skill.py')], + lastUpdated: 1, + } satisfies CodeSessionContext); + + const result = buildInitialToolSessions({ + skillSessions, + agents: [agent('stateful', undefined, undefined, statefulKey)], + }); + + expect(result!.has(statefulKey)).toBe(false); + }); +}); + +describe('collectCodeExecutionProfileRoutes', () => { + it('groups reachable code agents by deployment and retains each trusted partition', () => { + const statelessContext = { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: Constants.EXECUTE_CODE, + executionProfile: 'default' as const, + statefulSessions: false, + }; + const statefulContext = (key: string) => ({ + baseUrl: 'https://stateful.example.com/v1', + codeSessionKey: key, + executionProfile: 'stateful' as const, + runtimeSessionHint: key.slice('execute_code:stateful:'.length), + statefulSessions: true, + }); + const childKey = 'execute_code:stateful:v2:agent-user:child'; + const parentKey = 'execute_code:stateful:v2:user:shared'; + const child: CodeFilesAgent = { + codeEnvAvailable: true, + codeExecutionContext: statefulContext(childKey), + codeSessionKey: childKey, + }; + + const routes = collectCodeExecutionProfileRoutes([ + { + codeEnvAvailable: true, + codeExecutionContext: statelessContext, + codeSessionKey: Constants.EXECUTE_CODE, + }, + { + codeEnvAvailable: true, + codeExecutionContext: statefulContext(parentKey), + codeSessionKey: parentKey, + subagentAgentConfigs: [child], + }, + { codeEnvAvailable: false, codeExecutionContext: statefulContext('ignored') }, + ]); + + expect(routes).toEqual([ + { + codeExecutionContext: statelessContext, + codeSessionKeys: [Constants.EXECUTE_CODE], + }, + { + codeExecutionContext: statefulContext(parentKey), + codeSessionKeys: [parentKey, childKey], + }, + ]); + }); + + it('derives and includes the trusted profile for a lazy subagent descriptor', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'https://stateful.example.com/v1'; + const routes = collectCodeExecutionProfileRoutes( + [ + { + id: 'parent', + codeEnvAvailable: false, + lazySubagentConfigs: [ + { + id: 'lazy-child', + codeEnvAvailable: true, + statefulCodeSessions: true, + statefulCodeEnvironment: 'agent-user', + }, + ], + }, + ], + { userId: 'user-1', conversationId: 'conversation-1' }, + ); + + expect(routes).toHaveLength(1); + expect(routes[0].codeExecutionContext).toEqual( + expect.objectContaining({ + executionProfile: 'stateful', + statefulSessions: true, + }), + ); + expect(routes[0].codeSessionKeys).toEqual([ + expect.stringMatching(/^execute_code:stateful:v2:agent-user:/), + ]); + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + }); }); diff --git a/packages/api/src/agents/codeFilesSession.ts b/packages/api/src/agents/codeFilesSession.ts index 2b37b27f89..8ffcaeaa13 100644 --- a/packages/api/src/agents/codeFilesSession.ts +++ b/packages/api/src/agents/codeFilesSession.ts @@ -1,16 +1,82 @@ import { Constants } from '@librechat/agents'; import type { FileRefs, CodeEnvFile, ToolSessionMap, CodeSessionContext } from '@librechat/agents'; +import type { StatefulCodeEnvironment } from 'librechat-data-provider'; +import { resolveCodeExecutionContext, type CodeExecutionContext } from './execution'; /** - * Minimal shape for an agent that may contribute primed code files to the - * run-wide sandbox seed. Both `InitializedAgent` and `RunAgent` satisfy it, + * Minimal shape for an agent that may contribute primed code files to its + * execution-profile partition. Both `InitializedAgent` and `RunAgent` satisfy it, * and the recursive walk in {@link buildInitialToolSessions} traverses * `subagentAgentConfigs` so nested subagents (which aren't in the top-level * `agentConfigs` map after pure-subagent pruning) still contribute. */ export interface CodeFilesAgent { + id?: string; + codeEnvAvailable?: boolean; + codeExecutionContext?: CodeExecutionContext; + codeSessionKey?: string; primedCodeFiles?: CodeEnvFile[]; + statefulCodeSessions?: boolean; + statefulCodeEnvironment?: StatefulCodeEnvironment; subagentAgentConfigs?: CodeFilesAgent[]; + lazySubagentConfigs?: CodeFilesAgent[]; +} + +export interface CodeExecutionProfileRoute { + codeExecutionContext: CodeExecutionContext; + codeSessionKeys: string[]; +} + +/** Collects the distinct Code API deployments used by a run and every + * trusted session partition that must receive that deployment's immutable + * skill-file seed. */ +export function collectCodeExecutionProfileRoutes( + agents: Iterable, + scope?: { userId: string; conversationId?: string | null }, +): CodeExecutionProfileRoute[] { + const routes = new Map< + CodeExecutionContext['executionProfile'], + { codeExecutionContext: CodeExecutionContext; codeSessionKeys: Set } + >(); + const visited = new Set(); + const queue: CodeFilesAgent[] = []; + for (const agent of agents) { + if (agent) queue.push(agent); + } + while (queue.length > 0) { + const agent = queue.shift()!; + if (visited.has(agent)) continue; + visited.add(agent); + const context = + agent.codeExecutionContext ?? + (agent.codeEnvAvailable === true && scope + ? resolveCodeExecutionContext({ + statefulSessions: agent.statefulCodeSessions === true, + environment: agent.statefulCodeEnvironment, + userId: scope.userId, + agentId: agent.id, + conversationId: scope.conversationId, + }) + : undefined); + if (agent.codeEnvAvailable === true && context) { + const route = routes.get(context.executionProfile) ?? { + codeExecutionContext: context, + codeSessionKeys: new Set(), + }; + route.codeSessionKeys.add(agent.codeSessionKey ?? context.codeSessionKey); + routes.set(context.executionProfile, route); + } + for (const child of [ + ...(agent.subagentAgentConfigs ?? []), + ...(agent.lazySubagentConfigs ?? []), + ]) { + if (child && !visited.has(child)) queue.push(child); + } + } + return Array.from(routes.values(), (route) => ({ + codeExecutionContext: route.codeExecutionContext, + codeSessionKeys: Array.from(route.codeSessionKeys), + })); } /** @@ -44,13 +110,14 @@ export interface CodeFilesAgent { export function seedCodeFilesIntoSessions( files: CodeEnvFile[] | undefined, existing: ToolSessionMap | undefined, + sessionKey: string = Constants.EXECUTE_CODE, ): ToolSessionMap | undefined { if (!files || files.length === 0) { return existing; } const sessions: ToolSessionMap = existing ?? new Map(); - const prior = sessions.get(Constants.EXECUTE_CODE) as CodeSessionContext | undefined; + const prior = sessions.get(sessionKey) as CodeSessionContext | undefined; /** * Compose `(storage_session_id, id)` as a stable identity. `name` alone @@ -83,7 +150,7 @@ export function seedCodeFilesIntoSessions( return existing; } - sessions.set(Constants.EXECUTE_CODE, { + sessions.set(sessionKey, { session_id: representativeSessionId, files: mergedFiles, lastUpdated: Date.now(), @@ -92,23 +159,40 @@ export function seedCodeFilesIntoSessions( return sessions; } +/** Builds an isolated child-graph seed from the run's exact trusted partition + * plus files resolved specifically for that agent. Lazy subagents are resolved + * after the run-wide seed is built, so their attachments must be copied here + * when `AgentInputs` is created. */ +export function buildAgentInitialToolSessions( + agent: CodeFilesAgent, + runSessions: ToolSessionMap | undefined, +): ToolSessionMap | undefined { + const sessionKey = agent.codeSessionKey ?? Constants.EXECUTE_CODE; + const runContext = runSessions?.get(sessionKey) as CodeSessionContext | undefined; + let sessions: ToolSessionMap | undefined; + if (runContext) { + sessions = new Map([ + [ + sessionKey, + { + ...runContext, + files: runContext.files ? [...runContext.files] : undefined, + }, + ], + ]); + } + return seedCodeFilesIntoSessions(agent.primedCodeFiles, sessions, sessionKey); +} + /** - * Builds the run-wide initial `ToolSessionMap` for `Graph.sessions`, - * combining skill-priming output with code-resource files primed across - * every agent that may execute code in this run. + * Builds the run-wide `ToolSessionMap` for `Graph.sessions`, partitioned by + * each agent's trusted `codeSessionKey`. The legacy `execute_code` partition + * remains shared by stateless agents. Stateful agents share only when their + * configured environment resolves to the same key. * - * **Why "run-wide" (not per-agent):** `Graph.sessions` is a single map - * shared by every `ToolNode` instance in the run by design — the - * agents-library treats the code-execution sandbox as a conversation- - * scoped workspace, not an agent-scoped one. Two agents that both have - * code-execution enabled (a primary + a handoff target, or a parent + - * a subagent) implicitly share session_id and file refs through this - * map. This helper makes that explicit at the seeding boundary: every - * reachable agent's `primedCodeFiles` flows into the same - * `EXECUTE_CODE` entry. If per-agent isolation is ever needed, that - * has to land in the agents library first (per-agent `AgentContext` - * sessions); changing only this helper would diverge from how the - * sandbox actually behaves at runtime. + * Skill files are immutable input resources for the run, but their storage + * pointers are deployment-local. Callers therefore pre-seed each exact + * profile partition; this helper never copies a pointer across partitions. * * **Walk order:** primary first, then `agentConfigs` (handoff/addedConvo) * in iteration order, then recurse into each config's @@ -156,8 +240,9 @@ export function buildInitialToolSessions(params: { const agent = queue.shift()!; if (visited.has(agent)) continue; visited.add(agent); + const sessionKey = agent.codeSessionKey ?? Constants.EXECUTE_CODE; if (agent.primedCodeFiles && agent.primedCodeFiles.length > 0) { - sessions = seedCodeFilesIntoSessions(agent.primedCodeFiles, sessions); + sessions = seedCodeFilesIntoSessions(agent.primedCodeFiles, sessions, sessionKey); } if (agent.subagentAgentConfigs && agent.subagentAgentConfigs.length > 0) { for (const child of agent.subagentAgentConfigs) { diff --git a/packages/api/src/agents/execution.spec.ts b/packages/api/src/agents/execution.spec.ts new file mode 100644 index 0000000000..9dc02e4c59 --- /dev/null +++ b/packages/api/src/agents/execution.spec.ts @@ -0,0 +1,97 @@ +import { resolveCodeExecutionContext } from './execution'; + +jest.mock('@librechat/agents', () => ({ + Constants: { EXECUTE_CODE: 'execute_code' }, + getCodeBaseURL: jest.fn(() => 'http://code-default.test/v1///'), +})); + +describe('resolveCodeExecutionContext', () => { + const originalStatefulUrl = process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + + afterEach(() => { + if (originalStatefulUrl == null) { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + return; + } + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = originalStatefulUrl; + }); + + it('uses the AWS-free default profile when stateful sessions are off', () => { + expect(resolveCodeExecutionContext({ statefulSessions: false })).toEqual({ + baseUrl: 'http://code-default.test/v1', + codeSessionKey: 'execute_code', + executionProfile: 'default', + statefulSessions: false, + }); + }); + + it('fails closed when a stateful agent has no stateful endpoint', () => { + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; + + expect(() => resolveCodeExecutionContext({ statefulSessions: true, userId: 'user-1' })).toThrow( + 'LIBRECHAT_CODE_BASEURL_STATEFUL is not configured', + ); + }); + + it('fails closed when a stateful agent has no authenticated user', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1/'; + + expect(() => resolveCodeExecutionContext({ statefulSessions: true })).toThrow( + 'authenticated user ID', + ); + }); + + it('defaults stateful agents to one environment per authenticated user', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1///'; + + expect(resolveCodeExecutionContext({ statefulSessions: true, userId: 'user-1' })).toEqual({ + baseUrl: 'http://code-stateful.test/v1', + codeSessionKey: 'execute_code:stateful:v2:user:b5729fb0e3ca12e7a61ff6857b99d98e', + executionProfile: 'stateful', + runtimeSessionHint: 'v2:user:b5729fb0e3ca12e7a61ff6857b99d98e', + statefulSessions: true, + }); + }); + + it('supports agent-user and conversation isolation', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + + expect( + resolveCodeExecutionContext({ + statefulSessions: true, + environment: 'agent-user', + userId: 'user-1', + agentId: 'agent-1', + }), + ).toEqual( + expect.objectContaining({ + runtimeSessionHint: 'v2:agent-user:9cf1605ead4951d96f711e1b3db86642', + codeSessionKey: 'execute_code:stateful:v2:agent-user:9cf1605ead4951d96f711e1b3db86642', + }), + ); + expect( + resolveCodeExecutionContext({ + statefulSessions: true, + environment: 'conversation', + userId: 'user-1', + conversationId: 'conversation-1', + }), + ).toEqual( + expect.objectContaining({ + runtimeSessionHint: 'v2:conversation:ea98cd74d68a59d7c8dd012a62580520', + codeSessionKey: 'execute_code:stateful:v2:conversation:ea98cd74d68a59d7c8dd012a62580520', + }), + ); + }); + + it('partitions every stateful environment by authenticated user', () => { + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; + + const first = resolveCodeExecutionContext({ statefulSessions: true, userId: 'user-1' }); + const second = resolveCodeExecutionContext({ statefulSessions: true, userId: 'user-2' }); + + expect(first.runtimeSessionHint).not.toBe(second.runtimeSessionHint); + expect(first.runtimeSessionHint).not.toContain('user-1'); + expect(second.runtimeSessionHint).not.toContain('user-2'); + }); +}); diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts new file mode 100644 index 0000000000..846b86c2c2 --- /dev/null +++ b/packages/api/src/agents/execution.ts @@ -0,0 +1,105 @@ +import { createHash } from 'node:crypto'; +import { Constants, getCodeBaseURL } from '@librechat/agents'; +import type { StatefulCodeEnvironment } from 'librechat-data-provider'; + +export const CODE_API_EXPECTED_PROFILE_HEADER = 'X-CodeAPI-Expected-Profile'; + +export type CodeExecutionProfile = 'default' | 'stateful'; + +export interface CodeExecutionContext { + baseUrl: string; + codeSessionKey: string; + executionProfile: CodeExecutionProfile; + runtimeSessionHint?: string; + statefulSessions: boolean; +} + +export function normalizeStatefulCodeEnvironment( + environment?: StatefulCodeEnvironment | string | null, +): StatefulCodeEnvironment { + if (environment === 'agent-user') { + return 'agent-user'; + } + if (environment === 'conversation') { + return 'conversation'; + } + return 'user'; +} + +export function getCodeExecutionBaseUrl(profile: CodeExecutionProfile): string { + if (profile === 'default') { + return getCodeBaseURL().replace(/\/+$/, ''); + } + const baseUrl = process.env.LIBRECHAT_CODE_BASEURL_STATEFUL?.trim().replace(/\/+$/, ''); + if (baseUrl) { + return baseUrl; + } + throw new Error( + 'Stateful code execution is enabled for this agent, but LIBRECHAT_CODE_BASEURL_STATEFUL is not configured.', + ); +} + +function resolveRuntimeSessionHint(params: { + environment: StatefulCodeEnvironment; + userId: string; + agentId?: string | null; + conversationId?: string | null; +}): string { + const { environment, userId, agentId, conversationId } = params; + const scopeFingerprint = (...parts: string[]): string => + createHash('sha256').update(JSON.stringify(parts)).digest('hex').slice(0, 32); + if (environment === 'agent-user') { + if (!agentId) { + throw new Error('Agent-user code environments require an agent ID.'); + } + return `v2:agent-user:${scopeFingerprint(userId, agentId)}`; + } + if (environment === 'conversation') { + if (!conversationId) { + throw new Error('Conversation code environments require a conversation ID.'); + } + return `v2:conversation:${scopeFingerprint(userId, conversationId)}`; + } + return `v2:user:${scopeFingerprint(userId)}`; +} + +export function resolveCodeExecutionContext(params: { + statefulSessions: boolean; + environment?: StatefulCodeEnvironment | string | null; + userId?: string | null; + agentId?: string | null; + conversationId?: string | null; +}): CodeExecutionContext { + if (!params.statefulSessions) { + return { + baseUrl: getCodeExecutionBaseUrl('default'), + codeSessionKey: Constants.EXECUTE_CODE, + executionProfile: 'default', + statefulSessions: false, + }; + } + + const environment = normalizeStatefulCodeEnvironment(params.environment); + if (!params.userId) { + throw new Error('Stateful code environments require an authenticated user ID.'); + } + const runtimeSessionHint = resolveRuntimeSessionHint({ + environment, + userId: params.userId, + agentId: params.agentId, + conversationId: params.conversationId, + }); + return { + baseUrl: getCodeExecutionBaseUrl('stateful'), + codeSessionKey: `${Constants.EXECUTE_CODE}:stateful:${runtimeSessionHint}`, + executionProfile: 'stateful', + runtimeSessionHint, + statefulSessions: true, + }; +} + +export function codeExecutionHeaders( + context: Pick, +): Record { + return { [CODE_API_EXPECTED_PROFILE_HEADER]: context.executionProfile }; +} diff --git a/packages/api/src/agents/handlers.background.spec.ts b/packages/api/src/agents/handlers.background.spec.ts index 918244c44c..f44c05b924 100644 --- a/packages/api/src/agents/handlers.background.spec.ts +++ b/packages/api/src/agents/handlers.background.spec.ts @@ -597,7 +597,14 @@ describe('createToolExecuteHandler — backgrounded code execution', () => { emitted.push(attachment); }, }); - const configurable = buildConfig(['execute_code']); + const codeExecutionContext = { + baseUrl: 'https://code-stateful.example.com', + codeSessionKey: 'execute_code:stateful:convo-hint', + executionProfile: 'stateful' as const, + runtimeSessionHint: 'convo-hint', + statefulSessions: true, + }; + const configurable = { ...buildConfig(['execute_code']), codeExecutionContext }; const metadata = { thread_id: 'exec_convo_code', run_id: 'msg-dispatch' }; const dispatch = await runBatch(handler, { @@ -636,6 +643,7 @@ describe('createToolExecuteHandler — backgrounded code execution', () => { dispatchedAt: expect.any(Number), output: 'stdout:\nhello', artifact: CODE_ARTIFACT, + codeExecutionContext, }), ); // nothing rode the finalized dispatch turn's callback diff --git a/packages/api/src/agents/handlers.spec.ts b/packages/api/src/agents/handlers.spec.ts index c3a9da3298..b450e8a8f2 100644 --- a/packages/api/src/agents/handlers.spec.ts +++ b/packages/api/src/agents/handlers.spec.ts @@ -1,3 +1,7 @@ +jest.mock('./prewarm', () => ({ + markSandboxReady: jest.fn(), +})); + import { Readable } from 'stream'; import { Constants } from '@librechat/agents'; import { logger } from '@librechat/data-schemas'; @@ -6,7 +10,9 @@ import type { ToolExecuteResult, ToolCallRequest, } from '@librechat/agents'; +import type { CodeExecutionContext } from './execution'; import { createToolExecuteHandler, ToolExecuteOptions } from './handlers'; +import { markSandboxReady } from './prewarm'; function createMockTool( name: string, @@ -2591,6 +2597,7 @@ describe('createToolExecuteHandler', () => { activeSkillNames?: Set; skillPrimedIdsByName?: Record; skillAuthoringAvailable?: boolean; + codeExecutionContext?: CodeExecutionContext; req?: unknown; readSandboxFile?: ToolExecuteOptions['readSandboxFile']; readSandboxImage?: ToolExecuteOptions['readSandboxImage']; @@ -2606,6 +2613,7 @@ describe('createToolExecuteHandler', () => { activeSkillNames: params.activeSkillNames, skillPrimedIdsByName: params.skillPrimedIdsByName, skillAuthoringAvailable: params.skillAuthoringAvailable === true, + codeExecutionContext: params.codeExecutionContext, }, })); return createToolExecuteHandler({ @@ -2646,6 +2654,102 @@ describe('createToolExecuteHandler', () => { expect(result.content).toContain('hello-world'); }); + it('routes host file reads with the executing agent profile instead of a graph hint', async () => { + const readSandboxFile = jest.fn(async () => ({ content: 'stateful-data' })); + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + codeSessionKey: 'execute_code:stateful:v1:user', + executionProfile: 'stateful', + runtimeSessionHint: 'v1:user', + statefulSessions: true, + }, + readSandboxFile, + }); + + await invokeHandler(handler, [ + { + id: 'call_profiled_read', + name: Constants.READ_FILE, + args: { path: '/mnt/data/sentinel.txt' }, + runtimeSessionHint: 'legacy-graph-hint', + } as unknown as ToolCallRequest, + ]); + + expect(readSandboxFile).toHaveBeenCalledWith( + expect.objectContaining({ + codeApiBaseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + runtime_session_hint: 'v1:user', + }), + ); + expect(readSandboxFile).not.toHaveBeenCalledWith( + expect.objectContaining({ runtime_session_hint: 'legacy-graph-hint' }), + ); + }); + + it('marks an actual sandbox read warm without marking skill-backed reads', async () => { + const readSandboxFile = jest.fn(async () => ({ content: 'stateful-data' })); + const context: CodeExecutionContext = { + baseUrl: 'https://stateful-code.example.com', + codeSessionKey: 'execute_code:stateful:v2:user:abc', + executionProfile: 'stateful', + runtimeSessionHint: 'v2:user:abc', + statefulSessions: true, + }; + const handler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + codeExecutionContext: context, + readSandboxFile, + }); + + const [result] = await new Promise((resolve, reject) => { + handler.handle('on_tool_execute', { + toolCalls: [ + { + id: 'call_warm_read', + name: Constants.READ_FILE, + args: { path: '/mnt/data/sentinel.txt' }, + }, + ], + metadata: { thread_id: 'conversation-1' }, + resolve, + reject, + } as ToolExecuteBatchRequest); + }); + + expect(result.status).toBe('success'); + expect(markSandboxReady).toHaveBeenCalledWith('v2:user:abc'); + expect(markSandboxReady).toHaveBeenCalledWith('conversation-1'); + + jest.mocked(markSandboxReady).mockClear(); + const skillHandler = makeReadFileHandler({ + codeEnvAvailable: true, + accessibleSkillIds: skillsInScope(), + activeSkillNames: new Set(['docs']), + codeExecutionContext: context, + getSkillByName: jest.fn(async () => ({ + _id: '507f1f77bcf86cd799439011' as never, + name: 'docs', + body: '# Docs', + fileCount: 0, + version: 1, + })), + }); + + await invokeHandler(skillHandler, [ + { + id: 'call_skill_read', + name: Constants.READ_FILE, + args: { path: 'docs/SKILL.md' }, + }, + ]); + expect(markSandboxReady).not.toHaveBeenCalled(); + }); + it('returns a clear error for /mnt/data/ when codeEnv is not available', async () => { const readSandboxFile = jest.fn(); const handler = makeReadFileHandler({ diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 04ae2b978b..3cda8f3837 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -15,6 +15,7 @@ import type { StructuredToolInterface } from '@librechat/agents/langchain/tools' import type { ValidationIssue } from '@librechat/data-schemas'; import type { CodeEnvRef } from 'librechat-data-provider'; import type { SkillFileRecord, PrimeSkillFilesResult } from './skillFiles'; +import type { CodeExecutionContext } from './execution'; import type { ServerRequest } from '~/types'; import { backgroundTaskRegistry, @@ -101,6 +102,7 @@ export interface ToolExecuteOptions { dispatchedAt?: number; output?: string; artifact?: unknown; + codeExecutionContext?: CodeExecutionContext; attachments?: unknown[]; reapply?: boolean; }) => Promise<{ attachments?: unknown[] } | null>; @@ -233,12 +235,21 @@ export interface ToolExecuteOptions { id: string; version?: number; read_only?: boolean; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; }) => Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }>; }>; /** Checks if a code env file is still active. Returns lastModified or null. */ - getSessionInfo?: (ref: CodeEnvRef, req?: ServerRequest) => Promise; + getSessionInfo?: ( + ref: CodeEnvRef, + req?: ServerRequest, + route?: { + baseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; + }, + ) => Promise; /** 23-hour freshness check */ checkIfActive?: (dateString: string) => boolean; /** Persists `codeEnvRef` on skill files after upload */ @@ -285,6 +296,8 @@ export interface ToolExecuteOptions { * host file op that is the first sandbox call joins the same runtime session * as bash_tool instead of the Code API's default session. */ runtime_session_hint?: string; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; req?: ServerRequest; }) => Promise<{ content: string } | null>; /** @@ -303,6 +316,8 @@ export interface ToolExecuteOptions { files?: Array<{ id: string; name: string; session_id?: string; storage_session_id?: string }>; /** @see readSandboxFile.runtime_session_hint */ runtime_session_hint?: string; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; /** In-sandbox size cap; files larger than this return `tooLarge` without transferring bytes. */ maxBytes?: number; req?: ServerRequest; @@ -320,6 +335,8 @@ export interface ToolExecuteOptions { files?: Array<{ id: string; name: string; session_id?: string; storage_session_id?: string }>; /** @see readSandboxFile.runtime_session_hint */ runtime_session_hint?: string; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; req?: ServerRequest; }) => Promise<{ stdout?: string; @@ -356,6 +373,40 @@ const MAX_SKILL_WARNING_MESSAGE_CHARS = 300; const IMAGE_MIMES = new Set(['image/png', 'image/jpeg', 'image/gif', 'image/webp']); +function getCodeExecutionContext( + configurable: Record, +): CodeExecutionContext | undefined { + const context = configurable.codeExecutionContext; + if (context == null || typeof context !== 'object') { + return undefined; + } + const candidate = context as Partial; + if ( + typeof candidate.baseUrl !== 'string' || + typeof candidate.codeSessionKey !== 'string' || + (candidate.executionProfile !== 'default' && candidate.executionProfile !== 'stateful') || + typeof candidate.statefulSessions !== 'boolean' + ) { + return undefined; + } + return candidate as CodeExecutionContext; +} + +function codeExecutionRequestParams(context?: CodeExecutionContext): { + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; + runtime_session_hint?: string; +} { + if (!context) { + return {}; + } + return { + codeApiBaseUrl: context.baseUrl, + executionProfile: context.executionProfile, + ...(context.runtimeSessionHint ? { runtime_session_hint: context.runtimeSessionHint } : {}), + }; +} + type ToolInputSchemaKind = { object: boolean; string: boolean; @@ -1452,6 +1503,8 @@ async function handleSandboxImageRead( ext: string, options: ToolExecuteOptions, req?: ServerRequest, + codeExecutionContext?: CodeExecutionContext, + onSuccess?: () => void, ): Promise { const { readSandboxImage } = options; const binaryHint = (): ToolExecuteResult => ({ @@ -1472,7 +1525,7 @@ async function handleSandboxImageRead( session_id: ctx?.session_id, files: ctx?.files, maxBytes: MAX_SANDBOX_INLINE_IMAGE_BYTES, - ...(tc.runtimeSessionHint ? { runtime_session_hint: tc.runtimeSessionHint } : {}), + ...codeExecutionRequestParams(codeExecutionContext), ...(req ? { req } : {}), }); } catch (error) { @@ -1485,6 +1538,7 @@ async function handleSandboxImageRead( return binaryHint(); } if ('tooLarge' in read) { + onSuccess?.(); return { toolCallId: tc.id, status: 'success', @@ -1508,6 +1562,7 @@ async function handleSandboxImageRead( if (!mimeType || !isCompleteImage(buffer, mimeType)) { return binaryHint(); } + onSuccess?.(); return buildImageArtifactResult(tc.id, filePath, mimeType, buffer.length, read.base64); } @@ -1536,10 +1591,12 @@ async function handleSandboxFileFallback( filePath: string, options: ToolExecuteOptions, req?: ServerRequest, + codeExecutionContext?: CodeExecutionContext, + onSuccess?: () => void, ): Promise { const ext = lowercaseExtension(filePath); if (SANDBOX_IMAGE_EXTENSIONS.has(ext)) { - return handleSandboxImageRead(tc, filePath, ext, options, req); + return handleSandboxImageRead(tc, filePath, ext, options, req, codeExecutionContext, onSuccess); } if (BINARY_EXTENSIONS_NEVER_READABLE.has(ext)) { return { @@ -1566,7 +1623,7 @@ async function handleSandboxFileFallback( file_path: filePath, session_id: ctx?.session_id, files: ctx?.files, - ...(tc.runtimeSessionHint ? { runtime_session_hint: tc.runtimeSessionHint } : {}), + ...codeExecutionRequestParams(codeExecutionContext), ...(req ? { req } : {}), }); if (!result || result.content == null) { @@ -1604,6 +1661,7 @@ async function handleSandboxFileFallback( if (truncated) { numbered += `\n\n[truncated at ${MAX_READABLE_BYTES} bytes — use \`bash_tool\` (e.g. \`head -c\` / \`tail\`) to read the rest of "${filePath}"]`; } + onSuccess?.(); return { toolCallId: tc.id, status: 'success', @@ -1711,12 +1769,14 @@ async function loadSandboxTextForAuthoring({ options, req, sandboxContext, + codeExecutionContext, }: { filePath: string; tc: ToolCallRequest; options: ToolExecuteOptions; req?: ServerRequest; sandboxContext?: SandboxSessionContext; + codeExecutionContext?: CodeExecutionContext; }): Promise { const ext = lowercaseExtension(filePath); if (BINARY_EXTENSIONS_NEVER_READABLE.has(ext)) { @@ -1735,7 +1795,7 @@ async function loadSandboxTextForAuthoring({ file_path: filePath, session_id: ctx?.session_id, files: ctx?.files, - ...(tc.runtimeSessionHint ? { runtime_session_hint: tc.runtimeSessionHint } : {}), + ...codeExecutionRequestParams(codeExecutionContext), ...(req ? { req } : {}), }); if (!result || result.content == null) { @@ -1786,6 +1846,7 @@ async function writeSandboxTextForAuthoring({ oldContent, created, sandboxContext, + codeExecutionContext, }: { tc: ToolCallRequest; options: ToolExecuteOptions; @@ -1795,6 +1856,7 @@ async function writeSandboxTextForAuthoring({ oldContent?: string; created: boolean; sandboxContext?: SandboxSessionContext; + codeExecutionContext?: CodeExecutionContext; }): AuthoringResult { if (!options.writeSandboxFile) { return errorResult( @@ -1810,7 +1872,7 @@ async function writeSandboxTextForAuthoring({ content, session_id: ctx?.session_id, files: ctx?.files, - ...(tc.runtimeSessionHint ? { runtime_session_hint: tc.runtimeSessionHint } : {}), + ...codeExecutionRequestParams(codeExecutionContext), ...(req ? { req } : {}), }); } catch (error) { @@ -2583,6 +2645,7 @@ async function handleSandboxCreateFileCall({ content, overwrite, sandboxContext, + codeExecutionContext, }: { tc: ToolCallRequest; options: ToolExecuteOptions; @@ -2591,6 +2654,7 @@ async function handleSandboxCreateFileCall({ content: string; overwrite: boolean; sandboxContext?: SandboxSessionContext; + codeExecutionContext?: CodeExecutionContext; }): AuthoringResult { const pathError = invalidSandboxAuthoringPath(filePath); if (pathError) { @@ -2603,6 +2667,7 @@ async function handleSandboxCreateFileCall({ options, req, sandboxContext, + codeExecutionContext, }); if (current.status === 'error') { return errorResult(tc, current.message); @@ -2620,6 +2685,7 @@ async function handleSandboxCreateFileCall({ oldContent: current.status === 'loaded' ? current.content : undefined, created: current.status === 'missing', sandboxContext, + codeExecutionContext, }); } @@ -2630,6 +2696,7 @@ async function handleSandboxEditFileCall({ filePath, edits, sandboxContext, + codeExecutionContext, }: { tc: ToolCallRequest; options: ToolExecuteOptions; @@ -2637,6 +2704,7 @@ async function handleSandboxEditFileCall({ filePath: string; edits: TextEdit[]; sandboxContext?: SandboxSessionContext; + codeExecutionContext?: CodeExecutionContext; }): AuthoringResult { const pathError = invalidSandboxAuthoringPath(filePath); if (pathError) { @@ -2649,6 +2717,7 @@ async function handleSandboxEditFileCall({ options, req, sandboxContext, + codeExecutionContext, }); if (current.status === 'missing') { return errorResult(tc, `File not found: "${filePath}"`); @@ -2676,6 +2745,7 @@ async function handleSandboxEditFileCall({ oldContent: current.content, created: false, sandboxContext, + codeExecutionContext, }); if (result.status === 'success') { result.artifact = { @@ -2728,6 +2798,7 @@ async function handleCreateFileCall( content: args.content, overwrite, sandboxContext, + codeExecutionContext: getCodeExecutionContext(mergedConfigurable), }); } @@ -2837,6 +2908,7 @@ async function handleEditFileCall( filePath: args.path, edits, sandboxContext, + codeExecutionContext: getCodeExecutionContext(mergedConfigurable), }); } @@ -2939,6 +3011,7 @@ async function handleReadFileCall( mergedConfigurable: Record, options: ToolExecuteOptions, req?: ServerRequest, + onSandboxReadSuccess?: () => void, ): Promise { const { getSkillByName, getSkillFileByPath, getStrategyFunctions, updateSkillFileContent } = options; @@ -2953,6 +3026,7 @@ async function handleReadFileCall( } const codeEnvAvailable = mergedConfigurable?.codeEnvAvailable === true; + const codeExecutionContext = getCodeExecutionContext(mergedConfigurable); let accessibleIds = (mergedConfigurable?.accessibleSkillIds as Types.ObjectId[]) ?? []; /** @@ -2962,7 +3036,14 @@ async function handleReadFileCall( */ if (args.path.startsWith('/mnt/data/')) { if (codeEnvAvailable) { - return handleSandboxFileFallback(tc, args.path, options, req); + return handleSandboxFileFallback( + tc, + args.path, + options, + req, + codeExecutionContext, + onSandboxReadSuccess, + ); } return { toolCallId: tc.id, @@ -2991,7 +3072,14 @@ async function handleReadFileCall( const slashIdx = args.path.indexOf('/'); if (slashIdx < 1) { if (codeEnvAvailable) { - return handleSandboxFileFallback(tc, args.path, options, req); + return handleSandboxFileFallback( + tc, + args.path, + options, + req, + codeExecutionContext, + onSandboxReadSuccess, + ); } return { toolCallId: tc.id, @@ -3011,7 +3099,14 @@ async function handleReadFileCall( * dead-ending with a skill-centric error message. */ if (codeEnvAvailable) { - return handleSandboxFileFallback(tc, args.path, options, req); + return handleSandboxFileFallback( + tc, + args.path, + options, + req, + codeExecutionContext, + onSandboxReadSuccess, + ); } return { toolCallId: tc.id, @@ -3073,7 +3168,14 @@ async function handleReadFileCall( */ if (!skillsEffectivelyEnabled) { if (codeEnvAvailable && !explicitSkillNamespace) { - return handleSandboxFileFallback(tc, args.path, options, req); + return handleSandboxFileFallback( + tc, + args.path, + options, + req, + codeExecutionContext, + onSandboxReadSuccess, + ); } return { toolCallId: tc.id, @@ -3112,7 +3214,14 @@ async function handleReadFileCall( const recovered = await recoverAuthorSkill(); if (!recovered) { if (codeEnvAvailable && !explicitSkillNamespace) { - return handleSandboxFileFallback(tc, args.path, options, req); + return handleSandboxFileFallback( + tc, + args.path, + options, + req, + codeExecutionContext, + onSandboxReadSuccess, + ); } return { toolCallId: tc.id, @@ -3495,6 +3604,7 @@ async function handleSkillToolCall( // is enabled for this run. The flag is threaded via configurable upstream // so this gate cannot be bypassed. const codeEnvAvailable = mergedConfigurable?.codeEnvAvailable === true; + const codeExecutionContext = getCodeExecutionContext(mergedConfigurable); if ( codeEnvAvailable && skill.fileCount > 0 && @@ -3515,6 +3625,7 @@ async function handleSkillToolCall( getSessionInfo, checkIfActive, updateSkillFileCodeEnvIds, + codeExecutionContext, }); if (primeResult) { /* `session_id` at the top of the artifact is the (representative) @@ -3744,6 +3855,24 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand sourceConfigurable, loadedConfigurable, ); + const codeExecutionContext = getCodeExecutionContext(mergedConfigurable); + const runtimeSessionHint = codeExecutionContext?.runtimeSessionHint; + const sandboxConversationId = + ((metadata as Record)?.thread_id as string | undefined) ?? + (mergedConfigurable?.thread_id as string | undefined) ?? + ( + (mergedConfigurable?.req as ServerRequest | undefined)?.body as + | { conversationId?: string } + | undefined + )?.conversationId; + const markCodeSandboxWarm = (): void => { + if (runtimeSessionHint) { + void markSandboxReady(runtimeSessionHint); + } + if (sandboxConversationId) { + void markSandboxReady(sandboxConversationId); + } + }; const authoringQueues = new Map>(); const sandboxAuthoringContexts = new Map(); @@ -3854,6 +3983,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand * after a newer run wrote the same filename must not * overwrite it. */ dispatchedAt: task.createdAt, + codeExecutionContext, ...params, }); if (persisted == null) { @@ -3903,8 +4033,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand configurable: mergedConfigurable, metadata, } as Record)) as { content?: unknown; artifact?: unknown }; - if (tc.runtimeSessionHint != null && tc.runtimeSessionHint !== '') { - void markSandboxReady(tc.runtimeSessionHint); + if (isCodeCall) { + markCodeSandboxWarm(); } const content = isCodeCall && typeof result.content === 'string' @@ -3997,7 +4127,10 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand artifact: pending.artifact, }, }, - (metadata ?? {}) as ToolEndCallbackMetadata, + { + ...(metadata ?? {}), + executingAgentId: agentId, + } as ToolEndCallbackMetadata, ); } catch (callbackError) { /** Only synchronous callback throws land here (e.g. a @@ -4135,6 +4268,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand isFileAuthoringCall && typeof (tc.args as { path?: unknown }).path === 'string' && !(tc.args as { path: string }).path.startsWith(SKILL_FILE_PREFIX); + let sandboxReadSucceeded = false; if ( tc.name === Constants.SKILL_TOOL || tc.name === Constants.READ_FILE || @@ -4156,6 +4290,9 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand mergedConfigurable, options, req, + () => { + sandboxReadSucceeded = true; + }, ); } else if (tc.name === CREATE_FILE_TOOL_NAME && isFileAuthoringCall) { handlerResult = await handleCreateFileCall( @@ -4217,23 +4354,23 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand | string | undefined, ...metadata, + executingAgentId: agentId, + codeExecutionContext, }, ); } - /* Sandbox-routed create_file/edit_file return before the + /* Sandbox-routed host file tools return before the * generic invoke path's marker below, so refresh the warm - * window here. Gated on `isSandboxFileAuthoringCall`: - * skill-path writes and skill/read_file calls on this - * branch may resolve without touching the Code API, and - * under-marking only costs a redundant cold-boot label. */ + * window here. `sandboxReadSucceeded` is set only after an + * actual Code API read succeeds, so skill reads never mark + * the sandbox warm. */ if ( - isSandboxFileAuthoringCall && + (isSandboxFileAuthoringCall || sandboxReadSucceeded) && handlerResult.status === 'success' && - tc.runtimeSessionHint != null && - tc.runtimeSessionHint !== '' + (runtimeSessionHint || sandboxConversationId) ) { - void markSandboxReady(tc.runtimeSessionHint); + markCodeSandboxWarm(); } return handlerResult; @@ -4321,8 +4458,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand /* Only sandbox-bound calls carry a runtime session hint, so * this refreshes the prewarm module's warm window without * inspecting tool names. */ - if (tc.runtimeSessionHint != null && tc.runtimeSessionHint !== '') { - void markSandboxReady(tc.runtimeSessionHint); + if (isCodeSessionAwareToolCall(tc.name, mergedConfigurable)) { + markCodeSandboxWarm(); } // Code-execution tools emit per-call boilerplate @@ -4356,6 +4493,8 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand | string | undefined, ...metadata, + executingAgentId: agentId, + codeExecutionContext, }, ); } diff --git a/packages/api/src/agents/harvest.ts b/packages/api/src/agents/harvest.ts index 0ffabdc31e..a25846c1ba 100644 --- a/packages/api/src/agents/harvest.ts +++ b/packages/api/src/agents/harvest.ts @@ -1,4 +1,5 @@ import { logger } from '@librechat/data-schemas'; +import type { CodeExecutionContext } from './execution'; import type { ServerRequest } from '~/types'; /** @@ -52,6 +53,8 @@ export interface CodeHarvestDeps { agentId?: string; session_id?: string; freshClaimAfter?: number; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; }) => Promise; /** Host file service: runs the deferred office-preview extraction. */ runPreviewFinalize: (params: { @@ -75,6 +78,7 @@ export interface CodeHarvestParams { dispatchedAt?: number; output?: string; artifact?: unknown; + codeExecutionContext?: CodeExecutionContext; attachments?: unknown[]; reapply?: boolean; } @@ -111,6 +115,7 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa dispatchedAt, output, artifact, + codeExecutionContext, attachments: knownAttachments, reapply, }) => { @@ -161,6 +166,8 @@ export function createBackgroundCodeResultHandler(deps: CodeHarvestDeps): CodeHa agentId, session_id: file.storage_session_id ?? codeArtifact.session_id, freshClaimAfter, + codeApiBaseUrl: codeExecutionContext?.baseUrl, + executionProfile: codeExecutionContext?.executionProfile, }); if (result?.file) { attachments.push(result.file); diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index b7fa363488..c12d2bc93e 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -11,6 +11,7 @@ export * from './discovery'; export * from './edges'; export * from './errors'; export * from './envelope'; +export * from './execution'; export * from './handlers'; export * from './harvest'; export * from './initialize'; diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 03b21a0e17..c38d96260d 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -51,6 +51,11 @@ import { splitMCPToolKey, normalizeAgentToolKeys, } from '~/mcp/utils'; +import { + normalizeStatefulCodeEnvironment, + resolveCodeExecutionContext, + type CodeExecutionContext, +} from './execution'; import { optionalChainWithEmptyCheck, extractLibreChatParams, @@ -325,10 +330,16 @@ export type InitializedAgent = Agent & { * Whether stateful code sessions are active *for this agent*: the admin * `stateful_code_sessions` capability AND the agent's builder opt-in * (`agent.stateful_code_sessions`) AND `codeEnvAvailable`. Resolved once - * here; `createRun` walks this per-agent value to gate the run-level - * `toolExecution.sandbox` config. + * here and carried with the agent so execution routing never needs a + * graph-global stateful flag. */ statefulCodeSessions: boolean; + /** Sharing scope for this agent's stateful code environment. */ + statefulCodeEnvironment: Agent['stateful_code_environment']; + /** Trusted partition for transient code session ids and file references. */ + codeSessionKey: string; + /** Trusted endpoint/profile context for artifact processing and runtime tools. */ + codeExecutionContext: CodeExecutionContext; /** Whether host-side skill file authoring is available for this agent/run. */ skillAuthoringAvailable: boolean; /** Host-side file authoring tool names registered for this run. */ @@ -410,6 +421,8 @@ export interface InitializeAgentParams { model: string | null; tool_options: AgentToolOptions | undefined; tool_resources: AgentToolResources | undefined; + /** Trusted endpoint/profile resolved for this agent before any code-file priming. */ + codeExecutionContext: CodeExecutionContext; /** Full accessible MCP server names (operator + user DB) when the heal * already fetched them — lets execution-side collision guards see * cross-tier shadowing without another registry round-trip. */ @@ -700,6 +713,24 @@ export async function initializeAgent( const provider = agent.provider; agent.endpoint = provider; + /** Resolve the per-agent Code API route before resource/tool priming. A + * stateful agent must perform freshness checks and recovery uploads against + * the same isolated deployment its eventual `/exec` request will use. */ + const agentRequestsCodeExec = (agent.tools ?? []).includes(Tools.execute_code); + const effectiveCodeEnvAvailable = params.codeEnvAvailable === true && agentRequestsCodeExec; + const effectiveStatefulSessions = + effectiveCodeEnvAvailable && + params.statefulSessionsAvailable === true && + agent.stateful_code_sessions === true; + const statefulCodeEnvironment = normalizeStatefulCodeEnvironment(agent.stateful_code_environment); + const codeExecutionContext = resolveCodeExecutionContext({ + statefulSessions: effectiveStatefulSessions, + environment: statefulCodeEnvironment, + userId: requestFileOwnerId, + agentId: agent.id, + conversationId, + }); + /** * Load conversation files for ALL agents, not just the initial agent. * This enables handoff agents to access files that were uploaded earlier @@ -1017,6 +1048,7 @@ export async function initializeAgent( model: agent.model, tool_options: agent.tool_options, tool_resources, + codeExecutionContext, accessibleMcpServerNames: resolvedAuditNames, }); @@ -1181,8 +1213,6 @@ export async function initializeAgent( * code-only description to the skill-aware description without adding a * duplicate — exactly one copy of each tool reaches the LLM. */ - const agentRequestsCodeExec = (agent.tools ?? []).includes(Tools.execute_code); - const effectiveCodeEnvAvailable = params.codeEnvAvailable === true && agentRequestsCodeExec; /** * Capability marker → definition names its registration produced this run, * reported by the registrars themselves. `tool_options` entries keyed by a @@ -1199,14 +1229,6 @@ export async function initializeAgent( const existing = capabilityToolNames.get(capability); capabilityToolNames.set(capability, existing ? [...existing, ...toolNames] : toolNames); }; - /** Per-agent stateful-session truth: the admin capability AND the agent's - * own builder opt-in AND a working code env. Resolved once here so the - * registered bash description, the tool factories, and `createRun`'s - * `toolExecution.sandbox` gate all agree for this agent. */ - const effectiveStatefulSessions = - effectiveCodeEnvAvailable && - params.statefulSessionsAvailable === true && - agent.stateful_code_sessions === true; if (effectiveCodeEnvAvailable) { const codeExecResult = registerCodeExecutionTools({ toolRegistry, @@ -1527,6 +1549,9 @@ export async function initializeAgent( memoryToolsRegistered: inlineMemoryRegistered, codeEnvAvailable: effectiveCodeEnvAvailable, statefulCodeSessions: effectiveStatefulSessions, + statefulCodeEnvironment, + codeSessionKey: codeExecutionContext.codeSessionKey, + codeExecutionContext, reasoningKey: customEndpointConfig?.customParams?.reasoningKey, includeReasoningHistory: customEndpointConfig?.customParams?.includeReasoningHistory, skillAuthoringAvailable, diff --git a/packages/api/src/agents/lazySubagents.ts b/packages/api/src/agents/lazySubagents.ts index 0679ec76df..1d12a41894 100644 --- a/packages/api/src/agents/lazySubagents.ts +++ b/packages/api/src/agents/lazySubagents.ts @@ -19,6 +19,7 @@ type VersionedAgent = Pick< | 'skills' | 'skills_enabled' | 'stateful_code_sessions' + | 'stateful_code_environment' | 'artifacts' | 'recursion_limit' | 'agent_ids' @@ -87,6 +88,7 @@ export function selectLazySubagentConfig(agent: VersionedAgent): Omit[0]; interface TestAgent { id: string; statefulCodeSessions?: boolean; + statefulCodeEnvironment?: StatefulCodeEnvironment; subagentAgentConfigs?: TestAgent[]; + lazySubagentConfigs?: TestAgent[]; } -const req = {} as PrewarmParams['req']; +const req = { user: { id: 'user-1' } } as PrewarmParams['req']; const statefulAgent: TestAgent = { id: 'agent_stateful', statefulCodeSessions: true }; const plainAgent: TestAgent = { id: 'agent_plain', statefulCodeSessions: false }; @@ -35,6 +38,7 @@ describe('maybePrewarmCodeSandbox', () => { beforeEach(async () => { await resetSandboxStateForTests(); process.env.LIBRECHAT_CODE_BASEURL = 'http://code.test/v1'; + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; delete process.env.CODE_SANDBOX_PREWARM; delete process.env.CODE_SANDBOX_COLD_AFTER_MS; delete process.env.CODEAPI_JWT_ENABLED; @@ -48,6 +52,7 @@ describe('maybePrewarmCodeSandbox', () => { fetchMock.mockRestore(); jest.useRealTimers(); delete process.env.LIBRECHAT_CODE_BASEURL; + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; }); it('does nothing when no reachable agent has stateful sessions', async () => { @@ -70,17 +75,20 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).not.toHaveBeenCalled(); }); - it('fires one exec with the conversation as runtime_session_hint and marks ready', async () => { + it('fires one stateful-profile exec with the default user environment and marks ready', async () => { maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); expect(fetchMock).toHaveBeenCalledTimes(1); const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; - expect(url).toBe('http://code.test/v1/exec'); + expect(url).toBe('http://code-stateful.test/v1/exec'); + expect(init.headers).toEqual( + expect.objectContaining({ 'X-CodeAPI-Expected-Profile': 'stateful' }), + ); expect(JSON.parse(init.body as string)).toEqual({ lang: 'bash', code: 'true', - runtime_session_hint: 'convo-1', + runtime_session_hint: 'v2:user:b5729fb0e3ca12e7a61ff6857b99d98e', }); await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); }); @@ -92,6 +100,94 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).toHaveBeenCalledTimes(1); }); + it('walks lazy subagents and deduplicates user-scoped environments', async () => { + const parent = { + id: 'agent_parent', + statefulCodeSessions: true, + lazySubagentConfigs: [statefulAgent], + }; + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(parent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it('prewarms distinct per-agent environments independently', async () => { + const first: TestAgent = { + id: 'agent-1', + statefulCodeSessions: true, + statefulCodeEnvironment: 'agent-user', + }; + const second: TestAgent = { + id: 'agent-2', + statefulCodeSessions: true, + statefulCodeEnvironment: 'agent-user', + }; + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(first, second) }); + await flushAsync(); + + expect(fetchMock).toHaveBeenCalledTimes(2); + const hints = fetchMock.mock.calls.map( + ([, init]) => JSON.parse(init.body).runtime_session_hint, + ); + expect(hints).toEqual( + expect.arrayContaining([ + 'v2:agent-user:9cf1605ead4951d96f711e1b3db86642', + 'v2:agent-user:f2a396a5aa5e99ce8e423f5ba6c323a3', + ]), + ); + }); + + it('does not share prewarm cache entries between authenticated users', async () => { + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + const otherReq = { user: { id: 'user-2' } } as PrewarmParams['req']; + maybePrewarmCodeSandbox({ + req: otherReq, + conversationId: 'convo-2', + agents: agents(statefulAgent), + }); + await flushAsync(); + + expect(fetchMock).toHaveBeenCalledTimes(2); + const hints = fetchMock.mock.calls.map( + ([, init]) => JSON.parse(init.body).runtime_session_hint, + ); + expect(new Set(hints).size).toBe(2); + }); + + it('keeps the conversation start signal active until every selected environment is warm', async () => { + const resolvers: Array<(response: Response) => void> = []; + fetchMock.mockImplementation( + () => + new Promise((resolve) => { + resolvers.push(resolve); + }), + ); + const first = { + id: 'agent-1', + statefulCodeSessions: true, + statefulCodeEnvironment: 'agent-user' as const, + }; + const second = { + id: 'agent-2', + statefulCodeSessions: true, + statefulCodeEnvironment: 'agent-user' as const, + }; + + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(first, second) }); + await flushAsync(); + expect(resolvers).toHaveLength(2); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + + resolvers[0](mockResponse({ ok: true, status: 200 })); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + + resolvers[1](mockResponse({ ok: true, status: 200 })); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); + }); + it('does not refire while the warm marker is fresh', async () => { maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); @@ -109,6 +205,18 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).toHaveBeenCalledTimes(1); }); + it('keeps a second conversation cold while it joins an in-flight user prewarm', async () => { + fetchMock.mockImplementation(() => new Promise(() => undefined)); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-2', agents: agents(statefulAgent) }); + await flushAsync(); + + expect(fetchMock).toHaveBeenCalledTimes(1); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + await expect(shouldSignalSandboxStart('convo-2')).resolves.toBe(true); + }); + it('refires once the warm marker has expired', async () => { jest.useFakeTimers({ doNotFake: ['setImmediate'] }); jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); @@ -184,12 +292,14 @@ describe('maybePrewarmCodeSandbox', () => { describe('shouldSignalSandboxStart / markSandboxReady', () => { beforeEach(async () => { await resetSandboxStateForTests(); + process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'http://code-stateful.test/v1'; delete process.env.CODE_SANDBOX_PREWARM; delete process.env.CODE_SANDBOX_COLD_AFTER_MS; }); afterEach(() => { jest.useRealTimers(); + delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL; }); it('never signals for untracked conversations (stateless deployments)', async () => { diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts index 0e202cd54e..7b04b3904c 100644 --- a/packages/api/src/agents/prewarm.ts +++ b/packages/api/src/agents/prewarm.ts @@ -1,13 +1,25 @@ import { logger } from '@librechat/data-schemas'; -import { getCodeBaseURL } from '@librechat/agents'; import { CacheKeys } from 'librechat-data-provider'; +import type { StatefulCodeEnvironment } from 'librechat-data-provider'; import type { Keyv } from 'keyv'; import type { ServerRequest } from '~/types'; +import { + codeExecutionHeaders, + resolveCodeExecutionContext, + type CodeExecutionContext, +} from './execution'; import { getCodeApiAuthHeaders } from '~/auth/codeapi'; import { standardCache } from '~/cache/cacheFactory'; -import { anyAgentHasStatefulSessions } from './run'; -type PrewarmAgents = Parameters[0]; +type PrewarmAgent = { + id: string; + statefulCodeSessions?: boolean; + statefulCodeEnvironment?: StatefulCodeEnvironment; + subagentAgentConfigs?: PrewarmAgent[]; + lazySubagentConfigs?: PrewarmAgent[]; +}; + +type PrewarmAgents = Array; const PREWARM_INFLIGHT_COOLDOWN_MS = 120_000; const PREWARM_REQUEST_TIMEOUT_MS = 120_000; @@ -29,9 +41,10 @@ function prewarmDisabled(): boolean { } /** - * Per-conversation sandbox state, shared across replicas when Redis is - * configured and falling back to a process-local store otherwise. Two keys - * per conversation (= runtime_session_hint): + * Sandbox state, shared across replicas when Redis is configured and falling + * back to a process-local store otherwise. Runtime-session keys include a + * one-way fingerprint of the authenticated user so they cannot collide across + * users; the conversation key below controls only that conversation's UI signal. * - `inflight:` — a prewarm was fired and no completion has landed yet; * the TTL doubles as the retry backoff when a prewarm fails or hangs. * - `ready:` — the sandbox completed a request (prewarm or real exec) @@ -90,16 +103,24 @@ export async function shouldSignalSandboxStart(conversationId?: string | null): return inflight != null && ready == null; } -async function sendPrewarmRequest(req: ServerRequest, conversationId: string): Promise { +async function sendPrewarmRequest( + req: ServerRequest, + context: CodeExecutionContext, +): Promise { const authHeaders = await getCodeApiAuthHeaders(req); - const response = await fetch(`${getCodeBaseURL()}/exec`, { + const response = await fetch(`${context.baseUrl}/exec`, { method: 'POST', headers: { 'Content-Type': 'application/json', 'User-Agent': 'LibreChat/1.0', ...authHeaders, + ...codeExecutionHeaders(context), }, - body: JSON.stringify({ lang: 'bash', code: 'true', runtime_session_hint: conversationId }), + body: JSON.stringify({ + lang: 'bash', + code: 'true', + runtime_session_hint: context.runtimeSessionHint, + }), signal: AbortSignal.timeout(PREWARM_REQUEST_TIMEOUT_MS), }); if (!response.ok) { @@ -112,12 +133,71 @@ async function sendPrewarmRequest(req: ServerRequest, conversationId: string): P * Draining also releases the socket instead of leaving the body for * undici to reap. */ await response.arrayBuffer(); - await markSandboxReady(conversationId); - logger.debug(`[prewarmCodeSandbox] Sandbox warm for conversation ${conversationId}`); + await markSandboxReady(context.runtimeSessionHint ?? ''); + logger.debug(`[prewarmCodeSandbox] Sandbox warm for ${context.runtimeSessionHint}`); +} + +function collectPrewarmContexts( + agents: PrewarmAgents, + conversationId: string, + userId: string, +): CodeExecutionContext[] { + const visited = new Set(); + const contexts = new Map(); + const pending: PrewarmAgents = [...agents]; + + for (let index = 0; index < pending.length; index++) { + const agent = pending[index]; + if (!agent || visited.has(agent.id)) { + continue; + } + visited.add(agent.id); + if (agent.statefulCodeSessions === true) { + const context = resolveCodeExecutionContext({ + statefulSessions: true, + environment: agent.statefulCodeEnvironment, + userId, + agentId: agent.id, + conversationId, + }); + contexts.set(`${context.baseUrl}:${context.runtimeSessionHint}`, context); + } + pending.push(...(agent.subagentAgentConfigs ?? []), ...(agent.lazySubagentConfigs ?? [])); + } + return [...contexts.values()]; +} + +async function maybePrewarmContext( + req: ServerRequest, + context: CodeExecutionContext, + conversationId: string, +): Promise { + const runtimeSessionHint = context.runtimeSessionHint; + if (!runtimeSessionHint) { + return true; + } + const cache = sandboxCache(); + const [ready, inflight] = await Promise.all([ + cache.get(readyKey(runtimeSessionHint)), + cache.get(inflightKey(runtimeSessionHint)), + ]); + if (ready != null) { + return true; + } + if (inflight != null) { + await cache.set(inflightKey(conversationId), true, PREWARM_INFLIGHT_COOLDOWN_MS); + return false; + } + await Promise.all([ + cache.set(inflightKey(runtimeSessionHint), true, PREWARM_INFLIGHT_COOLDOWN_MS), + cache.set(inflightKey(conversationId), true, PREWARM_INFLIGHT_COOLDOWN_MS), + ]); + await sendPrewarmRequest(req, context); + return true; } /** - * Fire-and-forget boot of the per-conversation stateful code sandbox so it + * Fire-and-forget boot of each selected stateful code environment so it * comes up in parallel with model generation instead of on the first * execute_code/bash call (~4s cold, worse on heavy first imports). No-op * unless a reachable agent resolved `statefulCodeSessions` and neither a @@ -134,23 +214,24 @@ export function maybePrewarmCodeSandbox(params: { agents: PrewarmAgents; }): void { const { req, conversationId, agents } = params; - if (prewarmDisabled() || !conversationId || !anyAgentHasStatefulSessions(agents)) { + if (prewarmDisabled() || !conversationId) { return; } void (async () => { - const cache = sandboxCache(); - const [ready, inflight] = await Promise.all([ - cache.get(readyKey(conversationId)), - cache.get(inflightKey(conversationId)), - ]); - if (ready != null || inflight != null) { - return; + const userId = req.user?.id; + if (!userId) { + throw new Error('Stateful code prewarm requires an authenticated user ID.'); + } + const contexts = collectPrewarmContexts(agents, conversationId, userId); + const ready = await Promise.all( + contexts.map((context) => maybePrewarmContext(req, context, conversationId)), + ); + if (ready.length > 0 && ready.every(Boolean)) { + await markSandboxReady(conversationId); } - await cache.set(inflightKey(conversationId), true, PREWARM_INFLIGHT_COOLDOWN_MS); - await sendPrewarmRequest(req, conversationId); })().catch((error) => { logger.debug( - `[prewarmCodeSandbox] Prewarm failed for conversation ${conversationId}: ${ + `[prewarmCodeSandbox] Prewarm failed: ${ error instanceof Error ? error.message : String(error) }`, ); diff --git a/packages/api/src/agents/resources.ts b/packages/api/src/agents/resources.ts index a7c6e9546b..c66f51bbca 100644 --- a/packages/api/src/agents/resources.ts +++ b/packages/api/src/agents/resources.ts @@ -119,7 +119,7 @@ const categorizeFileForToolResources = ({ requestFileSet: Set; processedResourceFiles: Set; }): void => { - if (file.metadata?.codeEnvRef) { + if (file.metadata?.codeEnvRef || file.metadata?.codeEnvRefs) { addFileToResource({ file, resourceType: EToolResources.execute_code, diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index e59fd6ee50..e9f7f19c51 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -55,6 +55,7 @@ import { isSteeringSupported, isSteerPreemptSupported } from '~/agents/steering/ import { getLLMConfig as getAnthropicLLMConfig } from '~/endpoints/anthropic/llm'; import { resolveStreamLimits, resolveSubagentMaxTurns } from '~/agents/config'; import { CREATE_FILE_TOOL_NAME, EDIT_FILE_TOOL_NAME } from '~/agents/tools'; +import { buildAgentInitialToolSessions } from '~/agents/codeFilesSession'; import { getProviderConfig } from '~/endpoints/config/providers'; import { extractDefaultParams } from '~/endpoints/openai/llm'; import { resolveHeaders, createSafeUser } from '~/utils/env'; @@ -383,9 +384,13 @@ type RunAgent = Omit & { /** * Per-agent stateful-session gate set by `initializeAgent`: the admin * `stateful_code_sessions` capability AND the agent's builder opt-in AND - * `codeEnvAvailable`. Walked here to gate `toolExecution.sandbox`. + * `codeEnvAvailable`. Carried into per-agent tool loading and prewarming. */ statefulCodeSessions?: boolean; + /** Per-agent stateful workspace sharing scope. */ + statefulCodeEnvironment?: Agent['stateful_code_environment']; + /** Trusted partition for transient code session ids and file references. */ + codeSessionKey?: string; /** Optional per-agent summarization overrides */ summarization?: SummarizationConfig; /** Response field to read model reasoning from for custom OpenAI-compatible endpoints. */ @@ -421,6 +426,8 @@ type LazySubagentAgent = Pick< | 'subagents' | 'codeEnvAvailable' | 'statefulCodeSessions' + | 'statefulCodeEnvironment' + | 'codeSessionKey' | 'includeReasoningHistory' > & { configId: string; @@ -437,6 +444,8 @@ type SubagentTreeNode = Pick< | 'model_parameters' | 'codeEnvAvailable' | 'statefulCodeSessions' + | 'statefulCodeEnvironment' + | 'codeSessionKey' | 'includeReasoningHistory' > & { subagentAgentConfigs?: SubagentTreeNode[]; @@ -945,44 +954,6 @@ function isAskUserQuestionAdminDisabled(appConfig?: AppConfig): boolean { return appConfig?.filteredTools?.includes(ASK_USER_QUESTION_TOOL_NAME) === true; } -/** - * Whether any agent reachable in the run — primary, handoff/parallel, or a - * nested subagent — resolved `statefulCodeSessions` during initialization - * (admin `stateful_code_sessions` capability AND the agent's builder opt-in - * AND a working code env). Walks `subagentAgentConfigs` like - * {@link anyAgentHasCodeEnv}; when true, `createRun` opts the run's remote - * sandbox tools into stateful runtime sessions via `toolExecution.sandbox`. - * Off by default: the capability is absent from `defaultAgentCapabilities`, - * agents opt in individually, and the SDK derives the session hint from - * `thread_id` (= conversationId), so this is never a trust boundary. - */ -export function anyAgentHasStatefulSessions(agents: Array): boolean { - const visited = new Set(); - const pending: Array = [...agents]; - - for (let index = 0; index < pending.length; index++) { - const agent = pending[index]; - if (agent == null || visited.has(agent.id)) { - continue; - } - visited.add(agent.id); - if (agent.statefulCodeSessions === true) { - return true; - } - for (const child of agent.subagentAgentConfigs ?? []) { - if (child != null && !visited.has(child.id)) { - pending.push(child); - } - } - for (const child of agent.lazySubagentConfigs ?? []) { - if (!visited.has(child.id)) { - pending.push(child); - } - } - } - return false; -} - /** * Whether any agent reachable in the run — primary, handoff/parallel, or a * nested subagent — opts into cross-turn `reasoning_content` reconstruction. @@ -1496,6 +1467,8 @@ export async function createRun({ initialSummary: isSubagent ? undefined : initialSummary, contextPruningConfig: summarization.contextPruning, maxToolResultChars: agent.maxToolResultChars, + initialSessions: buildAgentInitialToolSessions(agent, initialSessions), + codeSessionKey: agent.codeSessionKey, }; if (askGraphTools) { /** @@ -1572,9 +1545,6 @@ export async function createRun({ * and the resume route). When disabled, nothing attaches and the run is identical * to before this feature shipped. */ - // Per-agent truth resolved by initializeAgent (admin capability AND builder - // opt-in AND code env) — the run opts in when any reachable agent did. - const statefulCodeSessions = anyAgentHasStatefulSessions(agents); // Resolve the effective policy through the single seam so per-agent / per-skill // sources can layer in later without touching this call site (see // `resolveToolApprovalPolicy`). Only the endpoint layer is wired today, so this @@ -1761,15 +1731,6 @@ export async function createRun({ ...(enableToolOutputReferences && { toolOutputReferences: { enabled: true }, }), - // Best-effort stateful runtime sessions on the remote Code API. The SDK - // stamps a per-conversation session hint on execute_code/bash requests and - // hedges those tools' descriptions; the transport is otherwise unchanged. - // `engine` is omitted (defaults to `sandbox`) and the hint defaults to - // thread_id. Requires @librechat/agents with `toolExecution.sandbox`; - // older versions ignore the field. - ...(statefulCodeSessions && { - toolExecution: { sandbox: { statefulSessions: true } }, - }), // HITL opt-in: the `humanInTheLoop` switch + the PreToolUse policy hook. Spread // here (not just `compileOptions.checkpointer` above) so an `ask` decision raises // a real interrupt — without these the run would never pause. Absent when disabled. diff --git a/packages/api/src/agents/skillFiles.spec.ts b/packages/api/src/agents/skillFiles.spec.ts index d02d91fe5c..bcd0eed7db 100644 --- a/packages/api/src/agents/skillFiles.spec.ts +++ b/packages/api/src/agents/skillFiles.spec.ts @@ -14,8 +14,12 @@ jest.mock('./run', () => ({ import { Readable } from 'stream'; import { Types } from 'mongoose'; -import { primeInvokedSkills, primeSkillFiles } from './skillFiles'; -import type { PrimeInvokedSkillsDeps, PrimeSkillFilesParams } from './skillFiles'; +import { primeInvokedSkills, primeInvokedSkillsForProfiles, primeSkillFiles } from './skillFiles'; +import type { + PrimeInvokedSkillsDeps, + PrimeInvokedSkillsForProfilesDeps, + PrimeSkillFilesParams, +} from './skillFiles'; const SKILL_ID = new Types.ObjectId(); const SKILL_VERSION = 7; @@ -93,6 +97,10 @@ describe('primeInvokedSkills — execute_code capability gate', () => { listSkillFiles, getStrategyFunctions, batchUploadCodeEnvFiles, + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }, }); await primeInvokedSkills(deps); @@ -110,6 +118,8 @@ describe('primeInvokedSkills — execute_code capability gate', () => { expect(uploadArgs.kind).toBe('skill'); expect(uploadArgs.id).toBe(SKILL_ID.toString()); expect(uploadArgs.version).toBe(SKILL_VERSION); + expect(uploadArgs.codeApiBaseUrl).toBe('https://stateful-code.example.com'); + expect(uploadArgs.executionProfile).toBe('stateful'); expect(uploadArgs.files).toHaveLength(fileRecords.length + 1); expect(uploadArgs.files.map((f: { filename: string }) => f.filename)).toEqual( expect.arrayContaining([ @@ -249,6 +259,7 @@ describe('primeInvokedSkills — execute_code capability gate', () => { storage_session_id: 'session-42', file_id: 'file-1', version: SKILL_VERSION, + executionProfile: 'default', }, }, ]); @@ -298,6 +309,7 @@ describe('primeInvokedSkills — execute_code capability gate', () => { version: SKILL_VERSION, }, deps.req, + undefined, ); const codeSession = result.initialSessions?.get('execute_code'); expect(codeSession?.files).toEqual([ @@ -315,6 +327,95 @@ describe('primeInvokedSkills — execute_code capability gate', () => { }); }); +describe('primeInvokedSkillsForProfiles', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockExtract.mockReturnValue(new Set(['brand-guidelines'])); + }); + + it('uploads and seeds historical skill files separately for default and stateful profiles', async () => { + const listSkillFiles = jest.fn().mockResolvedValue([ + { + relativePath: 'references/style.md', + filename: 'style.md', + filepath: '/storage/brand-guidelines/references/style.md', + source: 's3', + bytes: 256, + }, + ]); + const getStrategyFunctions = jest.fn().mockReturnValue({ + getDownloadStream: jest.fn().mockResolvedValue(Readable.from(Buffer.from('style'))), + }); + const batchUploadCodeEnvFiles = jest.fn().mockImplementation(({ executionProfile }) => ({ + storage_session_id: `${executionProfile}-session`, + files: [ + { + fileId: `${executionProfile}-file`, + filename: 'skills/brand-guidelines/references/style.md', + }, + ], + })); + const updateSkillFileCodeEnvIds = jest.fn().mockResolvedValue({ + matchedCount: 1, + modifiedCount: 1, + }); + const { + codeEnvAvailable: _codeEnvAvailable, + codeExecutionContext: _codeExecutionContext, + ...baseDeps + } = makeDeps({ + listSkillFiles, + getStrategyFunctions, + batchUploadCodeEnvFiles, + updateSkillFileCodeEnvIds, + }); + const statefulKey = 'execute_code:stateful:v2:user:abc'; + const deps: PrimeInvokedSkillsForProfilesDeps = { + ...baseDeps, + executionProfiles: [ + { + codeExecutionContext: { + baseUrl: 'https://code.example.com/v1', + codeSessionKey: 'execute_code', + executionProfile: 'default', + statefulSessions: false, + }, + codeSessionKeys: ['execute_code'], + }, + { + codeExecutionContext: { + baseUrl: 'https://stateful.example.com/v1', + codeSessionKey: statefulKey, + executionProfile: 'stateful', + runtimeSessionHint: 'v2:user:abc', + statefulSessions: true, + }, + codeSessionKeys: [statefulKey], + }, + ], + }; + + const result = await primeInvokedSkillsForProfiles(deps); + + expect(batchUploadCodeEnvFiles).toHaveBeenCalledTimes(2); + expect( + batchUploadCodeEnvFiles.mock.calls.map(([args]) => args.executionProfile).sort(), + ).toEqual(['default', 'stateful']); + expect(result.initialSessions?.get('execute_code')?.files?.map((file) => file.id)).toEqual([ + 'default-file', + ]); + expect(result.initialSessions?.get(statefulKey)?.files?.map((file) => file.id)).toEqual([ + 'stateful-file', + ]); + expect(updateSkillFileCodeEnvIds).toHaveBeenCalledTimes(2); + expect( + updateSkillFileCodeEnvIds.mock.calls + .map(([updates]) => updates[0].codeEnvRef.executionProfile) + .sort(), + ).toEqual(['default', 'stateful']); + }); +}); + /* The tool-invoked skill loader (`handle_skill` -> `primeSkillFiles`) * is a separate code path from `primeInvokedSkills` (the NL-detected * loader). Both feed `_injected_files` on the next /exec; both must @@ -384,6 +485,7 @@ describe('primeSkillFiles — resource identity propagation', () => { storage_session_id: 'session-cached', file_id: 'file-cached', version: SKILL_VERSION, + executionProfile: 'stateful' as const, }; const batchUploadCodeEnvFiles = jest.fn(); const deps = makeSkillFilesDeps({ @@ -400,11 +502,19 @@ describe('primeSkillFiles — resource identity propagation', () => { batchUploadCodeEnvFiles, getSessionInfo: jest.fn().mockResolvedValue('2026-05-06T00:00:00Z'), checkIfActive: jest.fn().mockReturnValue(true), + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }, }); const result = await primeSkillFiles(deps); expect(batchUploadCodeEnvFiles).not.toHaveBeenCalled(); + expect(deps.getSessionInfo).toHaveBeenCalledWith(cachedRef, deps.req, { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }); expect(result?.files).toEqual([ { id: 'file-cached', @@ -416,6 +526,54 @@ describe('primeSkillFiles — resource identity propagation', () => { }, ]); }); + + it('reuploads a cached ref that belongs to the other execution profile', async () => { + const batchUploadCodeEnvFiles = jest.fn().mockResolvedValue({ + storage_session_id: 'stateful-session', + files: [ + { + fileId: 'stateful-file', + filename: 'skills/brand-guidelines/references/style.md', + }, + { fileId: 'skill-md', filename: 'skills/brand-guidelines/SKILL.md' }, + ], + }); + const getSessionInfo = jest.fn().mockResolvedValue('2026-05-06T00:00:00Z'); + const deps = makeSkillFilesDeps({ + skillFiles: [ + { + relativePath: 'references/style.md', + filename: 'style.md', + filepath: '/storage/brand-guidelines/references/style.md', + source: 's3', + bytes: 256, + codeEnvRef: { + kind: 'skill', + id: SKILL_ID.toString(), + storage_session_id: 'default-session', + file_id: 'default-file', + version: SKILL_VERSION, + executionProfile: 'default', + }, + }, + ], + batchUploadCodeEnvFiles, + getSessionInfo, + checkIfActive: jest.fn().mockReturnValue(true), + codeExecutionContext: { + baseUrl: 'https://stateful-code.example.com', + executionProfile: 'stateful', + }, + }); + + const result = await primeSkillFiles(deps); + + expect(getSessionInfo).not.toHaveBeenCalled(); + expect(batchUploadCodeEnvFiles).toHaveBeenCalledWith( + expect.objectContaining({ executionProfile: 'stateful' }), + ); + expect(result?.files[0].id).toBe('stateful-file'); + }); }); /* Codeapi's upload limiter defaults to 30 requests per user per 5 minutes, diff --git a/packages/api/src/agents/skillFiles.ts b/packages/api/src/agents/skillFiles.ts index d167ebce2b..f626f29581 100644 --- a/packages/api/src/agents/skillFiles.ts +++ b/packages/api/src/agents/skillFiles.ts @@ -2,10 +2,16 @@ import { Readable } from 'stream'; import { isAxiosError } from 'axios'; import { Constants } from '@librechat/agents'; import { logger } from '@librechat/data-schemas'; -import type { ToolSessionMap, CodeSessionContext } from '@librechat/agents'; -import type { CodeEnvRef } from 'librechat-data-provider'; +import { + getCodeEnvRefForProfile, + type CodeEnvRef, + type CodeEnvRefMap, +} from 'librechat-data-provider'; +import type { CodeEnvFile, ToolSessionMap, CodeSessionContext } from '@librechat/agents'; import type { Types } from 'mongoose'; +import type { CodeExecutionContext } from './execution'; import type { ServerRequest } from '~/types'; +import { seedCodeFilesIntoSessions, type CodeExecutionProfileRoute } from './codeFilesSession'; import { createConcurrencyLimiter, logAxiosError } from '~/utils'; import { extractInvokedSkillsFromPayload } from './run'; import { SKILL_FILE_PREFIX } from './skills'; @@ -17,6 +23,7 @@ export interface SkillFileRecord { source: string; bytes: number; codeEnvRef?: CodeEnvRef; + codeEnvRefs?: CodeEnvRefMap; } export interface PrimeSkillFilesParams { @@ -49,12 +56,23 @@ export interface PrimeSkillFilesParams { * (read-only inputs that must never surface as generated artifacts, * even if sandboxed code mutates the bytes on disk). */ read_only?: boolean; + codeApiBaseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; }) => Promise<{ storage_session_id: string; files: Array<{ fileId: string; filename: string }>; }>; /** Checks if a code env file is still active. Returns lastModified timestamp or null. */ - getSessionInfo?: (ref: CodeEnvRef, req?: ServerRequest) => Promise; + getSessionInfo?: ( + ref: CodeEnvRef, + req?: ServerRequest, + route?: { + baseUrl?: string; + executionProfile?: CodeExecutionContext['executionProfile']; + }, + ) => Promise; + /** Trusted Code API route selected for the executing agent. */ + codeExecutionContext?: Pick; /** 23-hour freshness check */ checkIfActive?: (dateString: string) => boolean; /** Persists `codeEnvRef` on skill files after upload. Implementations @@ -190,7 +208,7 @@ export async function primeSkillFiles( * resource-scoped (`:skill::v:`), so sharing the * result across requests is sound. Per-process best-effort; the awaited * codeEnvRef persist covers cross-turn and cross-node dedupe. */ - const flightKey = `${params.skill._id}:v:${params.skill.version}`; + const flightKey = `${params.codeExecutionContext?.executionProfile ?? 'default'}:${params.skill._id}:v:${params.skill.version}`; const inflight = inflightPrimes.get(flightKey); if (inflight) { return inflight; @@ -213,7 +231,9 @@ async function executePrimeSkillFiles( getSessionInfo, checkIfActive, updateSkillFileCodeEnvIds, + codeExecutionContext, } = params; + const executionProfile = codeExecutionContext?.executionProfile ?? 'default'; /* Cache-hit path: every skillFile carries a `codeEnvRef` from the * previous prime. Check freshness against codeapi for every distinct @@ -222,11 +242,13 @@ async function executePrimeSkillFiles( * skill is edited, the upsert clears the ref and forces a fresh * upload on the next prime. */ if (getSessionInfo && checkIfActive && skillFiles.length > 0) { - const allHaveRefs = skillFiles.every((sf) => sf.codeEnvRef !== undefined); + const allHaveRefs = skillFiles.every( + (sf) => getCodeEnvRefForProfile(sf, executionProfile) !== undefined, + ); if (allHaveRefs) { const refsBySession = new Map(); for (const sf of skillFiles) { - const ref = sf.codeEnvRef; + const ref = getCodeEnvRefForProfile(sf, executionProfile); if (ref && !refsBySession.has(ref.storage_session_id)) { refsBySession.set(ref.storage_session_id, ref); } @@ -235,7 +257,7 @@ async function executePrimeSkillFiles( try { const checkResults = await Promise.all( Array.from(refsBySession.values()).map(async (ref) => { - const lastModified = await getSessionInfo(ref, req); + const lastModified = await getSessionInfo(ref, req, codeExecutionContext); return !!(lastModified && checkIfActive(lastModified)); }), ); @@ -244,7 +266,7 @@ async function executePrimeSkillFiles( if (allActive) { const files: PrimeSkillFilesResult['files'] = []; for (const sf of skillFiles) { - const ref = sf.codeEnvRef; + const ref = getCodeEnvRefForProfile(sf, executionProfile); if (!ref) continue; /* Cache-hit refs already carry resource identity (kind / id / * version) — pull them through so the artifact emitted by @@ -305,6 +327,8 @@ async function executePrimeSkillFiles( * skill files surface as ghost generated artifacts the user has no * authority to download. */ read_only: true, + codeApiBaseUrl: codeExecutionContext?.baseUrl, + executionProfile: codeExecutionContext?.executionProfile, }); return { filesToUpload, result }; }, `skill "${skill.name}"`), @@ -373,6 +397,7 @@ async function executePrimeSkillFiles( storage_session_id: result.storage_session_id, file_id: f.fileId, version: skill.version, + executionProfile, }; return { skillId: skill._id, @@ -429,6 +454,7 @@ export interface PrimeInvokedSkillsDeps { getSessionInfo?: PrimeSkillFilesParams['getSessionInfo']; checkIfActive?: PrimeSkillFilesParams['checkIfActive']; updateSkillFileCodeEnvIds?: PrimeSkillFilesParams['updateSkillFileCodeEnvIds']; + codeExecutionContext?: PrimeSkillFilesParams['codeExecutionContext']; } export interface PrimeInvokedSkillsResult { @@ -438,6 +464,11 @@ export interface PrimeInvokedSkillsResult { skills?: Map; } +export interface PrimeInvokedSkillsForProfilesDeps + extends Omit { + executionProfiles: CodeExecutionProfileRoute[]; +} + /** * Extracts previously invoked skills from message history, resolves their * bodies from DB, and re-primes their files to the code env. @@ -502,8 +533,13 @@ export async function primeInvokedSkills( // ALL distinct sessions for freshness. If all are active, return cached // references with zero re-uploads. If any expired, re-upload everything. if (deps.getSessionInfo && deps.checkIfActive) { + const executionProfile = deps.codeExecutionContext?.executionProfile ?? 'default'; const allResolved = fileListResults.flatMap((r) => - r.files.map((f) => ({ skillName: r.skill.name, file: f, ref: f.codeEnvRef })), + r.files.map((f) => ({ + skillName: r.skill.name, + file: f, + ref: getCodeEnvRefForProfile(f, executionProfile), + })), ); const resolvedWithRef = allResolved.filter((x) => x.ref !== undefined); @@ -519,7 +555,11 @@ export async function primeInvokedSkills( const checkResults = await Promise.all( Array.from(refsBySession.values()).map(async (ref) => { try { - const lastModified = await deps.getSessionInfo?.(ref, deps.req); + const lastModified = await deps.getSessionInfo?.( + ref, + deps.req, + deps.codeExecutionContext, + ); return !!(lastModified && deps.checkIfActive?.(lastModified)); } catch { return false; @@ -589,6 +629,7 @@ export async function primeInvokedSkills( getSessionInfo: deps.getSessionInfo, checkIfActive: deps.checkIfActive, updateSkillFileCodeEnvIds: deps.updateSkillFileCodeEnvIds, + codeExecutionContext: deps.codeExecutionContext, }); return { skill, result }; }), @@ -637,3 +678,49 @@ export async function primeInvokedSkills( skills: skills.size > 0 ? skills : undefined, }; } + +/** Primes historical skill files once per selected Code API deployment and + * seeds only the trusted session partitions that execute on that deployment. */ +export async function primeInvokedSkillsForProfiles( + deps: PrimeInvokedSkillsForProfilesDeps, +): Promise { + if (deps.executionProfiles.length === 0) { + return primeInvokedSkills({ ...deps, codeEnvAvailable: false }); + } + + const profileResults = await Promise.all( + deps.executionProfiles.map(async (profile) => ({ + profile, + result: await primeInvokedSkills({ + ...deps, + codeEnvAvailable: true, + codeExecutionContext: profile.codeExecutionContext, + updateSkillFileCodeEnvIds: deps.updateSkillFileCodeEnvIds, + }), + })), + ); + + let initialSessions: ToolSessionMap | undefined; + const skills = new Map(); + for (const { profile, result } of profileResults) { + for (const [name, body] of result.skills ?? []) { + skills.set(name, body); + } + const skillFiles = result.initialSessions?.get(Constants.EXECUTE_CODE)?.files; + if (!skillFiles?.length) { + continue; + } + for (const sessionKey of profile.codeSessionKeys) { + initialSessions = seedCodeFilesIntoSessions( + skillFiles as CodeEnvFile[], + initialSessions, + sessionKey, + ); + } + } + + return { + initialSessions, + skills: skills.size > 0 ? skills : undefined, + }; +} diff --git a/packages/api/src/agents/statefulCodeSessions.spec.ts b/packages/api/src/agents/statefulCodeSessions.spec.ts deleted file mode 100644 index a299922cea..0000000000 --- a/packages/api/src/agents/statefulCodeSessions.spec.ts +++ /dev/null @@ -1,53 +0,0 @@ -import { anyAgentHasStatefulSessions } from './run'; - -type WalkInput = Parameters[0]; - -interface TestAgent { - id: string; - statefulCodeSessions?: boolean; - subagentAgentConfigs?: Array; - lazySubagentConfigs?: Array; -} - -function agent( - id: string, - statefulCodeSessions?: boolean, - subagentAgentConfigs?: Array, -): TestAgent { - return { id, statefulCodeSessions, subagentAgentConfigs }; -} - -function walk(agents: Array): boolean { - return anyAgentHasStatefulSessions(agents as WalkInput); -} - -describe('anyAgentHasStatefulSessions', () => { - it('is true when a top-level agent resolved the per-agent flag at initialization', () => { - expect(walk([agent('a', true)])).toBe(true); - expect(walk([agent('a', false), agent('b', true)])).toBe(true); - }); - - it('stays off (default) when no agent opted in or the flag is absent', () => { - expect(walk([])).toBe(false); - expect(walk([agent('a')])).toBe(false); - expect(walk([agent('a', false)])).toBe(false); - }); - - it('walks nested subagent configs so a stateful subagent activates the run', () => { - const grandchild = agent('c', true); - const child = agent('b', false, [grandchild]); - expect(walk([agent('a', false, [child])])).toBe(true); - }); - - it('includes inert lazy descriptors when deciding whether to prewarm', () => { - expect(walk([{ id: 'a', lazySubagentConfigs: [agent('lazy-child', true)] }])).toBe(true); - }); - - it('tolerates null entries and cycles in the subagent graph', () => { - const a = agent('a', false); - const b = agent('b', false); - a.subagentAgentConfigs = [b, null]; - b.subagentAgentConfigs = [a]; - expect(walk([a, undefined, null])).toBe(false); - }); -}); diff --git a/packages/api/src/agents/tools.ts b/packages/api/src/agents/tools.ts index c89c750677..fe43cdbb1c 100644 --- a/packages/api/src/agents/tools.ts +++ b/packages/api/src/agents/tools.ts @@ -97,8 +97,8 @@ export interface RegisterCodeExecutionToolsParams { /** * When `true`, the registered `bash_tool` description is the hedged * stateful-session variant (workspace usually persists across calls, may - * reset at any time). Paired with `toolExecution.sandbox.statefulSessions` - * in `createRun`; resolved per-agent during initialization. + * reset at any time). Transport routing is resolved independently from the + * actually executing agent at tool-load time. */ statefulSessions?: boolean; } diff --git a/packages/api/src/agents/validation.spec.ts b/packages/api/src/agents/validation.spec.ts index fe9f17bfb8..7edfef482b 100644 --- a/packages/api/src/agents/validation.spec.ts +++ b/packages/api/src/agents/validation.spec.ts @@ -99,6 +99,26 @@ describe('agentCreateSchema with subagents', () => { }); }); +describe('stateful code environments', () => { + it.each(['user', 'agent-user', 'conversation'])('accepts %s', (environment) => { + const result = agentCreateSchema.safeParse({ + provider: 'openAI', + model: 'gpt-4o-mini', + tools: [], + stateful_code_sessions: true, + stateful_code_environment: environment, + }); + expect(result.success).toBe(true); + }); + + it('rejects unknown environment scopes', () => { + const result = agentUpdateSchema.safeParse({ + stateful_code_environment: 'agent', + }); + expect(result.success).toBe(false); + }); +}); + describe('agentUpdateSchema with subagents', () => { it('accepts a partial update with only the disabled flag set', () => { const result = agentUpdateSchema.safeParse({ diff --git a/packages/api/src/agents/validation.ts b/packages/api/src/agents/validation.ts index 0b4c6cd2f6..8f05615487 100644 --- a/packages/api/src/agents/validation.ts +++ b/packages/api/src/agents/validation.ts @@ -293,6 +293,7 @@ export const agentBaseSchema: z.ZodObject< end_after_tools: z.ZodOptional; hide_sequential_outputs: z.ZodOptional; stateful_code_sessions: z.ZodOptional; + stateful_code_environment: z.ZodOptional>; artifacts: z.ZodOptional; recursion_limit: z.ZodOptional; conversation_starters: z.ZodOptional>; @@ -384,6 +385,7 @@ export const agentBaseSchema: z.ZodObject< end_after_tools: z.boolean().optional(), hide_sequential_outputs: z.boolean().optional(), stateful_code_sessions: z.boolean().optional(), + stateful_code_environment: z.enum(['user', 'agent-user', 'conversation']).optional(), artifacts: z.string().optional(), recursion_limit: z.number().optional(), conversation_starters: z.array(z.string()).optional(), @@ -478,6 +480,7 @@ export const agentCreateSchema: z.ZodObject< end_after_tools: z.ZodOptional; hide_sequential_outputs: z.ZodOptional; stateful_code_sessions: z.ZodOptional; + stateful_code_environment: z.ZodOptional>; artifacts: z.ZodOptional; recursion_limit: z.ZodOptional; conversation_starters: z.ZodOptional>; @@ -628,6 +631,7 @@ export const agentUpdateSchema: z.ZodObject< end_after_tools: z.ZodOptional; hide_sequential_outputs: z.ZodOptional; stateful_code_sessions: z.ZodOptional; + stateful_code_environment: z.ZodOptional>; artifacts: z.ZodOptional; recursion_limit: z.ZodOptional; conversation_starters: z.ZodOptional>; diff --git a/packages/api/src/skills/deployment.ts b/packages/api/src/skills/deployment.ts index 3152a53246..1cf18e0efc 100644 --- a/packages/api/src/skills/deployment.ts +++ b/packages/api/src/skills/deployment.ts @@ -3,6 +3,7 @@ import path from 'path'; import yaml from 'js-yaml'; import crypto from 'crypto'; import { Types } from 'mongoose'; +import { mergeCodeEnvRef, type CodeEnvRef, type CodeEnvRefMap } from 'librechat-data-provider'; import { logger, partitionIssues, @@ -16,7 +17,6 @@ import { normalizeSkillFrontmatterKeys, } from '@librechat/data-schemas'; import type { ValidationIssue } from '@librechat/data-schemas'; -import type { CodeEnvRef } from 'librechat-data-provider'; import { parseFrontmatter, guessMimeType } from './import'; export const DEPLOYMENT_SKILLS_DIR_ENV = 'DEPLOYMENT_SKILLS_DIR'; @@ -53,6 +53,7 @@ export type DeploymentSkillFile = { content?: string; isBinary?: boolean; codeEnvRef?: CodeEnvRef; + codeEnvRefs?: CodeEnvRefMap; createdAt: Date; updatedAt: Date; }; @@ -153,7 +154,10 @@ type ListAlwaysApplyResult = { after?: string | null; }; -type SkillFileRow = Omit & { +type SkillFileRow = Omit< + DeploymentSkillFile, + 'codeEnvRef' | 'codeEnvRefs' | 'content' | 'isBinary' +> & { storageKey?: string; storageRegion?: string; tenantId?: string; @@ -161,6 +165,7 @@ type SkillFileRow = Omit Promise> | Record; + /** Trusted Code API route selected for the executing agent. */ + codeExecutionContext?: CodeExecutionContext; } /** Result from building tool classification */ @@ -267,6 +270,7 @@ export async function buildToolClassification( programmaticToolsEnabled = false, codeExecutionEnabled = false, authHeaders, + codeExecutionContext, } = params; const isGoogle = provider === Providers.GOOGLE || provider === Providers.VERTEXAI; const additionalTools: GenericTool[] = []; @@ -374,9 +378,18 @@ export async function buildToolClassification( } try { - const ptcTool = createBashProgrammaticToolCallingTool({ authHeaders } as Parameters< - typeof createBashProgrammaticToolCallingTool - >[0] & { authHeaders?: BuildToolClassificationParams['authHeaders'] }); + const profileParams = codeExecutionContext + ? { + baseUrl: codeExecutionContext.baseUrl, + executionProfile: codeExecutionContext.executionProfile, + runtimeSessionHint: codeExecutionContext.runtimeSessionHint, + } + : {}; + const ptcTool = createBashProgrammaticToolCallingTool({ + authHeaders, + ...profileParams, + } as Parameters[0] & + typeof profileParams & { authHeaders?: BuildToolClassificationParams['authHeaders'] }); additionalTools.push(ptcTool); /** Add PTC definition for event-driven mode */ diff --git a/packages/data-provider/src/codeEnvRef.spec.ts b/packages/data-provider/src/codeEnvRef.spec.ts index 561850c502..e8e7651f59 100644 --- a/packages/data-provider/src/codeEnvRef.spec.ts +++ b/packages/data-provider/src/codeEnvRef.spec.ts @@ -1,11 +1,10 @@ -/* `CodeEnvRef` is a plain typed struct (no helpers, no resolvers). - * Behavioral coverage lives at consumer sites — `processCodeOutput` - * (write), `primeFiles` (read+reupload), `primeSkillFiles` (read+write), - * agents `ToolNode` (forward to codeapi). This file just pins the - * shape so a future refactor can't silently widen or narrow the - * fields without surfacing here. */ -import { CODE_ENV_KINDS } from './codeEnvRef'; import type { CodeEnvKind, CodeEnvRef } from './codeEnvRef'; +import { + CODE_ENV_KINDS, + getCodeEnvRefForProfile, + getCodeEnvRefs, + mergeCodeEnvRef, +} from './codeEnvRef'; describe('CodeEnvRef', () => { it('accepts the canonical shape for kind: skill', () => { @@ -52,4 +51,41 @@ describe('CodeEnvRef', () => { const kinds: CodeEnvKind[] = [...CODE_ENV_KINDS]; expect(kinds).toEqual(['skill', 'agent', 'user']); }); + + it('retains independent pointers for default and stateful deployments', () => { + const defaultRef: CodeEnvRef = { + kind: 'user', + id: 'user-1', + storage_session_id: 'default-session', + file_id: 'default-file', + executionProfile: 'default', + }; + const statefulRef: CodeEnvRef = { + ...defaultRef, + storage_session_id: 'stateful-session', + file_id: 'stateful-file', + executionProfile: 'stateful', + }; + + const refs = mergeCodeEnvRef(mergeCodeEnvRef(undefined, defaultRef), statefulRef); + + expect(getCodeEnvRefForProfile(refs, 'default')).toBe(defaultRef); + expect(getCodeEnvRefForProfile(refs, 'stateful')).toBe(statefulRef); + expect(getCodeEnvRefs(refs)).toEqual([ + ['default', defaultRef], + ['stateful', statefulRef], + ]); + expect(refs.codeEnvRef).toBe(defaultRef); + }); + + it('treats a legacy pointer without a profile as default-only', () => { + const legacy: CodeEnvRef = { + kind: 'agent', + id: 'agent-1', + storage_session_id: 'legacy-session', + file_id: 'legacy-file', + }; + expect(getCodeEnvRefForProfile({ codeEnvRef: legacy }, 'default')).toBe(legacy); + expect(getCodeEnvRefForProfile({ codeEnvRef: legacy }, 'stateful')).toBeUndefined(); + }); }); diff --git a/packages/data-provider/src/codeEnvRef.ts b/packages/data-provider/src/codeEnvRef.ts index 81187ff78f..23fc5287f2 100644 --- a/packages/data-provider/src/codeEnvRef.ts +++ b/packages/data-provider/src/codeEnvRef.ts @@ -51,9 +51,72 @@ interface CodeEnvRefBase { id: string; storage_session_id: string; file_id: string; + /** Code API deployment that owns this storage pointer. Legacy refs omit + * the field and are treated as `default`; new writes always persist it. */ + executionProfile?: CodeExecutionProfile; } +export type CodeExecutionProfile = 'default' | 'stateful'; + export type CodeEnvRef = | (CodeEnvRefBase & { kind: 'skill'; version: number }) | (CodeEnvRefBase & { kind: 'agent' }) | (CodeEnvRefBase & { kind: 'user' }); + +/** Deployment-local pointers for a resource that may be used by both Code API profiles. */ +export type CodeEnvRefMap = Partial>; + +export interface CodeEnvReferenceSet { + /** Compatibility pointer for readers that have not adopted profile-aware refs yet. */ + codeEnvRef?: CodeEnvRef; + /** Canonical deployment-local pointers, keyed by trusted execution profile. */ + codeEnvRefs?: CodeEnvRefMap; +} + +export function getCodeEnvRefForProfile( + refs: CodeEnvReferenceSet | null | undefined, + profile: CodeExecutionProfile, +): CodeEnvRef | undefined { + const profileRef = refs?.codeEnvRefs?.[profile]; + if (profileRef) { + return profileRef; + } + const legacyRef = refs?.codeEnvRef; + if (legacyRef && (legacyRef.executionProfile ?? 'default') === profile) { + return legacyRef; + } + return undefined; +} + +/** Adds one profile pointer without discarding the other profile's storage object. */ +export function mergeCodeEnvRef( + refs: CodeEnvReferenceSet | null | undefined, + ref: CodeEnvRef, +): Required> { + const codeEnvRefs: CodeEnvRefMap = { ...refs?.codeEnvRefs }; + const legacyRef = refs?.codeEnvRef; + if (legacyRef) { + codeEnvRefs[legacyRef.executionProfile ?? 'default'] ??= legacyRef; + } + const profile = ref.executionProfile ?? 'default'; + codeEnvRefs[profile] = ref; + return { + codeEnvRef: codeEnvRefs.default ?? codeEnvRefs.stateful!, + codeEnvRefs, + }; +} + +/** Enumerates every deployment-local pointer, including legacy single-pointer records. */ +export function getCodeEnvRefs( + refs: CodeEnvReferenceSet | null | undefined, +): Array<[CodeExecutionProfile, CodeEnvRef]> { + const merged: CodeEnvRefMap = { ...refs?.codeEnvRefs }; + const legacyRef = refs?.codeEnvRef; + if (legacyRef) { + merged[legacyRef.executionProfile ?? 'default'] ??= legacyRef; + } + return (['default', 'stateful'] as const).flatMap((profile) => { + const ref = merged[profile]; + return ref ? [[profile, ref] as [CodeExecutionProfile, CodeEnvRef]] : []; + }); +} diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts index 5b7e3ddd91..df6baaa41b 100644 --- a/packages/data-provider/src/schemas.ts +++ b/packages/data-provider/src/schemas.ts @@ -347,6 +347,7 @@ export const defaultAgentFormValues = { [Tools.file_search]: false, [Tools.web_search]: false, [Tools.memory]: false, + stateful_code_environment: 'user' as const, category: 'general', support_contact: { name: '', diff --git a/packages/data-provider/src/types/assistants.ts b/packages/data-provider/src/types/assistants.ts index cf875763d0..1a8cfecbdb 100644 --- a/packages/data-provider/src/types/assistants.ts +++ b/packages/data-provider/src/types/assistants.ts @@ -301,6 +301,8 @@ export type Agent = { hide_sequential_outputs?: boolean; /** Per-agent opt-in for stateful code sessions (requires the app-level capability). */ stateful_code_sessions?: boolean; + /** Stateful workspace sharing scope. Defaults to one workspace per user. */ + stateful_code_environment?: StatefulCodeEnvironment; artifacts?: ArtifactModes; recursion_limit?: number; isPublic?: boolean; @@ -353,6 +355,7 @@ export type AgentCreateParams = { | 'end_after_tools' | 'hide_sequential_outputs' | 'stateful_code_sessions' + | 'stateful_code_environment' | 'artifacts' | 'recursion_limit' | 'category' @@ -382,6 +385,7 @@ export type AgentUpdateParams = { | 'end_after_tools' | 'hide_sequential_outputs' | 'stateful_code_sessions' + | 'stateful_code_environment' | 'artifacts' | 'recursion_limit' | 'category' @@ -534,6 +538,8 @@ export enum AnnotationTypes { FILE_PATH = 'file_path', } +export type StatefulCodeEnvironment = 'user' | 'agent-user' | 'conversation'; + export enum StepStatus { IN_PROGRESS = 'in_progress', CANCELLED = 'cancelled', diff --git a/packages/data-provider/src/types/files.ts b/packages/data-provider/src/types/files.ts index f518df24be..9ab4e64192 100644 --- a/packages/data-provider/src/types/files.ts +++ b/packages/data-provider/src/types/files.ts @@ -1,4 +1,4 @@ -import type { CodeEnvRef } from '../codeEnvRef'; +import type { CodeEnvRef, CodeEnvRefMap } from '../codeEnvRef'; import { EToolResources } from './assistants'; export enum FileSources { @@ -162,6 +162,7 @@ export type TFile = { * resolve via `resolveCodeEnvRef`. */ codeEnvRef?: CodeEnvRef; + codeEnvRefs?: CodeEnvRefMap; }; createdAt?: string | Date; updatedAt?: string | Date; diff --git a/packages/data-schemas/src/methods/file.spec.ts b/packages/data-schemas/src/methods/file.spec.ts index 54eece631b..d4235b03fd 100644 --- a/packages/data-schemas/src/methods/file.spec.ts +++ b/packages/data-schemas/src/methods/file.spec.ts @@ -86,6 +86,38 @@ describe('File Methods', () => { expect(file?.file_id).toBe(fileId); expect(file?.expiresAt).toBeUndefined(); }); + + it('persists independent Code API pointers for both execution profiles', async () => { + const defaultRef = { + kind: 'user' as const, + id: 'user-1', + storage_session_id: 'default-session', + file_id: 'default-file', + executionProfile: 'default' as const, + }; + const statefulRef = { + ...defaultRef, + storage_session_id: 'stateful-session', + file_id: 'stateful-file', + executionProfile: 'stateful' as const, + }; + + const file = await fileMethods.createFile({ + file_id: uuidv4(), + user: new mongoose.Types.ObjectId(), + filename: 'dual-profile.txt', + filepath: '/uploads/dual-profile.txt', + type: 'text/plain', + bytes: 10, + metadata: { + codeEnvRef: defaultRef, + codeEnvRefs: { default: defaultRef, stateful: statefulRef }, + }, + }); + + expect(file?.metadata?.codeEnvRefs?.default?.file_id).toBe('default-file'); + expect(file?.metadata?.codeEnvRefs?.stateful?.file_id).toBe('stateful-file'); + }); }); describe('claimCodeFile', () => { diff --git a/packages/data-schemas/src/methods/file.ts b/packages/data-schemas/src/methods/file.ts index cb544862c1..b145ea65f9 100644 --- a/packages/data-schemas/src/methods/file.ts +++ b/packages/data-schemas/src/methods/file.ts @@ -248,7 +248,10 @@ export function createFileMethods(mongoose: typeof import('mongoose')): { conversationId, context: FileContext.execute_code, file_id: { $in: threadFileIds }, - 'metadata.codeEnvRef': { $exists: true }, + $or: [ + { 'metadata.codeEnvRef': { $exists: true } }, + { 'metadata.codeEnvRefs': { $exists: true } }, + ], }, ownerScope, ); @@ -285,7 +288,10 @@ export function createFileMethods(mongoose: typeof import('mongoose')): { { file_id: { $in: fileIds }, context: { $ne: FileContext.execute_code }, - 'metadata.codeEnvRef': { $exists: true }, + $or: [ + { 'metadata.codeEnvRef': { $exists: true } }, + { 'metadata.codeEnvRefs': { $exists: true } }, + ], }, ownerScope, ); diff --git a/packages/data-schemas/src/methods/skill.spec.ts b/packages/data-schemas/src/methods/skill.spec.ts index c71bafc71c..aaefd5795b 100644 --- a/packages/data-schemas/src/methods/skill.spec.ts +++ b/packages/data-schemas/src/methods/skill.spec.ts @@ -1867,7 +1867,7 @@ describe('SkillFile methods', () => { } }); - it('clears codeEnvRef when a skill file is upserted (replacement)', async () => { + it('clears all code environment refs when a skill file is upserted (replacement)', async () => { /* A re-upload of a skill file replaces the row's contents — but the * cached `codeEnvRef` refers to the OLD bytes living in codeapi. * Leaving it populated would make the next prime resolve a stale @@ -1926,6 +1926,7 @@ describe('SkillFile methods', () => { expect(after).toHaveLength(1); expect(after[0].file_id).toBe('file-2'); expect(after[0].codeEnvRef).toBeUndefined(); + expect(after[0].codeEnvRefs).toBeUndefined(); }); it('deleteSkillFile recounts and bumps version', async () => { @@ -2024,6 +2025,60 @@ describe('SkillFile methods', () => { file_id: 'file-1', version: 1, }); + expect(files[0].codeEnvRefs?.default).toMatchObject({ + kind: 'skill', + id: entityId, + storage_session_id: 'session-1', + file_id: 'file-1', + version: 1, + }); + }); + + it('retains pointers for both execution profiles', async () => { + const { skill } = await methods.createSkill(makeSkillInput()); + await methods.upsertSkillFile({ + skillId: skill._id, + relativePath: 'scripts/a.sh', + file_id: 'f1', + filename: 'a.sh', + filepath: '/a', + source: 'local', + mimeType: 'text/plain', + bytes: 1, + author: owner._id, + }); + const base = { + kind: 'skill' as const, + id: skill._id.toString(), + version: 1, + }; + + await methods.updateSkillFileCodeEnvIds([ + { + skillId: skill._id, + relativePath: 'scripts/a.sh', + codeEnvRef: { + ...base, + storage_session_id: 'default-session', + file_id: 'default-file', + executionProfile: 'default', + }, + }, + { + skillId: skill._id, + relativePath: 'scripts/a.sh', + codeEnvRef: { + ...base, + storage_session_id: 'stateful-session', + file_id: 'stateful-file', + executionProfile: 'stateful', + }, + }, + ]); + + const [file] = await methods.listSkillFiles(skill._id); + expect(file.codeEnvRefs?.default?.file_id).toBe('default-file'); + expect(file.codeEnvRefs?.stateful?.file_id).toBe('stateful-file'); }); it('reports modifiedCount=0 when no SkillFile rows match the (skillId, relativePath) filter', async () => { diff --git a/packages/data-schemas/src/methods/skill.ts b/packages/data-schemas/src/methods/skill.ts index ae6852cec8..15db4b09fb 100644 --- a/packages/data-schemas/src/methods/skill.ts +++ b/packages/data-schemas/src/methods/skill.ts @@ -1827,7 +1827,7 @@ export function createSkillMethods( author: row.author, tenantId: row.tenantId, }, - $unset: { content: '', isBinary: '', codeEnvRef: '' }, + $unset: { content: '', isBinary: '', codeEnvRef: '', codeEnvRefs: '' }, }, { new: true, upsert: true, includeResultMetadata: true }, ).lean()) as unknown as SkillFileUpsertResult; @@ -1883,12 +1883,20 @@ export function createSkillMethods( ): Promise<{ matchedCount: number; modifiedCount: number }> { if (updates.length === 0) return { matchedCount: 0, modifiedCount: 0 }; const SkillFile = mongoose.models.SkillFile as Model; - const ops = updates.map((u) => ({ - updateOne: { - filter: { skillId: u.skillId, relativePath: u.relativePath }, - update: { $set: { codeEnvRef: u.codeEnvRef } }, - }, - })); + const ops = updates.map((u) => { + const profile = u.codeEnvRef.executionProfile ?? 'default'; + return { + updateOne: { + filter: { skillId: u.skillId, relativePath: u.relativePath }, + update: { + $set: { + codeEnvRef: u.codeEnvRef, + [`codeEnvRefs.${profile}`]: u.codeEnvRef, + }, + }, + }, + }; + }); /** * The returned `{matchedCount, modifiedCount}` lets callers warn on diff --git a/packages/data-schemas/src/schema/agent.ts b/packages/data-schemas/src/schema/agent.ts index 4ea93acba7..a49c15e5c0 100644 --- a/packages/data-schemas/src/schema/agent.ts +++ b/packages/data-schemas/src/schema/agent.ts @@ -77,6 +77,10 @@ const agentSchema: Schema = new Schema( stateful_code_sessions: { type: Boolean, }, + stateful_code_environment: { + type: String, + enum: ['user', 'agent-user', 'conversation'], + }, /** @deprecated Use edges instead */ agent_ids: { type: [String], diff --git a/packages/data-schemas/src/schema/codeEnvRef.ts b/packages/data-schemas/src/schema/codeEnvRef.ts new file mode 100644 index 0000000000..a468d10093 --- /dev/null +++ b/packages/data-schemas/src/schema/codeEnvRef.ts @@ -0,0 +1,28 @@ +import { Schema } from 'mongoose'; + +export const codeEnvRefSchema: Schema = new Schema( + { + kind: { + type: String, + enum: ['skill', 'agent', 'user'], + required: true, + }, + id: { type: String, required: true }, + storage_session_id: { type: String, required: true }, + file_id: { type: String, required: true }, + version: { type: Number }, + executionProfile: { + type: String, + enum: ['default', 'stateful'], + }, + }, + { _id: false }, +); + +export const codeEnvRefMapSchema: Schema = new Schema( + { + default: { type: codeEnvRefSchema, default: undefined }, + stateful: { type: codeEnvRefSchema, default: undefined }, + }, + { _id: false }, +); diff --git a/packages/data-schemas/src/schema/file.ts b/packages/data-schemas/src/schema/file.ts index 8855c7e93f..208ad67bdf 100644 --- a/packages/data-schemas/src/schema/file.ts +++ b/packages/data-schemas/src/schema/file.ts @@ -1,6 +1,7 @@ import mongoose, { Schema } from 'mongoose'; import { FileContext, FileSources } from 'librechat-data-provider'; import type { IMongoFile } from '~/types'; +import { codeEnvRefMapSchema, codeEnvRefSchema } from './codeEnvRef'; const file: Schema = new Schema( { @@ -119,20 +120,11 @@ const file: Schema = new Schema( height: Number, metadata: { codeEnvRef: { - type: new Schema( - { - kind: { - type: String, - enum: ['skill', 'agent', 'user'], - required: true, - }, - id: { type: String, required: true }, - storage_session_id: { type: String, required: true }, - file_id: { type: String, required: true }, - version: { type: Number }, - }, - { _id: false }, - ), + type: codeEnvRefSchema, + default: undefined, + }, + codeEnvRefs: { + type: codeEnvRefMapSchema, default: undefined, }, /** Dispatch-order stamp of the last writer (or claimant, on insert): diff --git a/packages/data-schemas/src/schema/skillFile.ts b/packages/data-schemas/src/schema/skillFile.ts index 27f38847b9..e4c86c066b 100644 --- a/packages/data-schemas/src/schema/skillFile.ts +++ b/packages/data-schemas/src/schema/skillFile.ts @@ -1,5 +1,6 @@ import { Schema } from 'mongoose'; import type { ISkillFileDocument } from '~/types/skill'; +import { codeEnvRefMapSchema, codeEnvRefSchema } from './codeEnvRef'; /** Max length for a skill file's relative path (e.g. "scripts/parse.sh"). */ const SKILL_FILE_PATH_MAX_LENGTH = 500; @@ -106,20 +107,11 @@ const skillFileSchema: Schema = new Schema( type: Boolean, }, codeEnvRef: { - type: new Schema( - { - kind: { - type: String, - enum: ['skill', 'agent', 'user'], - required: true, - }, - id: { type: String, required: true }, - storage_session_id: { type: String, required: true }, - file_id: { type: String, required: true }, - version: { type: Number }, - }, - { _id: false }, - ), + type: codeEnvRefSchema, + default: undefined, + }, + codeEnvRefs: { + type: codeEnvRefMapSchema, default: undefined, }, }, diff --git a/packages/data-schemas/src/types/agent.ts b/packages/data-schemas/src/types/agent.ts index 4bb47766f1..092c1a70c7 100644 --- a/packages/data-schemas/src/types/agent.ts +++ b/packages/data-schemas/src/types/agent.ts @@ -37,6 +37,7 @@ export interface IAgent extends Omit { hide_sequential_outputs?: boolean; end_after_tools?: boolean; stateful_code_sessions?: boolean; + stateful_code_environment?: 'user' | 'agent-user' | 'conversation'; /** @deprecated Use edges instead */ agent_ids?: string[]; edges?: GraphEdge[]; diff --git a/packages/data-schemas/src/types/file.ts b/packages/data-schemas/src/types/file.ts index 9ccdf16034..edca65e87c 100644 --- a/packages/data-schemas/src/types/file.ts +++ b/packages/data-schemas/src/types/file.ts @@ -1,5 +1,5 @@ import { Document, Types } from 'mongoose'; -import type { CodeEnvRef } from 'librechat-data-provider'; +import type { CodeEnvRef, CodeEnvRefMap } from 'librechat-data-provider'; export interface IMongoFile extends Omit { user: Types.ObjectId; @@ -70,6 +70,7 @@ export interface IMongoFile extends Omit { * derive the sessionKey explicitly. */ codeEnvRef?: CodeEnvRef; + codeEnvRefs?: CodeEnvRefMap; }; expiresAt?: Date; expiredAt?: Date | null; diff --git a/packages/data-schemas/src/types/skill.ts b/packages/data-schemas/src/types/skill.ts index e10c5bfda7..d3cf701826 100644 --- a/packages/data-schemas/src/types/skill.ts +++ b/packages/data-schemas/src/types/skill.ts @@ -1,4 +1,4 @@ -import type { CodeEnvRef } from 'librechat-data-provider'; +import type { CodeEnvRef, CodeEnvRefMap } from 'librechat-data-provider'; import type { Document, Types } from 'mongoose'; /** @@ -137,6 +137,7 @@ export interface ISkillFile { * when the skill file is re-uploaded to storage. */ codeEnvRef?: CodeEnvRef; + codeEnvRefs?: CodeEnvRefMap; createdAt?: Date; updatedAt?: Date; }