From ebbcd9c674df711e3f3779bddf7f1449310d8f78 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 22 Mar 2026 12:54:26 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B9=20chore:=20Optimize=20provisioning?= =?UTF-8?q?=20=E2=80=94=20single=20credential=20load,=20deferred=20DB=20wr?= =?UTF-8?q?ites?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - loadCodeApiKey: load CODE_API_KEY once per request, pass to both checkSessionsAlive and provisionToCodeEnv (was N+1 lookups) - provisionToCodeEnv/provisionToVectorDB now return fileUpdate objects instead of writing to DB immediately - primeResources batches all DB updates via Promise.allSettled after provisioning completes - Remove updateFile import from provision.js (no longer writes directly) --- .../services/Endpoints/agents/initialize.js | 5 ++ api/server/services/Files/provision.js | 85 ++++++++++--------- packages/api/src/agents/initialize.ts | 8 ++ packages/api/src/agents/resources.ts | 59 +++++++++++-- 4 files changed, 109 insertions(+), 48 deletions(-) diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 7a8b13874e..56bf9f5dac 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -24,6 +24,7 @@ const { const { loadAgentTools, loadToolsForExecution } = require('~/server/services/ToolService'); const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions'); const { + loadCodeApiKey, provisionToCodeEnv, provisionToVectorDB, checkSessionsAlive, @@ -224,6 +225,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { provisionToCodeEnv, provisionToVectorDB, checkSessionsAlive, + loadCodeApiKey, + updateFile: db.updateFile, }, ); @@ -308,6 +311,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { provisionToCodeEnv, provisionToVectorDB, checkSessionsAlive, + loadCodeApiKey, + updateFile: db.updateFile, }, ); diff --git a/api/server/services/Files/provision.js b/api/server/services/Files/provision.js index 176275be5b..1db7afd341 100644 --- a/api/server/services/Files/provision.js +++ b/api/server/services/Files/provision.js @@ -12,10 +12,21 @@ const { logger } = require('@librechat/data-schemas'); const { FileSources } = require('librechat-data-provider'); const { loadAuthValues } = require('~/server/services/Tools/credentials'); const { getStrategyFunctions } = require('./strategies'); -const { updateFile } = require('~/models'); const axios = createAxiosInstance(); +/** + * Loads the CODE_API_KEY for a user. Call once per request and pass the result + * to provisionToCodeEnv / checkSessionsAlive to avoid redundant lookups. + * + * @param {string} userId + * @returns {Promise} The CODE_API_KEY + */ +async function loadCodeApiKey(userId) { + const result = await loadAuthValues({ userId, authFields: [EnvVar.CODE_API_KEY] }); + return result[EnvVar.CODE_API_KEY]; +} + /** * Provisions a file to the code execution environment. * Gets a read stream from our storage and uploads to the code env. @@ -24,9 +35,10 @@ const axios = createAxiosInstance(); * @param {object} params.req - Express request object (needs req.user.id) * @param {import('librechat-data-provider').TFile} params.file - The file record from DB * @param {string} [params.entity_id] - Optional entity ID (agent_id) - * @returns {Promise} The fileIdentifier from the code env + * @param {string} [params.apiKey] - Pre-loaded CODE_API_KEY (avoids redundant loadAuthValues) + * @returns {Promise<{ fileIdentifier: string, fileUpdate: object }>} Result with deferred DB update */ -async function provisionToCodeEnv({ req, file, entity_id = '' }) { +async function provisionToCodeEnv({ req, file, entity_id = '', apiKey }) { const { getDownloadStream } = getStrategyFunctions(file.source); if (!getDownloadStream) { throw new Error( @@ -34,33 +46,26 @@ async function provisionToCodeEnv({ req, file, entity_id = '' }) { ); } + const resolvedApiKey = apiKey ?? (await loadCodeApiKey(req.user.id)); const { handleFileUpload: uploadCodeEnvFile } = getStrategyFunctions(FileSources.execute_code); - const result = await loadAuthValues({ userId: req.user.id, authFields: [EnvVar.CODE_API_KEY] }); const stream = await getDownloadStream(req, file.filepath); const fileIdentifier = await uploadCodeEnvFile({ req, stream, filename: file.filename, - apiKey: result[EnvVar.CODE_API_KEY], + apiKey: resolvedApiKey, entity_id, }); - const updatedMetadata = { - ...file.metadata, - fileIdentifier, - }; - - await updateFile({ - file_id: file.file_id, - metadata: updatedMetadata, - }); - logger.debug( `[provisionToCodeEnv] Provisioned file "${file.filename}" (${file.file_id}) to code env`, ); - return fileIdentifier; + return { + fileIdentifier, + fileUpdate: { file_id: file.file_id, metadata: { ...file.metadata, fileIdentifier } }, + }; } /** @@ -71,27 +76,31 @@ async function provisionToCodeEnv({ req, file, entity_id = '' }) { * @param {object} params.req - Express request object * @param {import('librechat-data-provider').TFile} params.file - The file record from DB * @param {string} [params.entity_id] - Optional entity ID (agent_id) - * @returns {Promise<{ embedded: boolean }>} Embedding result + * @param {import('stream').Readable} [params.existingStream] - Pre-fetched download stream (avoids duplicate storage fetch) + * @returns {Promise<{ embedded: boolean, fileUpdate: object | null }>} Result with deferred DB update */ -async function provisionToVectorDB({ req, file, entity_id }) { +async function provisionToVectorDB({ req, file, entity_id, existingStream }) { if (!process.env.RAG_API_URL) { logger.warn('[provisionToVectorDB] RAG_API_URL not defined, skipping vector provisioning'); - return { embedded: false }; + return { embedded: false, fileUpdate: null }; } - const { getDownloadStream } = getStrategyFunctions(file.source); - if (!getDownloadStream) { - throw new Error( - `Cannot provision file "${file.filename}" to vector DB: storage source "${file.source}" does not support download streams`, - ); - } - - // The uploadVectors function expects a file-like object with a `path` property for fs.createReadStream. - // Since we're provisioning from storage (not a multer upload), we need to stream to a temp file first. const tmpPath = path.join(os.tmpdir(), `provision-${file.file_id}${path.extname(file.filename)}`); try { - const stream = await getDownloadStream(req, file.filepath); + let stream = existingStream; + if (!stream) { + const { getDownloadStream } = getStrategyFunctions(file.source); + if (!getDownloadStream) { + throw new Error( + `Cannot provision file "${file.filename}" to vector DB: storage source "${file.source}" does not support download streams`, + ); + } + stream = await getDownloadStream(req, file.filepath); + } + + // uploadVectors expects a file-like object with a `path` property for fs.createReadStream. + // Since we're provisioning from storage (not a multer upload), we stream to a temp file first. await new Promise((resolve, reject) => { const writeStream = fs.createWriteStream(tmpPath); stream.pipe(writeStream); @@ -117,18 +126,15 @@ async function provisionToVectorDB({ req, file, entity_id }) { const embedded = embeddingResult?.embedded ?? false; - await updateFile({ - file_id: file.file_id, - embedded, - }); - logger.debug( `[provisionToVectorDB] Provisioned file "${file.filename}" (${file.file_id}) to vector DB, embedded=${embedded}`, ); - return { embedded }; + return { + embedded, + fileUpdate: embedded ? { file_id: file.file_id, embedded } : null, + }; } finally { - // Clean up temp file try { fs.unlinkSync(tmpPath); } catch { @@ -190,13 +196,11 @@ async function checkCodeEnvFileAlive({ file, apiKey }) { * * @param {object} params * @param {import('librechat-data-provider').TFile[]} params.files - Files with metadata.fileIdentifier - * @param {string} params.userId - User ID for loading CODE_API_KEY + * @param {string} params.apiKey - Pre-loaded CODE_API_KEY * @param {number} [params.staleSafeWindowMs=21600000] - Skip check if file updated within this window (default 6h) * @returns {Promise>} Set of file_ids that are confirmed alive */ -async function checkSessionsAlive({ files, userId, staleSafeWindowMs = 6 * 60 * 60 * 1000 }) { - const result = await loadAuthValues({ userId, authFields: [EnvVar.CODE_API_KEY] }); - const apiKey = result[EnvVar.CODE_API_KEY]; +async function checkSessionsAlive({ files, apiKey, staleSafeWindowMs = 6 * 60 * 60 * 1000 }) { const aliveFileIds = new Set(); const now = Date.now(); @@ -263,6 +267,7 @@ async function checkSessionsAlive({ files, userId, staleSafeWindowMs = 6 * 60 * } module.exports = { + loadCodeApiKey, provisionToCodeEnv, provisionToVectorDB, checkCodeEnvFileAlive, diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index d1d4f5ad78..c3f67c9478 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -32,10 +32,12 @@ import { generateArtifactsPrompt } from '~/prompts'; import { getProviderConfig } from '~/endpoints'; import { primeResources } from './resources'; import type { + TFileUpdate, TFilterFilesByAgentAccess, TProvisionToCodeEnv, TProvisionToVectorDB, TCheckSessionsAlive, + TLoadCodeApiKey, } from './resources'; /** @@ -156,6 +158,10 @@ export interface InitializeAgentDbMethods extends EndpointDbMethods { provisionToVectorDB?: TProvisionToVectorDB; /** Optional: batch-check code env file liveness */ checkSessionsAlive?: TCheckSessionsAlive; + /** Optional: load CODE_API_KEY once per request */ + loadCodeApiKey?: TLoadCodeApiKey; + /** Optional: persist file metadata updates after provisioning */ + updateFile?: (data: TFileUpdate) => Promise; } /** @@ -316,6 +322,8 @@ export async function initializeAgent( provisionToCodeEnv: db.provisionToCodeEnv, provisionToVectorDB: db.provisionToVectorDB, checkSessionsAlive: db.checkSessionsAlive, + loadCodeApiKey: db.loadCodeApiKey, + updateFile: db.updateFile as ((data: TFileUpdate) => Promise) | undefined, }); const { diff --git a/packages/api/src/agents/resources.ts b/packages/api/src/agents/resources.ts index 04b10fb267..45f258a223 100644 --- a/packages/api/src/agents/resources.ts +++ b/packages/api/src/agents/resources.ts @@ -5,25 +5,34 @@ import type { IMongoFile, AppConfig, IUser } from '@librechat/data-schemas'; import type { FilterQuery, QueryOptions, ProjectionType } from 'mongoose'; import type { Request as ServerRequest } from 'express'; +/** Deferred DB update from provisioning (batched after all files are provisioned) */ +export type TFileUpdate = { + file_id: string; + metadata?: Record; + embedded?: boolean; +}; + /** * Function type for provisioning a file to the code execution environment. - * @returns The fileIdentifier from the code env + * @returns The fileIdentifier and a deferred DB update object */ export type TProvisionToCodeEnv = (params: { req: ServerRequest & { user?: IUser }; file: TFile; entity_id?: string; -}) => Promise; + apiKey?: string; +}) => Promise<{ fileIdentifier: string; fileUpdate: TFileUpdate }>; /** * Function type for provisioning a file to the vector DB for file_search. - * @returns Object with embedded status + * @returns Object with embedded status and a deferred DB update object */ export type TProvisionToVectorDB = (params: { req: ServerRequest & { user?: IUser }; file: TFile; entity_id?: string; -}) => Promise<{ embedded: boolean }>; + existingStream?: unknown; +}) => Promise<{ embedded: boolean; fileUpdate: TFileUpdate | null }>; /** * Function type for batch-checking code env file liveness. @@ -32,10 +41,13 @@ export type TProvisionToVectorDB = (params: { */ export type TCheckSessionsAlive = (params: { files: TFile[]; - userId: string; + apiKey: string; staleSafeWindowMs?: number; }) => Promise>; +/** Loads CODE_API_KEY for a user. Call once per request. */ +export type TLoadCodeApiKey = (userId: string) => Promise; + /** * Function type for retrieving files from the database * @param filter - MongoDB filter query for files @@ -197,6 +209,8 @@ export const primeResources = async ({ provisionToCodeEnv, provisionToVectorDB, checkSessionsAlive, + loadCodeApiKey, + updateFile, }: { req: ServerRequest & { user?: IUser }; appConfig?: AppConfig; @@ -214,6 +228,10 @@ export const primeResources = async ({ provisionToVectorDB?: TProvisionToVectorDB; /** Optional callback to batch-check code env file liveness by session */ checkSessionsAlive?: TCheckSessionsAlive; + /** Optional callback to load CODE_API_KEY once per request */ + loadCodeApiKey?: TLoadCodeApiKey; + /** Optional callback to persist file metadata updates after provisioning */ + updateFile?: (data: TFileUpdate) => Promise; }): Promise<{ attachments: Array; tool_resources: AgentToolResources | undefined; @@ -366,20 +384,34 @@ export const primeResources = async ({ enabledToolResources.has(EToolResources.file_search) && provisionToVectorDB != null; if (needsCodeEnv || needsVectorDB) { + // Load CODE_API_KEY once for all code env operations + let codeApiKey: string | undefined; + if (needsCodeEnv && loadCodeApiKey && req.user?.id) { + try { + codeApiKey = await loadCodeApiKey(req.user.id); + } catch (error) { + logger.error('[primeResources] Failed to load CODE_API_KEY', error); + warnings.push('Code execution file provisioning unavailable'); + } + } + // Batch staleness check: verify code env files are still alive let aliveFileIds: Set = new Set(); - if (needsCodeEnv && checkSessionsAlive && req.user?.id) { + if (needsCodeEnv && codeApiKey && checkSessionsAlive) { const filesWithIdentifiers = attachments.filter( (f) => f?.metadata?.fileIdentifier && f.file_id, ); if (filesWithIdentifiers.length > 0) { aliveFileIds = await checkSessionsAlive({ files: filesWithIdentifiers as TFile[], - userId: req.user.id, + apiKey: codeApiKey, }); } } + // Collect deferred DB updates from provisioning + const pendingUpdates: TFileUpdate[] = []; + // Provision files in parallel const provisionResults = await Promise.allSettled( attachments.map(async (file) => { @@ -393,6 +425,7 @@ export const primeResources = async ({ // Code env provisioning (with staleness check) if ( needsCodeEnv && + codeApiKey && !processedResourceFiles.has(`${EToolResources.execute_code}:${file.file_id}`) ) { const hasFileIdentifier = !!file.metadata?.fileIdentifier; @@ -408,12 +441,14 @@ export const primeResources = async ({ } try { - const fileIdentifier = await provisionToCodeEnv({ + const { fileIdentifier, fileUpdate } = await provisionToCodeEnv({ req: typedReq, file, entity_id: agentId, + apiKey: codeApiKey, }); file.metadata = { ...file.metadata, fileIdentifier }; + pendingUpdates.push(fileUpdate); addFileToResource({ file, resourceType: EToolResources.execute_code, @@ -451,6 +486,9 @@ export const primeResources = async ({ }); if (result.embedded) { file.embedded = true; + if (result.fileUpdate) { + pendingUpdates.push(result.fileUpdate); + } addFileToResource({ file, resourceType: EToolResources.file_search, @@ -473,6 +511,11 @@ export const primeResources = async ({ logger.error('[primeResources] Unexpected provisioning rejection', result.reason); } } + + // Batch DB updates after all provisioning completes + if (pendingUpdates.length > 0 && updateFile) { + await Promise.allSettled(pendingUpdates.map((update) => updateFile(update))); + } } }