From ac984cc0b29ee5b64c98ac408733ed48ba413c73 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 13 Jul 2026 13:56:45 -0400 Subject: [PATCH 1/6] =?UTF-8?q?=E2=99=A8=EF=B8=8F=20feat:=20Prewarm=20Stat?= =?UTF-8?q?eful=20Code=20Sandboxes=20with=20Cold-Boot=20UX=20Feedback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/controllers/agents/callbacks.js | 33 ++++ api/server/controllers/agents/client.js | 11 ++ .../components/Chat/Messages/Content/Part.tsx | 5 + .../Chat/Messages/Content/Parts/BashCall.tsx | 23 ++- .../Messages/Content/Parts/ExecuteCode.tsx | 9 +- client/src/hooks/SSE/useStepHandler.ts | 35 +++- client/src/locales/en/translation.json | 1 + client/src/store/index.ts | 1 + client/src/store/sandbox.ts | 13 ++ packages/api/src/agents/handlers.ts | 8 + packages/api/src/agents/index.ts | 1 + packages/api/src/agents/prewarm.spec.ts | 159 ++++++++++++++++++ packages/api/src/agents/prewarm.ts | 159 ++++++++++++++++++ packages/data-provider/src/types/runs.ts | 8 + 14 files changed, 457 insertions(+), 9 deletions(-) create mode 100644 client/src/store/sandbox.ts create mode 100644 packages/api/src/agents/prewarm.spec.ts create mode 100644 packages/api/src/agents/prewarm.ts diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 196d6b8978a..50b34ddebf9 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -3,6 +3,7 @@ const { logger } = require('@librechat/data-schemas'); const { Tools, StepTypes, + StepEvents, FileContext, ErrorTypes, UsageEvents, @@ -21,6 +22,7 @@ const { createToolExecuteHandler, HOST_FILE_AUTHORING_ARTIFACT_KEY, isCodeSessionToolName, + shouldSignalSandboxStart, } = require('@librechat/api'); const { processFileCitations } = require('~/server/services/Files/Citations'); const { processCodeOutput, runPreviewFinalize } = require('~/server/services/Files/Code/process'); @@ -228,6 +230,36 @@ async function emitEvent(res, streamId, eventData) { } } +/** + * Emits `on_sandbox_starting` for each code-execution tool call in the run + * step when the conversation's stateful sandbox is still cold-booting, so the + * UI can explain the first call's boot latency instead of showing a generic + * running state. No-op on stateless deployments ({@link shouldSignalSandboxStart} + * only tracks conversations that fired a prewarm). + * @param {ServerResponse} res - The server response object + * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode + * @param {StreamEventData} data - The `on_run_step` event data + * @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata + * @returns {Promise} + */ +async function maybeEmitSandboxStarting(res, streamId, data, metadata) { + const conversationId = metadata?.thread_id; + if (!conversationId || !shouldSignalSandboxStart(conversationId)) { + return; + } + const toolCalls = data?.stepDetails?.tool_calls ?? []; + for (const toolCall of toolCalls) { + const name = toolCall?.name ?? toolCall?.function?.name; + if (!toolCall?.id || name == null || !isCodeSessionToolName(name)) { + continue; + } + await emitEvent(res, streamId, { + event: StepEvents.ON_SANDBOX_STARTING, + data: { tool_call_id: toolCall.id, runId: metadata?.run_id }, + }); + } +} + /** * Maps a {@link SubagentUpdateEvent} phase to the corresponding * {@link GraphEvents} name that the SDK's `createContentAggregator` @@ -363,6 +395,7 @@ function getDefaultHandlers({ aggregateContent({ event, data }); if (data?.stepDetails.type === StepTypes.TOOL_CALLS) { await emitEvent(res, streamId, { event, data }); + await maybeEmitSandboxStarting(res, streamId, data, metadata); } else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { await emitEvent(res, streamId, { event, data }); } else if (!metadata?.hide_sequential_outputs) { diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 4a2768f809b..d6026adfe39 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -68,6 +68,7 @@ const { appendYouTubeVideoParts, resolveYouTubeInjectionConfig, decrementPendingRequest, + maybePrewarmCodeSandbox, } = require('@librechat/api'); const { Callback, @@ -1386,6 +1387,16 @@ class AgentClient extends BaseClient { abortController = new AbortController(); } + /** Fire-and-forget: boot the per-conversation stateful sandbox in + * parallel with generation so the first execute_code/bash call lands + * on a warm VM. No-op unless a reachable agent resolved + * `statefulCodeSessions`. */ + maybePrewarmCodeSandbox({ + req: this.options.req, + conversationId: this.conversationId, + agents: [this.options.agent, ...(this.agentConfigs?.values() ?? [])], + }); + /** @type {AppConfig['endpoints']['agents']} */ const agentsEConfig = appConfig.endpoints?.[EModelEndpoint.agents]; diff --git a/client/src/components/Chat/Messages/Content/Part.tsx b/client/src/components/Chat/Messages/Content/Part.tsx index 9173d8c896e..38555d15961 100644 --- a/client/src/components/Chat/Messages/Content/Part.tsx +++ b/client/src/components/Chat/Messages/Content/Part.tsx @@ -151,6 +151,8 @@ const Part = memo(function Part({ const isToolCall = 'args' in toolCall && (!toolCall.type || toolCall.type === ToolCallTypes.TOOL_CALL); if (isToolCall) { + const toolCallId = + 'id' in toolCall && typeof toolCall.id === 'string' ? toolCall.id : undefined; const card = (() => { if (isBashProgrammaticToolCall(toolCall.name, toolCall.args)) { return ( @@ -163,6 +165,7 @@ const Part = memo(function Part({ commandField="code" hideAttachments={hideAttachments} onExpand={onToolExpand} + toolCallId={toolCallId} /> ); } else if ( @@ -179,6 +182,7 @@ const Part = memo(function Part({ args={toolCall.args} hideAttachments={hideAttachments} onExpand={onToolExpand} + toolCallId={toolCallId} /> ); } else if ( @@ -279,6 +283,7 @@ const Part = memo(function Part({ attachments={attachments} hideAttachments={hideAttachments} onExpand={onToolExpand} + toolCallId={toolCallId} /> ); } else if (toolCall.name === Tools.web_search) { diff --git a/client/src/components/Chat/Messages/Content/Parts/BashCall.tsx b/client/src/components/Chat/Messages/Content/Parts/BashCall.tsx index 7221a0f7f01..67356fcf77c 100644 --- a/client/src/components/Chat/Messages/Content/Parts/BashCall.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/BashCall.tsx @@ -1,14 +1,16 @@ import { useMemo, useRef, useState, useCallback, useEffect } from 'react'; import copy from 'copy-to-clipboard'; +import { useRecoilValue } from 'recoil'; import type { TAttachment } from 'librechat-data-provider'; import ProgressText from '~/components/Chat/Messages/Content/ProgressText'; +import parseJsonField, { areToolCallArgsComplete } from './parseJsonField'; import CopyButton from '~/components/Messages/Content/CopyButton'; import LangIcon from '~/components/Messages/Content/LangIcon'; +import { sandboxStartingByToolCallId } from '~/store'; import useToolCallState from './useToolCallState'; import useLazyHighlight from './useLazyHighlight'; import { ERROR_PATTERNS } from './ExecuteCode'; import { AttachmentGroup } from './Attachment'; -import parseJsonField, { areToolCallArgsComplete } from './parseJsonField'; import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; @@ -21,6 +23,7 @@ export default function BashCall({ commandField = 'command', hideAttachments = false, onExpand, + toolCallId, }: { initialProgress: number; isSubmitting: boolean; @@ -30,10 +33,12 @@ export default function BashCall({ commandField?: string; hideAttachments?: boolean; onExpand?: () => void; + toolCallId?: string; }) { const localize = useLocalize(); const command = useMemo(() => parseJsonField(args, commandField), [args, commandField]); const isWritingCommand = !command || !areToolCallArgsComplete(args); + const sandboxStarting = useRecoilValue(sandboxStartingByToolCallId(toolCallId ?? '')); const { showCode, toggleCode, expandStyle, expandRef, progress, cancelled, hasError, hasOutput } = useToolCallState(initialProgress, isSubmitting, output, !!command, onExpand); @@ -52,17 +57,23 @@ export default function BashCall({ timerRef.current = setTimeout(() => setIsCopied(false), 3000); }, [command]); + const inProgressText = (() => { + if (isWritingCommand) { + return localize('com_ui_writing_command'); + } + if (sandboxStarting) { + return localize('com_ui_sandbox_starting'); + } + return localize('com_ui_running_command'); + })(); + return ( <>
void; + toolCallId?: string; }) { const localize = useLocalize(); const { lang = 'py', code } = useParseArgs(args) ?? ({} as ParsedArgs); + const sandboxStarting = useRecoilValue(sandboxStartingByToolCallId(toolCallId ?? '')); const { showCode, toggleCode, expandStyle, expandRef, progress, cancelled, hasError, hasOutput } = useToolCallState(initialProgress, isSubmitting, output, !!code, onExpand); @@ -82,7 +87,9 @@ export default function ExecuteCode({ ()); + + const setSandboxStarting = useRecoilCallback( + ({ set }) => + (toolCallId: string): void => { + knownSandboxAtomKeys.current.add(toolCallId); + set(sandboxStartingByToolCallId(toolCallId), true); + }, + [], + ); + + const clearSandboxStarting = useRecoilCallback( + ({ reset }) => + (toolCallId?: string | null): void => { + if (!toolCallId || !knownSandboxAtomKeys.current.has(toolCallId)) { + return; + } + knownSandboxAtomKeys.current.delete(toolCallId); + reset(sandboxStartingByToolCallId(toolCallId)); + }, + [], + ); + /** * Calculate content index for a run step. * For edited content scenarios, offset by initialContent length. @@ -934,6 +960,7 @@ export default function useStepHandler({ const { result } = stepEvent.data; const { id: stepId } = result; + clearSandboxStarting(result.tool_call?.id); const runStep = stepMap.current.get(stepId); let responseMessageId = runStep?.runId ?? ''; @@ -977,6 +1004,8 @@ export default function useStepHandler({ }), ); } + } else if (stepEvent.event === StepEvents.ON_SANDBOX_STARTING) { + setSandboxStarting(stepEvent.data.tool_call_id); } else if (stepEvent.event === StepEvents.ON_SUBAGENT_UPDATE) { applySubagentUpdate(stepEvent.data); } else if (stepEvent.event === StepEvents.ON_SUMMARIZE_START) { @@ -1082,6 +1111,8 @@ export default function useStepHandler({ calculateContentIndex, getCurrentMessages, applySubagentUpdate, + setSandboxStarting, + clearSandboxStarting, onSkillAuthoringComplete, ], ); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index ca7d59afb53..f81c899b0fe 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1639,6 +1639,7 @@ "com_ui_running": "Running...", "com_ui_running_command": "Running command", "com_ui_running_n_agents": "Running {{0}} agents", + "com_ui_sandbox_starting": "Starting sandbox environment", "com_ui_save": "Save", "com_ui_save_badge_changes": "Save badge changes?", "com_ui_save_changes": "Save Changes", diff --git a/client/src/store/index.ts b/client/src/store/index.ts index 7d7d9cfe614..84dcfb45f8f 100644 --- a/client/src/store/index.ts +++ b/client/src/store/index.ts @@ -16,6 +16,7 @@ export * from './agents'; export * from './mcp'; export * from './favorites'; export * from './subagents'; +export * from './sandbox'; export * from './usage'; export default { diff --git a/client/src/store/sandbox.ts b/client/src/store/sandbox.ts new file mode 100644 index 00000000000..54232d90b33 --- /dev/null +++ b/client/src/store/sandbox.ts @@ -0,0 +1,13 @@ +import { atomFamily } from 'recoil'; + +/** + * True while the backend reported the stateful code sandbox is cold-booting + * for this tool call (`on_sandbox_starting` SSE event). Keyed by + * `tool_call_id`; `ExecuteCode`/`BashCall` swap their in-progress label to a + * "starting sandbox" message while set. Cleared when the tool call's run + * step completes. + */ +export const sandboxStartingByToolCallId = atomFamily({ + key: 'sandboxStartingByToolCallId', + default: false, +}); diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index a60581591a8..88956158337 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -40,6 +40,7 @@ import { buildSkillPrimeMessage, SKILL_FILE_PREFIX } from './skills'; import { parseFrontmatter } from '../skills/import'; import { cleanCodeToolOutput } from './cleanup'; import { primeSkillFiles } from './skillFiles'; +import { markSandboxReady } from './prewarm'; export interface ToolEndCallbackData { output: { @@ -3718,6 +3719,13 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand } as Record, ); + /* 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 !== '') { + markSandboxReady(tc.runtimeSessionHint); + } + // Code-execution tools emit per-call boilerplate // ("Note: ..." paragraphs and `| ` per-file // suffixes) that wastes tokens when re-injected into diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 8865ef8762d..9aa84e03b81 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -15,6 +15,7 @@ export * from './memory'; export * from './orphans'; export * from './migration'; export * from './parameters'; +export * from './prewarm'; export * from './openai'; export * from './transactions'; export * from './usage'; diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts new file mode 100644 index 00000000000..c398c921507 --- /dev/null +++ b/packages/api/src/agents/prewarm.spec.ts @@ -0,0 +1,159 @@ +import { + markSandboxReady, + maybePrewarmCodeSandbox, + resetSandboxStateForTests, + shouldSignalSandboxStart, +} from './prewarm'; + +type PrewarmParams = Parameters[0]; + +const req = {} as PrewarmParams['req']; +const statefulAgent = { id: 'agent_stateful', statefulCodeSessions: true }; +const plainAgent = { id: 'agent_plain', statefulCodeSessions: false }; + +function agents(...list: Array>): PrewarmParams['agents'] { + return list as unknown as PrewarmParams['agents']; +} + +function flushAsync(): Promise { + return new Promise((resolve) => setImmediate(resolve)); +} + +describe('maybePrewarmCodeSandbox', () => { + let fetchMock: jest.SpyInstance; + + beforeEach(() => { + resetSandboxStateForTests(); + process.env.LIBRECHAT_CODE_BASEURL = 'http://code.test/v1'; + delete process.env.CODE_SANDBOX_PREWARM; + delete process.env.CODE_SANDBOX_COLD_AFTER_MS; + fetchMock = jest + .spyOn(globalThis, 'fetch') + .mockResolvedValue({ ok: true, status: 200 } as Response); + }); + + afterEach(() => { + fetchMock.mockRestore(); + jest.useRealTimers(); + }); + + it('does nothing when no reachable agent has stateful sessions', async () => { + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(plainAgent) }); + await flushAsync(); + expect(fetchMock).not.toHaveBeenCalled(); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + }); + + it('does nothing without a conversationId', async () => { + maybePrewarmCodeSandbox({ req, conversationId: null, agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('respects the CODE_SANDBOX_PREWARM=false kill switch', async () => { + process.env.CODE_SANDBOX_PREWARM = 'false'; + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('fires one exec with the conversation as runtime_session_hint 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(JSON.parse(init.body as string)).toEqual({ + lang: 'bash', + code: 'true', + runtime_session_hint: 'convo-1', + }); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + }); + + it('walks subagent configs for the stateful gate', async () => { + const parent = { id: 'agent_parent', subagentAgentConfigs: [statefulAgent] }; + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(parent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it('does not refire within the warm window', async () => { + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it('refires once the warm window has expired', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate'] }); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + + jest.setSystemTime(new Date('2026-07-13T01:00:00Z')); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it('keeps signaling when the prewarm request fails, without throwing', async () => { + fetchMock.mockRejectedValue(new Error('boom')); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + }); + + it('treats a non-2xx prewarm response as a failure', async () => { + fetchMock.mockResolvedValue({ ok: false, status: 503 } as Response); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + }); +}); + +describe('shouldSignalSandboxStart / markSandboxReady', () => { + beforeEach(() => { + resetSandboxStateForTests(); + delete process.env.CODE_SANDBOX_COLD_AFTER_MS; + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('never signals for untracked conversations (stateless deployments)', () => { + expect(shouldSignalSandboxStart('never-seen')).toBe(false); + expect(shouldSignalSandboxStart(null)).toBe(false); + expect(shouldSignalSandboxStart(undefined)).toBe(false); + }); + + it('stops signaling after a real tool call marks the sandbox ready', () => { + jest.useFakeTimers(); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + const fetchMock = jest + .spyOn(globalThis, 'fetch') + .mockImplementation(() => new Promise(() => undefined)); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + + markSandboxReady('convo-1'); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + fetchMock.mockRestore(); + }); + + it('honors CODE_SANDBOX_COLD_AFTER_MS overrides', () => { + jest.useFakeTimers(); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + process.env.CODE_SANDBOX_COLD_AFTER_MS = '1000'; + markSandboxReady('convo-1'); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + jest.setSystemTime(new Date('2026-07-13T00:00:02Z')); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + }); +}); diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts new file mode 100644 index 00000000000..ce1555fee0d --- /dev/null +++ b/packages/api/src/agents/prewarm.ts @@ -0,0 +1,159 @@ +import { logger } from '@librechat/data-schemas'; +import { getCodeBaseURL } from '@librechat/agents'; +import type { ServerRequest } from '~/types'; +import { getCodeApiAuthHeaders } from '~/auth/codeapi'; +import { anyAgentHasStatefulSessions } from './run'; + +type PrewarmAgents = Parameters[0]; + +interface SandboxState { + firedAt: number; + readyAt: number | null; +} + +/** + * Per-conversation sandbox lifecycle observations, keyed by conversationId + * (= runtime_session_hint). Only conversations whose run resolved + * `statefulCodeSessions` ever get an entry, so presence doubles as the + * "stateful deployment" gate for {@link shouldSignalSandboxStart}. + */ +const sandboxStateByConversation = new Map(); + +const MAX_TRACKED_CONVERSATIONS = 5000; +const PREWARM_INFLIGHT_COOLDOWN_MS = 120_000; +const PREWARM_REQUEST_TIMEOUT_MS = 120_000; + +/** + * How long a sandbox is assumed to survive without a touch before a fresh + * boot is required. Mirrors the Code API's idle + suspend windows + * (`LAMBDA_MICROVM_IDLE_SECONDS` + `LAMBDA_MICROVM_SUSPEND_SECONDS`, + * 300s + 1800s by default): within the window the VM is warm or resumes + * in ~1s, past it the next exec pays a full relaunch + checkpoint restore. + */ +function coldAfterMs(): number { + const parsed = Number(process.env.CODE_SANDBOX_COLD_AFTER_MS); + return Number.isFinite(parsed) && parsed > 0 ? parsed : 2_100_000; +} + +function prewarmDisabled(): boolean { + return process.env.CODE_SANDBOX_PREWARM === 'false'; +} + +function pruneOldestEntries(): void { + if (sandboxStateByConversation.size <= MAX_TRACKED_CONVERSATIONS) { + return; + } + let oldestKey: string | null = null; + let oldestFiredAt = Infinity; + for (const [key, state] of sandboxStateByConversation) { + if (state.firedAt < oldestFiredAt) { + oldestFiredAt = state.firedAt; + oldestKey = key; + } + } + if (oldestKey != null) { + sandboxStateByConversation.delete(oldestKey); + } +} + +/** + * Record that the conversation's sandbox answered a request (prewarm or a + * real execute_code/bash call). Refreshes the warm window used by both + * {@link maybePrewarmCodeSandbox} and {@link shouldSignalSandboxStart}. + */ +export function markSandboxReady(conversationId: string): void { + if (!conversationId) { + return; + } + const now = Date.now(); + const state = sandboxStateByConversation.get(conversationId); + if (state) { + state.readyAt = now; + state.firedAt = now; + return; + } + sandboxStateByConversation.set(conversationId, { firedAt: now, readyAt: now }); + pruneOldestEntries(); +} + +/** + * Whether the UI should be told the sandbox is cold-booting for this + * conversation's code tool call. True only for conversations tracked as + * stateful (an entry exists) whose sandbox has not confirmed ready within + * the warm window — stateless deployments never get an entry and never + * signal, preserving existing behavior. + */ +export function shouldSignalSandboxStart(conversationId?: string | null): boolean { + if (!conversationId) { + return false; + } + const state = sandboxStateByConversation.get(conversationId); + if (!state) { + return false; + } + return state.readyAt == null || Date.now() - state.readyAt > coldAfterMs(); +} + +async function sendPrewarmRequest(req: ServerRequest, conversationId: string): Promise { + const authHeaders = await getCodeApiAuthHeaders(req); + const response = await fetch(`${getCodeBaseURL()}/exec`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'User-Agent': 'LibreChat/1.0', + ...authHeaders, + }, + body: JSON.stringify({ lang: 'bash', code: 'true', runtime_session_hint: conversationId }), + signal: AbortSignal.timeout(PREWARM_REQUEST_TIMEOUT_MS), + }); + if (!response.ok) { + throw new Error(`prewarm exec returned ${response.status}`); + } + markSandboxReady(conversationId); + logger.debug(`[prewarmCodeSandbox] Sandbox warm for conversation ${conversationId}`); +} + +/** + * Fire-and-forget boot of the per-conversation stateful code sandbox 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`, the sandbox is + * outside its warm window, and no prewarm is already in flight. Failures + * are logged at debug level and never affect the chat request; the first + * real tool call simply pays the cold boot as before. + */ +export function maybePrewarmCodeSandbox(params: { + req: ServerRequest; + conversationId?: string | null; + agents: PrewarmAgents; +}): void { + const { req, conversationId, agents } = params; + if (prewarmDisabled() || !conversationId || !anyAgentHasStatefulSessions(agents)) { + return; + } + const now = Date.now(); + const state = sandboxStateByConversation.get(conversationId); + const withinWarmWindow = state?.readyAt != null && now - state.readyAt <= coldAfterMs(); + const prewarmInFlight = state != null && now - state.firedAt <= PREWARM_INFLIGHT_COOLDOWN_MS; + if (withinWarmWindow || prewarmInFlight) { + return; + } + if (state) { + state.firedAt = now; + } else { + sandboxStateByConversation.set(conversationId, { firedAt: now, readyAt: null }); + pruneOldestEntries(); + } + void sendPrewarmRequest(req, conversationId).catch((error) => { + logger.debug( + `[prewarmCodeSandbox] Prewarm failed for conversation ${conversationId}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + }); +} + +/** Test-only: clear tracked sandbox state between specs. */ +export function resetSandboxStateForTests(): void { + sandboxStateByConversation.clear(); +} diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index da386ac5f4b..e7fbed9adb4 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -40,8 +40,16 @@ export enum StepEvents { ON_SUMMARIZE_DELTA = 'on_summarize_delta', ON_SUMMARIZE_COMPLETE = 'on_summarize_complete', ON_SUBAGENT_UPDATE = 'on_subagent_update', + ON_SANDBOX_STARTING = 'on_sandbox_starting', } +/** Payload for {@link StepEvents.ON_SANDBOX_STARTING} — the stateful code + * sandbox is cold-booting for the given code tool call. */ +export type SandboxStartingEvent = { + tool_call_id: string; + runId?: string; +}; + /** Token-tracking event names streamed to the client (separate from StepEvents dispatch). */ export enum UsageEvents { ON_CONTEXT_USAGE = 'on_context_usage', From 26f5665d703a85d99c2a087296959af210bdbb71 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 13 Jul 2026 17:57:10 -0400 Subject: [PATCH 2/6] =?UTF-8?q?=F0=9F=A7=B9=20fix:=20Drain=20Prewarm=20Res?= =?UTF-8?q?ponse=20+=20Reset=20Sandbox=20Atoms=20on=20Stream=20Cleanup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- client/src/hooks/SSE/useStepHandler.ts | 18 +++++++++++++++++- packages/api/src/agents/prewarm.spec.ts | 8 ++++++-- packages/api/src/agents/prewarm.ts | 4 ++++ 3 files changed, 27 insertions(+), 3 deletions(-) diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index 01a0cecfffc..411d7e26d75 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -310,6 +310,17 @@ export default function useStepHandler({ [], ); + const resetSandboxAtoms = useRecoilCallback( + ({ reset }) => + (): void => { + for (const toolCallId of knownSandboxAtomKeys.current) { + reset(sandboxStartingByToolCallId(toolCallId)); + } + knownSandboxAtomKeys.current.clear(); + }, + [], + ); + /** * Calculate content index for a run step. * For edited content scenarios, offset by initialContent length. @@ -1129,6 +1140,11 @@ export default function useStepHandler({ subagentRunToToolCallId.current.clear(); claimedSubagentToolCallIds.current.clear(); pendingSubagentBuffer.current.clear(); + /** Unlike subagent atoms below, sandbox-starting flags are transient + * status with no audit value — reset them at this boundary so an + * interrupted cold boot can't leak a stale "starting" label onto a + * later tool call that reuses the same id (e.g. `call_0`). */ + resetSandboxAtoms(); /** Intentionally NOT calling `resetSubagentAtoms()` here — users need * to be able to reopen the SubagentCall dialog after completion to * audit what the child did. `resetSubagentAtoms` is returned below @@ -1137,7 +1153,7 @@ export default function useStepHandler({ * persisted `subagent_content` takes over for historical messages * once the conversation is saved, and we prevent unbounded * atomFamily growth across multi-conversation sessions. */ - }, []); + }, [resetSandboxAtoms]); /** * Sync a message into the step handler's messageMap. diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts index c398c921507..411e410ab31 100644 --- a/packages/api/src/agents/prewarm.spec.ts +++ b/packages/api/src/agents/prewarm.spec.ts @@ -19,6 +19,10 @@ function flushAsync(): Promise { return new Promise((resolve) => setImmediate(resolve)); } +function mockResponse(init: { ok: boolean; status: number }): Response { + return { ...init, arrayBuffer: async () => new ArrayBuffer(0) } as Response; +} + describe('maybePrewarmCodeSandbox', () => { let fetchMock: jest.SpyInstance; @@ -29,7 +33,7 @@ describe('maybePrewarmCodeSandbox', () => { delete process.env.CODE_SANDBOX_COLD_AFTER_MS; fetchMock = jest .spyOn(globalThis, 'fetch') - .mockResolvedValue({ ok: true, status: 200 } as Response); + .mockResolvedValue(mockResponse({ ok: true, status: 200 })); }); afterEach(() => { @@ -110,7 +114,7 @@ describe('maybePrewarmCodeSandbox', () => { }); it('treats a non-2xx prewarm response as a failure', async () => { - fetchMock.mockResolvedValue({ ok: false, status: 503 } as Response); + fetchMock.mockResolvedValue(mockResponse({ ok: false, status: 503 })); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); expect(shouldSignalSandboxStart('convo-1')).toBe(true); diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts index ce1555fee0d..96592394ede 100644 --- a/packages/api/src/agents/prewarm.ts +++ b/packages/api/src/agents/prewarm.ts @@ -106,6 +106,10 @@ async function sendPrewarmRequest(req: ServerRequest, conversationId: string): P body: JSON.stringify({ lang: 'bash', code: 'true', runtime_session_hint: conversationId }), signal: AbortSignal.timeout(PREWARM_REQUEST_TIMEOUT_MS), }); + /* Drain before acting on the status: fetch resolves at headers, and the + * sandbox is only warm once the exec's body has fully arrived — this also + * releases the socket instead of leaving the body for undici to reap. */ + await response.arrayBuffer().catch(() => undefined); if (!response.ok) { throw new Error(`prewarm exec returned ${response.status}`); } From 7300e20636b61dd3eb2411203eef6989d2705946 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 13 Jul 2026 18:14:59 -0400 Subject: [PATCH 3/6] =?UTF-8?q?=F0=9F=9A=BF=20fix:=20Propagate=20Prewarm?= =?UTF-8?q?=20Drain=20Failures=20+=20Warm=20Marker=20for=20Host=20File=20T?= =?UTF-8?q?ools?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/api/src/agents/handlers.ts | 11 +++++++++++ packages/api/src/agents/prewarm.spec.ts | 13 +++++++++++++ packages/api/src/agents/prewarm.ts | 11 +++++++---- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index 88956158337..ae8f05af3dd 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -3555,6 +3555,17 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand ); } + /* Host-handled read_file/create_file/edit_file reach the + * sandbox too (they return before the generic invoke path's + * marker below), so refresh the warm window here as well. */ + if ( + handlerResult.status === 'success' && + tc.runtimeSessionHint != null && + tc.runtimeSessionHint !== '' + ) { + markSandboxReady(tc.runtimeSessionHint); + } + return handlerResult; } diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts index 411e410ab31..0b7136de813 100644 --- a/packages/api/src/agents/prewarm.spec.ts +++ b/packages/api/src/agents/prewarm.spec.ts @@ -119,6 +119,19 @@ describe('maybePrewarmCodeSandbox', () => { await flushAsync(); expect(shouldSignalSandboxStart('convo-1')).toBe(true); }); + + it('does not mark the sandbox ready when the 2xx body fails to drain', async () => { + fetchMock.mockResolvedValue({ + ok: true, + status: 200, + arrayBuffer: async () => { + throw new Error('body aborted'); + }, + } as unknown as Response); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + }); }); describe('shouldSignalSandboxStart / markSandboxReady', () => { diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts index 96592394ede..e0bac0835c0 100644 --- a/packages/api/src/agents/prewarm.ts +++ b/packages/api/src/agents/prewarm.ts @@ -106,13 +106,16 @@ async function sendPrewarmRequest(req: ServerRequest, conversationId: string): P body: JSON.stringify({ lang: 'bash', code: 'true', runtime_session_hint: conversationId }), signal: AbortSignal.timeout(PREWARM_REQUEST_TIMEOUT_MS), }); - /* Drain before acting on the status: fetch resolves at headers, and the - * sandbox is only warm once the exec's body has fully arrived — this also - * releases the socket instead of leaving the body for undici to reap. */ - await response.arrayBuffer().catch(() => undefined); if (!response.ok) { + await response.arrayBuffer().catch(() => undefined); throw new Error(`prewarm exec returned ${response.status}`); } + /* fetch resolves at headers, but the sandbox is only warm once the exec's + * body has fully arrived — a failed drain means the exec did not complete, + * so it must propagate as a prewarm failure instead of marking ready. + * Draining also releases the socket instead of leaving the body for + * undici to reap. */ + await response.arrayBuffer(); markSandboxReady(conversationId); logger.debug(`[prewarmCodeSandbox] Sandbox warm for conversation ${conversationId}`); } From 4d201cb5cbfeb2cabc5eec42546770b6dc7590c9 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 13 Jul 2026 18:31:56 -0400 Subject: [PATCH 4/6] =?UTF-8?q?=F0=9F=8C=A1=EF=B8=8F=20fix:=20Decouple=20P?= =?UTF-8?q?rewarm=20In-Flight=20State=20from=20Warm=20Refreshes=20+=20Prec?= =?UTF-8?q?ise=20Ready=20Gates?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- packages/api/src/agents/handlers.ts | 10 +++++++--- packages/api/src/agents/prewarm.spec.ts | 26 +++++++++++++++++++++++++ packages/api/src/agents/prewarm.ts | 26 +++++++++++++++++-------- 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index ae8f05af3dd..fe7668b3fde 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -3555,10 +3555,14 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand ); } - /* Host-handled read_file/create_file/edit_file reach the - * sandbox too (they return before the generic invoke path's - * marker below), so refresh the warm window here as well. */ + /* Sandbox-routed create_file/edit_file 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. */ if ( + isSandboxFileAuthoringCall && handlerResult.status === 'success' && tc.runtimeSessionHint != null && tc.runtimeSessionHint !== '' diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts index 0b7136de813..e982520ecc8 100644 --- a/packages/api/src/agents/prewarm.spec.ts +++ b/packages/api/src/agents/prewarm.spec.ts @@ -120,6 +120,21 @@ describe('maybePrewarmCodeSandbox', () => { expect(shouldSignalSandboxStart('convo-1')).toBe(true); }); + it('prewarms again after a short cold-after window even within the fire cooldown', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate'] }); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + process.env.CODE_SANDBOX_COLD_AFTER_MS = '30000'; + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + + jest.setSystemTime(new Date('2026-07-13T00:00:45Z')); + expect(shouldSignalSandboxStart('convo-1')).toBe(true); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + it('does not mark the sandbox ready when the 2xx body fails to drain', async () => { fetchMock.mockResolvedValue({ ok: true, @@ -173,4 +188,15 @@ describe('shouldSignalSandboxStart / markSandboxReady', () => { jest.setSystemTime(new Date('2026-07-13T00:00:02Z')); expect(shouldSignalSandboxStart('convo-1')).toBe(true); }); + + it('never signals when the kill switch is on, even for tracked conversations', () => { + jest.useFakeTimers(); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + process.env.CODE_SANDBOX_PREWARM = 'false'; + process.env.CODE_SANDBOX_COLD_AFTER_MS = '1000'; + markSandboxReady('convo-1'); + jest.setSystemTime(new Date('2026-07-13T00:00:05Z')); + expect(shouldSignalSandboxStart('convo-1')).toBe(false); + delete process.env.CODE_SANDBOX_PREWARM; + }); }); diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts index e0bac0835c0..bf63d5dc8cc 100644 --- a/packages/api/src/agents/prewarm.ts +++ b/packages/api/src/agents/prewarm.ts @@ -7,7 +7,9 @@ import { anyAgentHasStatefulSessions } from './run'; type PrewarmAgents = Parameters[0]; interface SandboxState { + /** When the last prewarm request was fired; 0 when only real execs touched the entry. */ firedAt: number; + /** When the sandbox last confirmed a completed request (prewarm or real exec). */ readyAt: number | null; } @@ -44,10 +46,11 @@ function pruneOldestEntries(): void { return; } let oldestKey: string | null = null; - let oldestFiredAt = Infinity; + let oldestTouchedAt = Infinity; for (const [key, state] of sandboxStateByConversation) { - if (state.firedAt < oldestFiredAt) { - oldestFiredAt = state.firedAt; + const touchedAt = Math.max(state.firedAt, state.readyAt ?? 0); + if (touchedAt < oldestTouchedAt) { + oldestTouchedAt = touchedAt; oldestKey = key; } } @@ -69,10 +72,9 @@ export function markSandboxReady(conversationId: string): void { const state = sandboxStateByConversation.get(conversationId); if (state) { state.readyAt = now; - state.firedAt = now; return; } - sandboxStateByConversation.set(conversationId, { firedAt: now, readyAt: now }); + sandboxStateByConversation.set(conversationId, { firedAt: 0, readyAt: now }); pruneOldestEntries(); } @@ -81,10 +83,11 @@ export function markSandboxReady(conversationId: string): void { * conversation's code tool call. True only for conversations tracked as * stateful (an entry exists) whose sandbox has not confirmed ready within * the warm window — stateless deployments never get an entry and never - * signal, preserving existing behavior. + * signal, preserving existing behavior. The `CODE_SANDBOX_PREWARM=false` + * kill switch silences this too, reverting the full feature to baseline. */ export function shouldSignalSandboxStart(conversationId?: string | null): boolean { - if (!conversationId) { + if (!conversationId || prewarmDisabled()) { return false; } const state = sandboxStateByConversation.get(conversationId); @@ -141,7 +144,14 @@ export function maybePrewarmCodeSandbox(params: { const now = Date.now(); const state = sandboxStateByConversation.get(conversationId); const withinWarmWindow = state?.readyAt != null && now - state.readyAt <= coldAfterMs(); - const prewarmInFlight = state != null && now - state.firedAt <= PREWARM_INFLIGHT_COOLDOWN_MS; + /* In-flight means a fired prewarm that no completion (readyAt) has caught + * up with — real execs refreshing readyAt must not extend this, or a + * cold-after window shorter than the cooldown could suppress needed + * prewarms between the two thresholds. */ + const prewarmInFlight = + state != null && + now - state.firedAt <= PREWARM_INFLIGHT_COOLDOWN_MS && + (state.readyAt == null || state.readyAt < state.firedAt); if (withinWarmWindow || prewarmInFlight) { return; } From 88aa53170512b452e08ca8f729558451017b8a71 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 14 Jul 2026 08:32:38 -0400 Subject: [PATCH 5/6] =?UTF-8?q?=E2=98=81=EF=B8=8F=20refactor:=20Redis-Back?= =?UTF-8?q?ed=20Sandbox=20Prewarm=20State=20via=20standardCache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/controllers/agents/callbacks.js | 2 +- packages/api/src/agents/handlers.ts | 4 +- packages/api/src/agents/prewarm.spec.ts | 119 ++++++++-------- packages/api/src/agents/prewarm.ts | 149 ++++++++++----------- packages/data-provider/src/config.ts | 4 + 5 files changed, 141 insertions(+), 137 deletions(-) diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 50b34ddebf9..27b7a4f2991 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -244,7 +244,7 @@ async function emitEvent(res, streamId, eventData) { */ async function maybeEmitSandboxStarting(res, streamId, data, metadata) { const conversationId = metadata?.thread_id; - if (!conversationId || !shouldSignalSandboxStart(conversationId)) { + if (!conversationId || !(await shouldSignalSandboxStart(conversationId))) { return; } const toolCalls = data?.stepDetails?.tool_calls ?? []; diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index fe7668b3fde..bebabe2644c 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -3567,7 +3567,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand tc.runtimeSessionHint != null && tc.runtimeSessionHint !== '' ) { - markSandboxReady(tc.runtimeSessionHint); + void markSandboxReady(tc.runtimeSessionHint); } return handlerResult; @@ -3738,7 +3738,7 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand * this refreshes the prewarm module's warm window without * inspecting tool names. */ if (tc.runtimeSessionHint != null && tc.runtimeSessionHint !== '') { - markSandboxReady(tc.runtimeSessionHint); + void markSandboxReady(tc.runtimeSessionHint); } // Code-execution tools emit per-call boilerplate diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts index e982520ecc8..e0118431359 100644 --- a/packages/api/src/agents/prewarm.spec.ts +++ b/packages/api/src/agents/prewarm.spec.ts @@ -26,8 +26,8 @@ function mockResponse(init: { ok: boolean; status: number }): Response { describe('maybePrewarmCodeSandbox', () => { let fetchMock: jest.SpyInstance; - beforeEach(() => { - resetSandboxStateForTests(); + beforeEach(async () => { + await resetSandboxStateForTests(); process.env.LIBRECHAT_CODE_BASEURL = 'http://code.test/v1'; delete process.env.CODE_SANDBOX_PREWARM; delete process.env.CODE_SANDBOX_COLD_AFTER_MS; @@ -45,7 +45,7 @@ describe('maybePrewarmCodeSandbox', () => { maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(plainAgent) }); await flushAsync(); expect(fetchMock).not.toHaveBeenCalled(); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); }); it('does nothing without a conversationId', async () => { @@ -73,7 +73,7 @@ describe('maybePrewarmCodeSandbox', () => { code: 'true', runtime_session_hint: 'convo-1', }); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); }); it('walks subagent configs for the stateful gate', async () => { @@ -83,7 +83,7 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).toHaveBeenCalledTimes(1); }); - it('does not refire within the warm window', async () => { + it('does not refire while the warm marker is fresh', async () => { maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); @@ -91,33 +91,26 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).toHaveBeenCalledTimes(1); }); - it('refires once the warm window has expired', async () => { - jest.useFakeTimers({ doNotFake: ['setImmediate'] }); - jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + it('does not refire while a prewarm is in flight', async () => { + fetchMock.mockImplementation(() => new Promise(() => undefined)); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); - expect(fetchMock).toHaveBeenCalledTimes(1); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); - - jest.setSystemTime(new Date('2026-07-13T01:00:00Z')); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); - expect(fetchMock).toHaveBeenCalledTimes(2); + expect(fetchMock).toHaveBeenCalledTimes(1); }); - it('keeps signaling when the prewarm request fails, without throwing', async () => { - fetchMock.mockRejectedValue(new Error('boom')); + it('refires once the warm marker has expired', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate'] }); + jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); - }); + expect(fetchMock).toHaveBeenCalledTimes(1); - it('treats a non-2xx prewarm response as a failure', async () => { - fetchMock.mockResolvedValue(mockResponse({ ok: false, status: 503 })); + jest.setSystemTime(new Date('2026-07-13T01:00:00Z')); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); + expect(fetchMock).toHaveBeenCalledTimes(2); }); it('prewarms again after a short cold-after window even within the fire cooldown', async () => { @@ -129,12 +122,42 @@ describe('maybePrewarmCodeSandbox', () => { expect(fetchMock).toHaveBeenCalledTimes(1); jest.setSystemTime(new Date('2026-07-13T00:00:45Z')); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); expect(fetchMock).toHaveBeenCalledTimes(2); }); + it('signals while a prewarm is in flight and stops after it completes', async () => { + let resolveFetch: ((value: Response) => void) | undefined; + fetchMock.mockImplementation( + () => + new Promise((resolve) => { + resolveFetch = resolve; + }), + ); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + + resolveFetch?.(mockResponse({ ok: true, status: 200 })); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); + }); + + it('keeps signaling when the prewarm request fails, without throwing', async () => { + fetchMock.mockRejectedValue(new Error('boom')); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + }); + + it('treats a non-2xx prewarm response as a failure', async () => { + fetchMock.mockResolvedValue(mockResponse({ ok: false, status: 503 })); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + }); + it('does not mark the sandbox ready when the 2xx body fails to drain', async () => { fetchMock.mockResolvedValue({ ok: true, @@ -145,13 +168,14 @@ describe('maybePrewarmCodeSandbox', () => { } as unknown as Response); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); await flushAsync(); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); }); }); describe('shouldSignalSandboxStart / markSandboxReady', () => { - beforeEach(() => { - resetSandboxStateForTests(); + beforeEach(async () => { + await resetSandboxStateForTests(); + delete process.env.CODE_SANDBOX_PREWARM; delete process.env.CODE_SANDBOX_COLD_AFTER_MS; }); @@ -159,44 +183,33 @@ describe('shouldSignalSandboxStart / markSandboxReady', () => { jest.useRealTimers(); }); - it('never signals for untracked conversations (stateless deployments)', () => { - expect(shouldSignalSandboxStart('never-seen')).toBe(false); - expect(shouldSignalSandboxStart(null)).toBe(false); - expect(shouldSignalSandboxStart(undefined)).toBe(false); + it('never signals for untracked conversations (stateless deployments)', async () => { + await expect(shouldSignalSandboxStart('never-seen')).resolves.toBe(false); + await expect(shouldSignalSandboxStart(null)).resolves.toBe(false); + await expect(shouldSignalSandboxStart(undefined)).resolves.toBe(false); }); - it('stops signaling after a real tool call marks the sandbox ready', () => { - jest.useFakeTimers(); - jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + it('stops signaling after a real tool call marks the sandbox ready', async () => { const fetchMock = jest .spyOn(globalThis, 'fetch') .mockImplementation(() => new Promise(() => undefined)); maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); - markSandboxReady('convo-1'); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); + await markSandboxReady('convo-1'); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); fetchMock.mockRestore(); }); - it('honors CODE_SANDBOX_COLD_AFTER_MS overrides', () => { - jest.useFakeTimers(); - jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); - process.env.CODE_SANDBOX_COLD_AFTER_MS = '1000'; - markSandboxReady('convo-1'); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); - jest.setSystemTime(new Date('2026-07-13T00:00:02Z')); - expect(shouldSignalSandboxStart('convo-1')).toBe(true); - }); - - it('never signals when the kill switch is on, even for tracked conversations', () => { - jest.useFakeTimers(); - jest.setSystemTime(new Date('2026-07-13T00:00:00Z')); + it('never signals when the kill switch is on, even with an in-flight prewarm', async () => { + const fetchMock = jest + .spyOn(globalThis, 'fetch') + .mockImplementation(() => new Promise(() => undefined)); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); process.env.CODE_SANDBOX_PREWARM = 'false'; - process.env.CODE_SANDBOX_COLD_AFTER_MS = '1000'; - markSandboxReady('convo-1'); - jest.setSystemTime(new Date('2026-07-13T00:00:05Z')); - expect(shouldSignalSandboxStart('convo-1')).toBe(false); - delete process.env.CODE_SANDBOX_PREWARM; + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); + fetchMock.mockRestore(); }); }); diff --git a/packages/api/src/agents/prewarm.ts b/packages/api/src/agents/prewarm.ts index bf63d5dc8cc..0e202cd54e1 100644 --- a/packages/api/src/agents/prewarm.ts +++ b/packages/api/src/agents/prewarm.ts @@ -1,27 +1,14 @@ import { logger } from '@librechat/data-schemas'; import { getCodeBaseURL } from '@librechat/agents'; +import { CacheKeys } from 'librechat-data-provider'; +import type { Keyv } from 'keyv'; import type { ServerRequest } from '~/types'; import { getCodeApiAuthHeaders } from '~/auth/codeapi'; +import { standardCache } from '~/cache/cacheFactory'; import { anyAgentHasStatefulSessions } from './run'; type PrewarmAgents = Parameters[0]; -interface SandboxState { - /** When the last prewarm request was fired; 0 when only real execs touched the entry. */ - firedAt: number; - /** When the sandbox last confirmed a completed request (prewarm or real exec). */ - readyAt: number | null; -} - -/** - * Per-conversation sandbox lifecycle observations, keyed by conversationId - * (= runtime_session_hint). Only conversations whose run resolved - * `statefulCodeSessions` ever get an entry, so presence doubles as the - * "stateful deployment" gate for {@link shouldSignalSandboxStart}. - */ -const sandboxStateByConversation = new Map(); - -const MAX_TRACKED_CONVERSATIONS = 5000; const PREWARM_INFLIGHT_COOLDOWN_MS = 120_000; const PREWARM_REQUEST_TIMEOUT_MS = 120_000; @@ -41,60 +28,66 @@ function prewarmDisabled(): boolean { return process.env.CODE_SANDBOX_PREWARM === 'false'; } -function pruneOldestEntries(): void { - if (sandboxStateByConversation.size <= MAX_TRACKED_CONVERSATIONS) { - return; - } - let oldestKey: string | null = null; - let oldestTouchedAt = Infinity; - for (const [key, state] of sandboxStateByConversation) { - const touchedAt = Math.max(state.firedAt, state.readyAt ?? 0); - if (touchedAt < oldestTouchedAt) { - oldestTouchedAt = touchedAt; - oldestKey = key; - } - } - if (oldestKey != null) { - sandboxStateByConversation.delete(oldestKey); +/** + * 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): + * - `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) + * within the warm window. + */ +let cacheInstance: Keyv | undefined; + +function sandboxCache(): Keyv { + if (!cacheInstance) { + cacheInstance = standardCache(CacheKeys.SANDBOX_PREWARM); } + return cacheInstance; +} + +function readyKey(conversationId: string): string { + return `ready:${conversationId}`; +} + +function inflightKey(conversationId: string): string { + return `inflight:${conversationId}`; } /** * Record that the conversation's sandbox answered a request (prewarm or a - * real execute_code/bash call). Refreshes the warm window used by both - * {@link maybePrewarmCodeSandbox} and {@link shouldSignalSandboxStart}. + * real execute_code/bash call), refreshing the warm window and releasing + * any in-flight prewarm marker. Callers on hot paths should not await this; + * `void markSandboxReady(...)` is the expected usage. */ -export function markSandboxReady(conversationId: string): void { +export async function markSandboxReady(conversationId: string): Promise { if (!conversationId) { return; } - const now = Date.now(); - const state = sandboxStateByConversation.get(conversationId); - if (state) { - state.readyAt = now; - return; - } - sandboxStateByConversation.set(conversationId, { firedAt: 0, readyAt: now }); - pruneOldestEntries(); + const cache = sandboxCache(); + await Promise.all([ + cache.set(readyKey(conversationId), true, coldAfterMs()), + cache.delete(inflightKey(conversationId)), + ]); } /** * Whether the UI should be told the sandbox is cold-booting for this - * conversation's code tool call. True only for conversations tracked as - * stateful (an entry exists) whose sandbox has not confirmed ready within - * the warm window — stateless deployments never get an entry and never - * signal, preserving existing behavior. The `CODE_SANDBOX_PREWARM=false` - * kill switch silences this too, reverting the full feature to baseline. + * conversation's code tool call: a prewarm is in flight and no completion + * (prewarm or real exec) has landed. Deployments that never prewarm — + * stateless setups or the `CODE_SANDBOX_PREWARM=false` kill switch — never + * have an in-flight marker and never signal, preserving existing behavior. */ -export function shouldSignalSandboxStart(conversationId?: string | null): boolean { +export async function shouldSignalSandboxStart(conversationId?: string | null): Promise { if (!conversationId || prewarmDisabled()) { return false; } - const state = sandboxStateByConversation.get(conversationId); - if (!state) { - return false; - } - return state.readyAt == null || Date.now() - state.readyAt > coldAfterMs(); + const cache = sandboxCache(); + const [ready, inflight] = await Promise.all([ + cache.get(readyKey(conversationId)), + cache.get(inflightKey(conversationId)), + ]); + return inflight != null && ready == null; } async function sendPrewarmRequest(req: ServerRequest, conversationId: string): Promise { @@ -119,7 +112,7 @@ 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(); - markSandboxReady(conversationId); + await markSandboxReady(conversationId); logger.debug(`[prewarmCodeSandbox] Sandbox warm for conversation ${conversationId}`); } @@ -127,10 +120,13 @@ async function sendPrewarmRequest(req: ServerRequest, conversationId: string): P * Fire-and-forget boot of the per-conversation stateful code sandbox 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`, the sandbox is - * outside its warm window, and no prewarm is already in flight. Failures - * are logged at debug level and never affect the chat request; the first - * real tool call simply pays the cold boot as before. + * unless a reachable agent resolved `statefulCodeSessions` and neither a + * warm marker nor an in-flight prewarm exists. The existence check and + * marker write are not atomic, so concurrent turns (or replicas) can rarely + * double-fire — harmless, since the prewarm exec is a trivial idempotent + * `true` and the Code API serializes per-session work behind its own lock. + * Failures are logged at debug level and never affect the chat request; the + * in-flight marker's TTL then acts as the retry backoff. */ export function maybePrewarmCodeSandbox(params: { req: ServerRequest; @@ -141,27 +137,18 @@ export function maybePrewarmCodeSandbox(params: { if (prewarmDisabled() || !conversationId || !anyAgentHasStatefulSessions(agents)) { return; } - const now = Date.now(); - const state = sandboxStateByConversation.get(conversationId); - const withinWarmWindow = state?.readyAt != null && now - state.readyAt <= coldAfterMs(); - /* In-flight means a fired prewarm that no completion (readyAt) has caught - * up with — real execs refreshing readyAt must not extend this, or a - * cold-after window shorter than the cooldown could suppress needed - * prewarms between the two thresholds. */ - const prewarmInFlight = - state != null && - now - state.firedAt <= PREWARM_INFLIGHT_COOLDOWN_MS && - (state.readyAt == null || state.readyAt < state.firedAt); - if (withinWarmWindow || prewarmInFlight) { - return; - } - if (state) { - state.firedAt = now; - } else { - sandboxStateByConversation.set(conversationId, { firedAt: now, readyAt: null }); - pruneOldestEntries(); - } - void sendPrewarmRequest(req, conversationId).catch((error) => { + 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; + } + await cache.set(inflightKey(conversationId), true, PREWARM_INFLIGHT_COOLDOWN_MS); + await sendPrewarmRequest(req, conversationId); + })().catch((error) => { logger.debug( `[prewarmCodeSandbox] Prewarm failed for conversation ${conversationId}: ${ error instanceof Error ? error.message : String(error) @@ -171,6 +158,6 @@ export function maybePrewarmCodeSandbox(params: { } /** Test-only: clear tracked sandbox state between specs. */ -export function resetSandboxStateForTests(): void { - sandboxStateByConversation.clear(); +export async function resetSandboxStateForTests(): Promise { + await sandboxCache().clear(); } diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 203ac3b023d..35e35d8ed80 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -2300,6 +2300,10 @@ export enum CacheKeys { * Key for cached group memberships used to resolve ACL user principals. */ USER_PRINCIPALS = 'USER_PRINCIPALS', + /** + * Key for per-conversation stateful code sandbox prewarm/warm state. + */ + SANDBOX_PREWARM = 'SANDBOX_PREWARM', /** * Key for the title generation cache. */ From 22d052b940eb00a65e4d78a4d42e9ab2d51711cd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 14 Jul 2026 10:04:50 -0400 Subject: [PATCH 6/6] =?UTF-8?q?=F0=9F=A7=AA=20chore:=20Hermetic=20Prewarm?= =?UTF-8?q?=20Spec=20+=20Accurate=20Signal=20JSDoc=20(Copilot=20review)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/server/controllers/agents/callbacks.js | 5 +++-- packages/api/src/agents/prewarm.spec.ts | 17 +++++++++++++---- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 27b7a4f2991..e006d96417f 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -234,8 +234,9 @@ async function emitEvent(res, streamId, eventData) { * Emits `on_sandbox_starting` for each code-execution tool call in the run * step when the conversation's stateful sandbox is still cold-booting, so the * UI can explain the first call's boot latency instead of showing a generic - * running state. No-op on stateless deployments ({@link shouldSignalSandboxStart} - * only tracks conversations that fired a prewarm). + * running state. Only signals while a fired prewarm remains unresolved + * ({@link shouldSignalSandboxStart}); stateless deployments never fire one + * and completed boots clear the marker, so both stay on the generic label. * @param {ServerResponse} res - The server response object * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode * @param {StreamEventData} data - The `on_run_step` event data diff --git a/packages/api/src/agents/prewarm.spec.ts b/packages/api/src/agents/prewarm.spec.ts index e0118431359..c80212500ef 100644 --- a/packages/api/src/agents/prewarm.spec.ts +++ b/packages/api/src/agents/prewarm.spec.ts @@ -7,12 +7,18 @@ import { type PrewarmParams = Parameters[0]; +interface TestAgent { + id: string; + statefulCodeSessions?: boolean; + subagentAgentConfigs?: TestAgent[]; +} + const req = {} as PrewarmParams['req']; -const statefulAgent = { id: 'agent_stateful', statefulCodeSessions: true }; -const plainAgent = { id: 'agent_plain', statefulCodeSessions: false }; +const statefulAgent: TestAgent = { id: 'agent_stateful', statefulCodeSessions: true }; +const plainAgent: TestAgent = { id: 'agent_plain', statefulCodeSessions: false }; -function agents(...list: Array>): PrewarmParams['agents'] { - return list as unknown as PrewarmParams['agents']; +function agents(...list: TestAgent[]): PrewarmParams['agents'] { + return list as PrewarmParams['agents']; } function flushAsync(): Promise { @@ -31,6 +37,8 @@ describe('maybePrewarmCodeSandbox', () => { process.env.LIBRECHAT_CODE_BASEURL = 'http://code.test/v1'; delete process.env.CODE_SANDBOX_PREWARM; delete process.env.CODE_SANDBOX_COLD_AFTER_MS; + delete process.env.CODEAPI_JWT_ENABLED; + delete process.env.CODEAPI_AUTH_PROVIDER; fetchMock = jest .spyOn(globalThis, 'fetch') .mockResolvedValue(mockResponse({ ok: true, status: 200 })); @@ -39,6 +47,7 @@ describe('maybePrewarmCodeSandbox', () => { afterEach(() => { fetchMock.mockRestore(); jest.useRealTimers(); + delete process.env.LIBRECHAT_CODE_BASEURL; }); it('does nothing when no reachable agent has stateful sessions', async () => {