diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 196d6b8978..e006d96417 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,37 @@ 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. 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 + * @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata + * @returns {Promise} + */ +async function maybeEmitSandboxStarting(res, streamId, data, metadata) { + const conversationId = metadata?.thread_id; + if (!conversationId || !(await 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 +396,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 81364f5559..eb1a11fa26 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -73,6 +73,7 @@ const { appendYouTubeVideoParts, resolveYouTubeInjectionConfig, decrementPendingRequest, + maybePrewarmCodeSandbox, } = require('@librechat/api'); const { Callback, @@ -1515,6 +1516,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 ee0942fe53..18a8b472df 100644 --- a/client/src/components/Chat/Messages/Content/Part.tsx +++ b/client/src/components/Chat/Messages/Content/Part.tsx @@ -163,6 +163,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 ( @@ -175,6 +177,7 @@ const Part = memo(function Part({ commandField="code" hideAttachments={hideAttachments} onExpand={onToolExpand} + toolCallId={toolCallId} /> ); } else if ( @@ -191,6 +194,7 @@ const Part = memo(function Part({ args={toolCall.args} hideAttachments={hideAttachments} onExpand={onToolExpand} + toolCallId={toolCallId} /> ); } else if ( @@ -291,6 +295,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 7221a0f7f0..67356fcf77 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)); + }, + [], + ); + + 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. @@ -934,6 +971,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 +1015,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 +1122,8 @@ export default function useStepHandler({ calculateContentIndex, getCurrentMessages, applySubagentUpdate, + setSandboxStarting, + clearSandboxStarting, onSkillAuthoringComplete, ], ); @@ -1098,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 @@ -1106,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/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 3a34eb2fde..8e626a1ea0 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1653,6 +1653,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 7d7d9cfe61..84dcfb45f8 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 0000000000..54232d90b3 --- /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 a60581591a..bebabe2644 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: { @@ -3554,6 +3555,21 @@ export function createToolExecuteHandler(options: ToolExecuteOptions): EventHand ); } + /* 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 !== '' + ) { + void markSandboxReady(tc.runtimeSessionHint); + } + return handlerResult; } @@ -3718,6 +3734,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 !== '') { + void 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 6b35e014b3..cdd28d2218 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 0000000000..c80212500e --- /dev/null +++ b/packages/api/src/agents/prewarm.spec.ts @@ -0,0 +1,224 @@ +import { + markSandboxReady, + maybePrewarmCodeSandbox, + resetSandboxStateForTests, + shouldSignalSandboxStart, +} from './prewarm'; + +type PrewarmParams = Parameters[0]; + +interface TestAgent { + id: string; + statefulCodeSessions?: boolean; + subagentAgentConfigs?: TestAgent[]; +} + +const req = {} as PrewarmParams['req']; +const statefulAgent: TestAgent = { id: 'agent_stateful', statefulCodeSessions: true }; +const plainAgent: TestAgent = { id: 'agent_plain', statefulCodeSessions: false }; + +function agents(...list: TestAgent[]): PrewarmParams['agents'] { + return list as PrewarmParams['agents']; +} + +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; + + 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; + delete process.env.CODEAPI_JWT_ENABLED; + delete process.env.CODEAPI_AUTH_PROVIDER; + fetchMock = jest + .spyOn(globalThis, 'fetch') + .mockResolvedValue(mockResponse({ ok: true, status: 200 })); + }); + + afterEach(() => { + fetchMock.mockRestore(); + jest.useRealTimers(); + delete process.env.LIBRECHAT_CODE_BASEURL; + }); + + 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(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.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', + }); + await expect(shouldSignalSandboxStart('convo-1')).resolves.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 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) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + 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(); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + 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(fetchMock).toHaveBeenCalledTimes(1); + + jest.setSystemTime(new Date('2026-07-13T01:00:00Z')); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + 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')); + 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, + status: 200, + arrayBuffer: async () => { + throw new Error('body aborted'); + }, + } as unknown as Response); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + }); +}); + +describe('shouldSignalSandboxStart / markSandboxReady', () => { + beforeEach(async () => { + await resetSandboxStateForTests(); + delete process.env.CODE_SANDBOX_PREWARM; + delete process.env.CODE_SANDBOX_COLD_AFTER_MS; + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + 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', async () => { + const fetchMock = jest + .spyOn(globalThis, 'fetch') + .mockImplementation(() => new Promise(() => undefined)); + maybePrewarmCodeSandbox({ req, conversationId: 'convo-1', agents: agents(statefulAgent) }); + await flushAsync(); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(true); + + await markSandboxReady('convo-1'); + await expect(shouldSignalSandboxStart('convo-1')).resolves.toBe(false); + fetchMock.mockRestore(); + }); + + 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'; + 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 new file mode 100644 index 0000000000..0e202cd54e --- /dev/null +++ b/packages/api/src/agents/prewarm.ts @@ -0,0 +1,163 @@ +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]; + +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'; +} + +/** + * 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), 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 async function markSandboxReady(conversationId: string): Promise { + if (!conversationId) { + return; + } + 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: 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 async function shouldSignalSandboxStart(conversationId?: string | null): Promise { + if (!conversationId || prewarmDisabled()) { + return false; + } + 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 { + 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) { + 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(); + await 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` 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; + conversationId?: string | null; + agents: PrewarmAgents; +}): void { + const { req, conversationId, agents } = params; + if (prewarmDisabled() || !conversationId || !anyAgentHasStatefulSessions(agents)) { + 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; + } + 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) + }`, + ); + }); +} + +/** Test-only: clear tracked sandbox state between specs. */ +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 f710a3692f..96cebc9cb4 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -2303,6 +2303,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. */ diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index b01017e32f..fdf1ca5eef 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -42,8 +42,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',