From ade02054c84e4bbf1cd13cf32b625f69960ccbf0 Mon Sep 17 00:00:00 2001 From: MarcAmick <5194465+MarcAmick@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:27:09 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=9F=20fix:=20Keep=20File=20Uploads=20A?= =?UTF-8?q?live=20With=20SSE=20Heartbeats=20(#14295)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: Use SSE to upload files in order to avoid idle timeouts. Idle timeouts can occur for example from gateways and other services like cloudfare when uploading large files. For example during rag processing the file is uploaded to librechat which then sends it to rag. While librechat is waiting for the embeddings to come back from rag the file upload is sitting idle. Gateways tend to want to cancel the upload with an http 408 , 504, or 524. This change uses SSE to perform the upload so that while librechat is sending the file to rag, it consistently sends back a heartbeat event to the client to keep the connection alive. This is especially useful when utilizing EMBEDDING_BATCH_SIZE in librechat rag which will allow rag to process signifigantly larger files without running out of memory. * added tests to packages\api\src\files\sse.spec.ts in order to test the new sse.ts * fix: Harden SSE file upload lifecycle * style: Sort data provider imports --------- Co-authored-by: Marc Amick Co-authored-by: Danny Avila --- .env.example | 3 + api/server/routes/config.js | 1 + api/server/routes/files/files.js | 38 ++- api/server/routes/files/images.js | 37 ++- api/server/services/Files/process.js | 20 +- api/server/services/Files/process.spec.js | 7 + client/src/data-provider/Files/mutations.ts | 9 +- packages/api/src/files/index.ts | 1 + packages/api/src/files/sse.spec.ts | 218 ++++++++++++++++++ packages/api/src/files/sse.ts | 100 ++++++++ .../specs/request-interceptor.spec.ts | 66 +++++- packages/data-provider/src/config.ts | 1 + packages/data-provider/src/data-service.ts | 14 +- packages/data-provider/src/request.ts | 73 ++++-- packages/data-provider/src/types.ts | 4 +- packages/data-provider/src/upload.spec.ts | 152 ++++++++++++ packages/data-provider/src/upload.ts | 212 +++++++++++++++++ 17 files changed, 919 insertions(+), 37 deletions(-) create mode 100644 packages/api/src/files/sse.spec.ts create mode 100644 packages/api/src/files/sse.ts create mode 100644 packages/data-provider/src/upload.spec.ts create mode 100644 packages/data-provider/src/upload.ts diff --git a/.env.example b/.env.example index 44821dc0fe..2a15904b1f 100644 --- a/.env.example +++ b/.env.example @@ -577,6 +577,9 @@ TTS_API_KEY= # EMBEDDINGS_PROVIDER=openai # EMBEDDINGS_MODEL=text-embedding-3-small +# Stream upload responses with heartbeats during long-running file processing. +# FILE_UPLOAD_SSE_ENABLED=false + #===================================================# # User System # #===================================================# diff --git a/api/server/routes/config.js b/api/server/routes/config.js index 40787662d7..03044f533e 100644 --- a/api/server/routes/config.js +++ b/api/server/routes/config.js @@ -276,6 +276,7 @@ router.get('/', async function (req, res) { : 0, ...(cloudFront ? { cloudFront } : {}), ...(rum ? { rum } : {}), + fileUploadSseEnabled: isEnabled(process.env.FILE_UPLOAD_SSE_ENABLED), }; const webSearch = buildWebSearchConfig(appConfig); diff --git a/api/server/routes/files/files.js b/api/server/routes/files/files.js index fafa02e66d..d98ed3b6f8 100644 --- a/api/server/routes/files/files.js +++ b/api/server/routes/files/files.js @@ -5,6 +5,8 @@ const { logAxiosError, refreshS3FileUrls, handleFilesUsageRequest, + shouldUseUploadSse, + startUploadSseStream, resolveUploadErrorMessage, verifyAgentUploadPermission, } = require('@librechat/api'); @@ -633,6 +635,15 @@ router.post('/', async (req, res) => { const metadata = req.body; let cleanup = true; + /** Opened only once auth/validation has passed, right before the potentially + * long-running upload processing begins — see `startUploadSseStream`. */ + let sseStream = null; + const openSseStreamIfRequested = () => { + if (shouldUseUploadSse(req)) { + sseStream = startUploadSseStream(res); + } + }; + try { filterFile({ req }); @@ -640,7 +651,8 @@ router.post('/', async (req, res) => { metadata.file_id = req.file_id; if (isAssistantsEndpoint(metadata.endpoint)) { - return await processFileUpload({ req, res, metadata }); + openSseStreamIfRequested(); + return await processFileUpload({ req, res, metadata, sseStream }); } let skipUploadAuth = false; @@ -663,7 +675,8 @@ router.post('/', async (req, res) => { } } - return await processAgentFileUpload({ req, res, metadata }); + openSseStreamIfRequested(); + return await processAgentFileUpload({ req, res, metadata, sseStream }); } catch (error) { const message = resolveUploadErrorMessage(error); logger.error('[/files] Error processing file:', error); @@ -674,7 +687,23 @@ router.post('/', async (req, res) => { } catch (error) { logger.error('[/files] Error deleting file:', error); } - res.status(500).json({ message }); + + let errorStatusCode = 500; + if (error.userErrorStatusCode) { + errorStatusCode = error.userErrorStatusCode; + } + + if (sseStream) { + sseStream.sendError({ + message, + code: errorStatusCode, + temp_file_id: metadata.temp_file_id, + tool_resource: metadata.tool_resource, + display_to_user: true, + }); + } else { + res.status(errorStatusCode).json({ message }); + } } finally { if (cleanup) { try { @@ -685,6 +714,9 @@ router.post('/', async (req, res) => { } else { logger.debug('[/files] File processing completed without cleanup'); } + if (sseStream) { + sseStream.close(); + } } }); diff --git a/api/server/routes/files/images.js b/api/server/routes/files/images.js index 353557dc4f..7289a99310 100644 --- a/api/server/routes/files/images.js +++ b/api/server/routes/files/images.js @@ -2,7 +2,12 @@ const path = require('path'); const fs = require('fs').promises; const express = require('express'); const { logger } = require('@librechat/data-schemas'); -const { verifyAgentUploadPermission, resolveUploadErrorMessage } = require('@librechat/api'); +const { + shouldUseUploadSse, + startUploadSseStream, + resolveUploadErrorMessage, + verifyAgentUploadPermission, +} = require('@librechat/api'); const { isAssistantsEndpoint } = require('librechat-data-provider'); const { processAgentFileUpload, @@ -18,6 +23,15 @@ router.post('/', async (req, res) => { const metadata = req.body; const appConfig = req.config; + /** Opened only once auth/validation has passed, right before the potentially + * long-running upload processing begins — see `startUploadSseStream`. */ + let sseStream = null; + const openSseStreamIfRequested = () => { + if (shouldUseUploadSse(req)) { + sseStream = startUploadSseStream(res); + } + }; + try { filterFile({ req, image: true }); @@ -35,10 +49,12 @@ router.post('/', async (req, res) => { if (denied) { return; } - return await processAgentFileUpload({ req, res, metadata }); + openSseStreamIfRequested(); + return await processAgentFileUpload({ req, res, metadata, sseStream }); } - await processImageFile({ req, res, metadata }); + openSseStreamIfRequested(); + await processImageFile({ req, res, metadata, sseStream }); } catch (error) { // TODO: delete remote file if it exists logger.error('[/files/images] Error processing file:', error); @@ -55,7 +71,17 @@ router.post('/', async (req, res) => { } catch (error) { logger.error('[/files/images] Error deleting file:', error); } - res.status(500).json({ message }); + if (sseStream) { + sseStream.sendError({ + message, + code: 500, + temp_file_id: metadata.temp_file_id, + tool_resource: metadata.tool_resource, + display_to_user: true, + }); + } else { + res.status(500).json({ message }); + } } finally { try { await fs.unlink(req.file.path); @@ -63,6 +89,9 @@ router.post('/', async (req, res) => { } catch { logger.debug('[/files/images] Temp. image upload file already deleted'); } + if (sseStream) { + sseStream.close(); + } } }); diff --git a/api/server/services/Files/process.js b/api/server/services/Files/process.js index 5aa3443fe8..527e7b0e31 100644 --- a/api/server/services/Files/process.js +++ b/api/server/services/Files/process.js @@ -24,6 +24,7 @@ const { sanitizeFilename, parseText, processAudioFile, + sendUploadSuccess, getStorageMetadata, sweepExpiredFiles: sweepExpiredFilesWithDeps, startExpiredFileSweep: startExpiredFileSweepWithDeps, @@ -447,9 +448,10 @@ const processFileURL = async ({ * @param {Express.Response} [params.res] - The Express response object. * @param {ImageMetadata} params.metadata - Additional metadata for the file. * @param {boolean} params.returnFile - Whether to return the file metadata or return response as normal. + * @param {import('@librechat/api').UploadSseStream | null} [params.sseStream] - Active upload SSE stream, if enabled. * @returns {Promise} */ -const processImageFile = async ({ req, res, metadata, returnFile = false }) => { +const processImageFile = async ({ req, res, metadata, returnFile = false, sseStream }) => { const { file } = req; const appConfig = req.config; const source = getFileStrategy(appConfig, { isImage: true }); @@ -487,7 +489,7 @@ const processImageFile = async ({ req, res, metadata, returnFile = false }) => { if (returnFile) { return result; } - res.status(200).json({ message: 'File uploaded and processed successfully', ...result }); + sendUploadSuccess(res, sseStream, 'File uploaded and processed successfully', result); }; /** @@ -554,9 +556,10 @@ const uploadImageBuffer = async ({ req, context, metadata = {}, resize = true }) * @param {ServerRequest} params.req - The Express request object. * @param {Express.Response} params.res - The Express response object. * @param {FileMetadata} params.metadata - Additional metadata for the file. + * @param {import('@librechat/api').UploadSseStream | null} [params.sseStream] - Active upload SSE stream, if enabled. * @returns {Promise} */ -const processFileUpload = async ({ req, res, metadata }) => { +const processFileUpload = async ({ req, res, metadata, sseStream }) => { const appConfig = req.config; const isAssistantUpload = isAssistantsEndpoint(metadata.endpoint); const assistantSource = @@ -649,7 +652,7 @@ const processFileUpload = async ({ req, res, metadata }) => { }, true, ); - res.status(200).json({ message: 'File uploaded and processed successfully', ...result }); + sendUploadSuccess(res, sseStream, 'File uploaded and processed successfully', result); }; /** @@ -661,9 +664,10 @@ const processFileUpload = async ({ req, res, metadata }) => { * @param {ServerRequest} params.req - The Express request object. * @param {Express.Response} params.res - The Express response object. * @param {FileMetadata} params.metadata - Additional metadata for the file. + * @param {import('@librechat/api').UploadSseStream | null} [params.sseStream] - Active upload SSE stream, if enabled. * @returns {Promise} */ -const processAgentFileUpload = async ({ req, res, metadata }) => { +const processAgentFileUpload = async ({ req, res, metadata, sseStream }) => { const { file } = req; const appConfig = req.config; const { agent_id, tool_resource, file_id, temp_file_id = null } = metadata; @@ -787,9 +791,7 @@ const processAgentFileUpload = async ({ req, res, metadata }) => { }); } const result = await db.createFile(fileInfo, true); - return res - .status(200) - .json({ message: 'Agent file uploaded and processed successfully', ...result }); + sendUploadSuccess(res, sseStream, 'Agent file uploaded and processed successfully', result); }; const fileConfig = mergeFileConfig(appConfig.fileConfig); @@ -1035,7 +1037,7 @@ const processAgentFileUpload = async ({ req, res, metadata }) => { const result = await db.createFile(fileInfo, true); - res.status(200).json({ message: 'Agent file uploaded and processed successfully', ...result }); + sendUploadSuccess(res, sseStream, 'Agent file uploaded and processed successfully', result); }; /** diff --git a/api/server/services/Files/process.spec.js b/api/server/services/Files/process.spec.js index c5998fee4c..48c37166df 100644 --- a/api/server/services/Files/process.spec.js +++ b/api/server/services/Files/process.spec.js @@ -41,6 +41,13 @@ jest.mock('@librechat/api', () => { sanitizeFilename: jest.fn((n) => n), parseText: jest.fn().mockResolvedValue({ text: '', bytes: 0 }), processAudioFile: jest.fn(), + sendUploadSuccess: jest.fn((res, sseStream, message, result) => { + if (sseStream) { + sseStream.sendData({ message, ...result }); + return; + } + res.status(200).json({ message, ...result }); + }), getStorageMetadata: jest.fn(() => ({})), getRetentionExpiry, getAgentFileRetentionExpiry: jest.fn(({ req, messageAttachment, toolResource }) => { diff --git a/client/src/data-provider/Files/mutations.ts b/client/src/data-provider/Files/mutations.ts index d98da899f9..a538eea4b8 100644 --- a/client/src/data-provider/Files/mutations.ts +++ b/client/src/data-provider/Files/mutations.ts @@ -10,6 +10,7 @@ import { } from 'librechat-data-provider'; import type { UseMutationResult } from '@tanstack/react-query'; import type * as t from 'librechat-data-provider'; +import { useGetStartupConfig } from '../Endpoints'; import { useLocalize } from '~/hooks'; export const useUploadFileMutation = ( @@ -21,6 +22,8 @@ export const useUploadFileMutation = ( FormData, // request unknown // context > => { + const { data: startupConfig } = useGetStartupConfig(); + const sseEnabled = startupConfig?.fileUploadSseEnabled === true; const queryClient = useQueryClient(); const { onSuccess, ...options } = _options || {}; return useMutation([MutationKeys.fileUpload], { @@ -30,14 +33,14 @@ export const useUploadFileMutation = ( const version = body.get('version') ?? ''; const endpoint = (body.get('endpoint') ?? '') as string; if (isAssistantsEndpoint(endpoint) && version === '2') { - return dataService.uploadFile(body, signal); + return dataService.uploadFile(body, signal, sseEnabled); } if (width !== '' && height !== '') { - return dataService.uploadImage(body, signal); + return dataService.uploadImage(body, signal, sseEnabled); } - return dataService.uploadFile(body, signal); + return dataService.uploadFile(body, signal, sseEnabled); }, ...options, onSuccess: (data, formData, context) => { diff --git a/packages/api/src/files/index.ts b/packages/api/src/files/index.ts index 7cfc9aff13..cbf79614ab 100644 --- a/packages/api/src/files/index.ts +++ b/packages/api/src/files/index.ts @@ -10,6 +10,7 @@ export * from './ocr'; export * from './parse'; export * from './rag'; export * from './retention'; +export * from './sse'; export * from './sweep'; export * from './usage'; export * from './validation'; diff --git a/packages/api/src/files/sse.spec.ts b/packages/api/src/files/sse.spec.ts new file mode 100644 index 0000000000..dc42bd33f9 --- /dev/null +++ b/packages/api/src/files/sse.spec.ts @@ -0,0 +1,218 @@ +import { EventEmitter } from 'events'; +import type { Request, Response } from 'express'; +import { sendUploadSuccess, shouldUseUploadSse, startUploadSseStream } from './sse'; + +describe('sse', () => { + const createMockReq = (accept?: string): Request => + ({ headers: accept ? { accept } : {} }) as Request; + + const createMockRes = (): jest.Mocked => { + const res = new EventEmitter() as unknown as jest.Mocked; + res.writeHead = jest.fn().mockReturnValue(res); + res.flushHeaders = jest.fn().mockReturnValue(res); + res.write = jest.fn().mockReturnValue(true); + res.end = jest.fn().mockReturnValue(res); + res.status = jest.fn().mockReturnValue(res); + res.json = jest.fn().mockReturnValue(res); + Object.defineProperty(res, 'writableEnded', { value: false, writable: true }); + Object.defineProperty(res, 'destroyed', { value: false, writable: true }); + return res; + }; + + const parseEvents = (res: jest.Mocked): Array<{ event: string; data: unknown }> => + (res.write as jest.Mock).mock.calls.map(([chunk]: [string]) => { + const eventMatch = /event:(.*)\n/.exec(chunk); + const dataMatch = /data:(.*)\n\n/.exec(chunk); + return { + event: eventMatch ? eventMatch[1] : '', + data: dataMatch ? JSON.parse(dataMatch[1]) : undefined, + }; + }); + + describe('shouldUseUploadSse', () => { + const originalValue = process.env.FILE_UPLOAD_SSE_ENABLED; + + afterEach(() => { + if (originalValue === undefined) { + delete process.env.FILE_UPLOAD_SSE_ENABLED; + return; + } + process.env.FILE_UPLOAD_SSE_ENABLED = originalValue; + }); + + it('requires both the feature flag and an explicit event-stream accept value', () => { + process.env.FILE_UPLOAD_SSE_ENABLED = 'true'; + + expect(shouldUseUploadSse(createMockReq('application/json, text/event-stream'))).toBe(true); + expect(shouldUseUploadSse(createMockReq('text/event-stream; charset=utf-8'))).toBe(true); + expect(shouldUseUploadSse(createMockReq('application/json'))).toBe(false); + expect(shouldUseUploadSse(createMockReq('*/*'))).toBe(false); + expect(shouldUseUploadSse(createMockReq())).toBe(false); + }); + + it('keeps JSON responses when the server feature flag is disabled', () => { + process.env.FILE_UPLOAD_SSE_ENABLED = 'false'; + + expect(shouldUseUploadSse(createMockReq('text/event-stream'))).toBe(false); + }); + }); + + describe('startUploadSseStream', () => { + beforeEach(() => { + jest.useFakeTimers(); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('writes the SSE headers and flushes them immediately', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + + expect(res.writeHead).toHaveBeenCalledWith( + 200, + expect.objectContaining({ + 'Content-Type': 'text/event-stream', + Connection: 'keep-alive', + }), + ); + expect(res.flushHeaders).toHaveBeenCalledTimes(1); + stream.close(); + }); + + it('emits a heartbeat event on every interval tick', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + + jest.advanceTimersByTime(3000); + + const heartbeats = parseEvents(res).filter((e) => e.event === 'heartbeat'); + expect(heartbeats).toHaveLength(3); + expect(heartbeats.map((e) => e.data)).toEqual([ + { keepAlive: 1 }, + { keepAlive: 2 }, + { keepAlive: 3 }, + ]); + stream.close(); + }); + + it('stops the heartbeat once the response has already ended', () => { + const res = createMockRes(); + startUploadSseStream(res); + + jest.advanceTimersByTime(1000); + Object.defineProperty(res, 'writableEnded', { value: true }); + jest.advanceTimersByTime(2000); + + const heartbeats = parseEvents(res).filter((e) => e.event === 'heartbeat'); + expect(heartbeats).toHaveLength(1); + expect(jest.getTimerCount()).toBe(0); + }); + + it('clears the heartbeat interval on response disconnect and stops writes', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + expect(jest.getTimerCount()).toBe(1); + + Object.defineProperty(res, 'destroyed', { value: true }); + res.emit('close'); + stream.sendData({ file_id: 'ignored' }); + + expect(jest.getTimerCount()).toBe(0); + expect(res.write).not.toHaveBeenCalled(); + }); + + it('emits data and error events', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + + stream.sendData({ file_id: 'abc' }); + stream.sendError({ message: 'failed' }); + + const events = parseEvents(res); + expect(events.find((event) => event.event === 'data')?.data).toEqual({ file_id: 'abc' }); + expect(events.find((event) => event.event === 'error')?.data).toEqual({ + message: 'failed', + }); + stream.close(); + }); + + describe('close', () => { + it('emits a close event, ends the response, and clears the heartbeat', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + + stream.close(); + + const closeEvents = parseEvents(res).filter((e) => e.event === 'close'); + expect(closeEvents).toHaveLength(1); + expect(closeEvents[0].data).toEqual( + expect.objectContaining({ closedAt: expect.any(String) }), + ); + expect(res.end).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); + + it('does not write after the response already ended', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + Object.defineProperty(res, 'writableEnded', { value: true }); + + stream.close(); + + expect(parseEvents(res).filter((e) => e.event === 'close')).toHaveLength(0); + expect(res.end).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + it('does not emit further heartbeats after close', () => { + const res = createMockRes(); + const stream = startUploadSseStream(res); + + stream.close(); + jest.advanceTimersByTime(5000); + + expect(parseEvents(res).filter((e) => e.event === 'heartbeat')).toHaveLength(0); + }); + }); + }); + + describe('sendUploadSuccess', () => { + it('sends the payload over the SSE stream when a stream is active', () => { + const res = createMockRes(); + const sseStream = { + sendData: jest.fn(), + sendError: jest.fn(), + close: jest.fn(), + }; + + sendUploadSuccess(res, sseStream, 'Upload complete', { file_id: 'abc' }); + + expect(sseStream.sendData).toHaveBeenCalledWith({ + message: 'Upload complete', + file_id: 'abc', + }); + expect(res.status).not.toHaveBeenCalled(); + expect(res.json).not.toHaveBeenCalled(); + }); + + it('falls back to a plain JSON response when no stream is provided', () => { + const res = createMockRes(); + + sendUploadSuccess(res, null, 'Upload complete', { file_id: 'abc' }); + + expect(res.status).toHaveBeenCalledWith(200); + expect(res.json).toHaveBeenCalledWith({ message: 'Upload complete', file_id: 'abc' }); + }); + + it('falls back to a plain JSON response when the stream is undefined', () => { + const res = createMockRes(); + + sendUploadSuccess(res, undefined, 'Upload complete', { file_id: 'abc' }); + + expect(res.status).toHaveBeenCalledWith(200); + expect(res.json).toHaveBeenCalledWith({ message: 'Upload complete', file_id: 'abc' }); + }); + }); +}); diff --git a/packages/api/src/files/sse.ts b/packages/api/src/files/sse.ts new file mode 100644 index 0000000000..66b0397535 --- /dev/null +++ b/packages/api/src/files/sse.ts @@ -0,0 +1,100 @@ +import type { Request, Response } from 'express'; +import { isEnabled } from '~/utils'; + +const HEARTBEAT_INTERVAL_MS = 1000; +const EVENT_STREAM_MEDIA_TYPE = 'text/event-stream'; + +const SSE_HEADERS = { + 'Content-Type': EVENT_STREAM_MEDIA_TYPE, + 'Cache-Control': 'no-cache, no-transform', + Connection: 'keep-alive', + /** Required so Nginx/other proxies don't buffer the stream. */ + 'X-Accel-Buffering': 'no', +}; + +function writeSseEvent(res: Response, event: string, data: T): void { + if (res.writableEnded || res.destroyed) { + return; + } + res.write(`event:${event}\nid:${Date.now()}\ndata:${JSON.stringify(data)}\n\n`); +} + +export function shouldUseUploadSse(req: Request): boolean { + if (!isEnabled(process.env.FILE_UPLOAD_SSE_ENABLED)) { + return false; + } + + return ( + req.headers.accept?.split(',').some((value) => { + const [mediaType] = value.split(';', 1); + return mediaType.trim().toLowerCase() === EVENT_STREAM_MEDIA_TYPE; + }) ?? false + ); +} + +export interface UploadSseStream { + /** Emits the successful upload payload as an `event:data` message. */ + sendData: (data: T) => void; + /** Emits a failure payload as an `event:error` message. */ + sendError: (data: T) => void; + /** Emits the terminal `event:close` message, stops the heartbeat, and ends the response. */ + close: () => void; +} + +/** + * Opens a keep-alive SSE stream for a file-upload response: sends the SSE headers, starts a + * heartbeat interval so proxies/clients don't time out the connection during long-running + * uploads, and stops the heartbeat when the client disconnects early. + * + * Callers must only invoke this once all synchronous validation and permission checks that + * might still need to send a normal (non-SSE) response have already passed — once the headers + * are flushed here, the HTTP status is committed to 200 and can no longer be changed. + */ +export function startUploadSseStream(res: Response): UploadSseStream { + res.writeHead(200, SSE_HEADERS); + res.flushHeaders(); + + let counter = 1; + const intervalId = setInterval(() => { + if (res.writableEnded || res.destroyed) { + clearInterval(intervalId); + return; + } + writeSseEvent(res, 'heartbeat', { keepAlive: counter++ }); + }, HEARTBEAT_INTERVAL_MS); + + const stopHeartbeat = () => { + clearInterval(intervalId); + }; + res.once('close', stopHeartbeat); + + return { + sendData: (data) => writeSseEvent(res, 'data', data), + sendError: (data) => writeSseEvent(res, 'error', data), + close: () => { + clearInterval(intervalId); + res.off('close', stopHeartbeat); + if (!res.writableEnded && !res.destroyed) { + writeSseEvent(res, 'close', { closedAt: new Date().toISOString() }); + res.end(); + } + }, + }; +} + +/** + * Sends a successful upload response, either as an SSE `event:data` message (when an upload + * stream is active) or as a plain JSON response. + */ +export function sendUploadSuccess( + res: Response, + sseStream: UploadSseStream | null | undefined, + message: string, + result: T, +): void { + if (sseStream) { + sseStream.sendData({ message, ...result }); + return; + } + res.status(200).json({ message, ...result }); +} diff --git a/packages/data-provider/specs/request-interceptor.spec.ts b/packages/data-provider/specs/request-interceptor.spec.ts index ae4a58d484..654161342a 100644 --- a/packages/data-provider/specs/request-interceptor.spec.ts +++ b/packages/data-provider/specs/request-interceptor.spec.ts @@ -20,6 +20,7 @@ import { setTokenHeader } from '../src/headers-helpers'; const mockAdapter = jest.fn(); let originalAdapter: typeof axios.defaults.adapter; let savedLocation: Location; +let dataRequest: typeof import('../src/request').default; type RetryableAdapterConfig = InternalAxiosRequestConfig & { _retry?: boolean }; @@ -74,7 +75,7 @@ beforeAll(async () => { originalAdapter = axios.defaults.adapter; axios.defaults.adapter = mockAdapter; - await import('../src/request'); + dataRequest = (await import('../src/request')).default; }); beforeEach(() => { @@ -659,6 +660,69 @@ describe('axios 401 interceptor — Authorization header guard', () => { expect(mockAdapter.mock.calls[1][0].headers?.Authorization).toBe('Bearer fresh-token'); }); + it('uses shared proactive refresh for authenticated fetch requests', async () => { + expect.assertions(4); + setTokenHeader(createJwt(Date.now() + 60_000)); + + mockAdapter.mockImplementation((config: InternalAxiosRequestConfig) => { + if (config.url?.includes('/api/auth/refresh') === true) { + return createAdapterResponse(config, { token: 'fresh-token' }); + } + return createAdapterResponse(config, { ok: true }); + }); + const fetchSpy = jest.spyOn(globalThis, 'fetch').mockResolvedValue( + new Response(null, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }), + ); + + await dataRequest.authenticatedFetch('/api/files', { + method: 'POST', + body: new FormData(), + headers: { Accept: 'text/event-stream' }, + }); + + expect(getCallsForUrl('/api/auth/refresh')).toHaveLength(1); + expect(fetchSpy).toHaveBeenCalledTimes(1); + const uploadHeaders = new Headers(fetchSpy.mock.calls[0][1]?.headers); + expect(uploadHeaders.get('Authorization')).toBe('Bearer fresh-token'); + expect(uploadHeaders.get('Accept')).toBe('text/event-stream'); + }); + + it('refreshes and retries an authenticated fetch request after a 401', async () => { + expect.assertions(4); + setTokenHeader('expired-token'); + + mockAdapter.mockImplementation((config: InternalAxiosRequestConfig) => { + if (config.url?.includes('/api/auth/refresh') === true) { + return createAdapterResponse(config, { token: 'fresh-token' }); + } + return createAdapterResponse(config, { ok: true }); + }); + const fetchSpy = jest + .spyOn(globalThis, 'fetch') + .mockResolvedValueOnce(new Response(null, { status: 401 })) + .mockResolvedValueOnce( + new Response(null, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }), + ); + + await dataRequest.authenticatedFetch('/api/files', { + method: 'POST', + body: new FormData(), + }); + + expect(getCallsForUrl('/api/auth/refresh')).toHaveLength(1); + expect(fetchSpy).toHaveBeenCalledTimes(2); + const firstHeaders = new Headers(fetchSpy.mock.calls[0][1]?.headers); + const retriedHeaders = new Headers(fetchSpy.mock.calls[1][1]?.headers); + expect(firstHeaders.get('Authorization')).toBe('Bearer expired-token'); + expect(retriedHeaders.get('Authorization')).toBe('Bearer fresh-token'); + }); + it('does not wait on the in-flight recovery when the refresh request itself fails', async () => { expect.assertions(3); setTokenHeader(createJwt(Date.now() + 60_000)); diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 96cebc9cb4..7abc94c1a1 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1578,6 +1578,7 @@ export type TStartupConfig = { branch?: string | null; buildDate?: string | null; }; + fileUploadSseEnabled?: boolean; }; export type TSharedLinkStartupInterface = Pick< diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 756351368d..59a651659b 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -3,6 +3,7 @@ import type { TFileConfig } from './file-config'; import type * as t from './types'; import * as permissions from './accessPermissions'; import * as endpoints from './api-endpoints'; +import { uploadEventStream } from './upload'; import * as mcp from './types/mcpServers'; import * as a from './types/assistants'; import * as m from './types/mutations'; @@ -455,13 +456,24 @@ export const getFileConfig = (): Promise => { export const uploadImage = ( data: FormData, signal?: AbortSignal | null, + sseEnabled = false, ): Promise => { const requestConfig = signal ? { signal } : undefined; + if (sseEnabled) { + return uploadEventStream(endpoints.images(), data, signal); + } return request.postMultiPart(endpoints.images(), data, requestConfig); }; -export const uploadFile = (data: FormData, signal?: AbortSignal | null): Promise => { +export const uploadFile = ( + data: FormData, + signal?: AbortSignal | null, + sseEnabled = false, +): Promise => { const requestConfig = signal ? { signal } : undefined; + if (sseEnabled) { + return uploadEventStream(endpoints.files(), data, signal); + } return request.postMultiPart(endpoints.files(), data, requestConfig); }; diff --git a/packages/data-provider/src/request.ts b/packages/data-provider/src/request.ts index 528fd3d562..bde70e8ae8 100644 --- a/packages/data-provider/src/request.ts +++ b/packages/data-provider/src/request.ts @@ -1,5 +1,6 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import axios, { AxiosRequestConfig } from 'axios'; +import axios from 'axios'; +import type { AxiosRequestConfig } from 'axios'; import type * as t from './types'; import { setTokenHeader } from './headers-helpers'; import * as endpoints from './api-endpoints'; @@ -280,22 +281,63 @@ const shouldRefreshBeforeRequest = (url?: string) => { return timeUntilExpiry > 0 && timeUntilExpiry <= TOKEN_REFRESH_BUFFER_MS; }; +const refreshBeforeRequest = async (url?: string) => { + const state = getAuthRecoveryState(); + if (state.refreshPromise && !isAuthRecoveryEndpoint(url)) { + return state.refreshPromise.catch(() => null); + } + + if (!shouldRefreshBeforeRequest(url)) { + return null; + } + + return startAuthRecovery(false).catch(() => null); +}; + +const withAuthorization = (options: RequestInit | undefined, token: string | null): RequestInit => { + const headers = new Headers(options?.headers); + if (token) { + headers.set('Authorization', `Bearer ${token}`); + } + return { ...options, headers }; +}; + +async function _authenticatedFetch(url: string, options?: RequestInit): Promise { + if (typeof window === 'undefined') { + return fetch(url, options); + } + + const token = (await refreshBeforeRequest(url)) ?? getBearerToken(); + const response = await fetch(url, withAuthorization(options, token)); + if ( + response.status !== 401 || + isAuthRecoveryEndpoint(url) || + isAuthRedirectInProgress() || + !getBearerToken() + ) { + return response; + } + + let refreshedToken: string | null; + try { + refreshedToken = await startAuthRecovery(false); + } catch { + redirectToLoginOnce(); + return response; + } + + if (!refreshedToken) { + redirectToLoginOnce(); + return response; + } + + await response.body?.cancel().catch(() => undefined); + return fetch(url, withAuthorization(options, refreshedToken)); +} + if (typeof window !== 'undefined') { axios.interceptors.request.use(async (config) => { - const state = getAuthRecoveryState(); - if (state.refreshPromise && !isAuthRecoveryEndpoint(config.url)) { - const token = await state.refreshPromise.catch(() => null); - if (token) { - setRequestAuthorizationHeader(config, token); - } - return config; - } - - if (!shouldRefreshBeforeRequest(config.url)) { - return config; - } - - const token = await startAuthRecovery(false).catch(() => null); + const token = await refreshBeforeRequest(config.url); if (token) { setRequestAuthorizationHeader(config, token); } @@ -384,6 +426,7 @@ export default { delete: _delete, deleteWithOptions: _deleteWithOptions, patch: _patch, + authenticatedFetch: _authenticatedFetch, refreshToken, dispatchTokenUpdatedEvent, }; diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index e89ebaf51e..7316d2b33e 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -10,11 +10,11 @@ import type { ReasoningResponseKey, ReasoningParameterFormat, } from './schemas'; +import type { Agent, EToolResources } from './types/assistants'; import type { RefillIntervalUnit } from './balance'; import type { SettingDefinition } from './generate'; import type { TMinimalFeedback } from './feedback'; import type { ContentTypes } from './types/runs'; -import type { Agent } from './types/assistants'; export * from './schemas'; @@ -217,6 +217,8 @@ export type TMarketplaceCategory = TCategory & { export type TError = { message: string; code?: number | string; + file_id?: string; + tool_resource?: EToolResources; response?: { data?: { message?: string; diff --git a/packages/data-provider/src/upload.spec.ts b/packages/data-provider/src/upload.spec.ts new file mode 100644 index 0000000000..0b297e99d6 --- /dev/null +++ b/packages/data-provider/src/upload.spec.ts @@ -0,0 +1,152 @@ +import { uploadEventStream } from './upload'; +import request from './request'; + +jest.mock('./request', () => ({ + __esModule: true, + default: { + authenticatedFetch: jest.fn(), + }, +})); + +const authenticatedFetch = request.authenticatedFetch as jest.Mock; + +const createStream = (...chunks: string[]) => + new ReadableStream({ + start(controller) { + const encoder = new TextEncoder(); + for (const chunk of chunks) { + controller.enqueue(encoder.encode(chunk)); + } + controller.close(); + }, + }); + +const createFormData = () => { + const formData = new FormData(); + formData.set('file_id', 'temporary-id'); + return formData; +}; + +describe('uploadEventStream', () => { + beforeEach(() => { + authenticatedFetch.mockReset(); + }); + + it('parses SSE events split across chunks and waits for the terminal close event', async () => { + authenticatedFetch.mockResolvedValue( + new Response( + createStream( + 'event:heartbeat\r\ndata:{"keepAlive":1}\r\n\r\nevent:da', + 'ta\r\ndata:{"file_id":"stored-id","temp_file_id":"temporary-id"}\r\n\r\n', + 'event:close\r\ndata:{}\r\n\r\n', + ), + { headers: { 'Content-Type': 'text/event-stream; charset=utf-8' } }, + ), + ); + + await expect(uploadEventStream('/api/files', createFormData())).resolves.toMatchObject({ + file_id: 'stored-id', + temp_file_id: 'temporary-id', + }); + }); + + it('falls back to a JSON upload response for mixed-version servers', async () => { + authenticatedFetch.mockResolvedValue( + new Response('{"file_id":"stored-id","temp_file_id":"temporary-id"}', { + headers: { 'Content-Type': 'application/json; charset=utf-8' }, + }), + ); + + await expect(uploadEventStream('/api/files', createFormData())).resolves.toMatchObject({ + file_id: 'stored-id', + temp_file_id: 'temporary-id', + }); + }); + + it('maps streamed failures to the existing upload error shape', async () => { + authenticatedFetch.mockResolvedValue( + new Response( + createStream( + 'event:error\n', + 'data:{"message":"Unsupported file","temp_file_id":"temporary-id","code":422,"display_to_user":true}\n\n', + ), + { headers: { 'Content-Type': 'text/event-stream' } }, + ), + ); + + await expect(uploadEventStream('/api/files', createFormData())).rejects.toMatchObject({ + name: 'CustomAppError', + code: 422, + file_id: 'temporary-id', + display_to_user: true, + response: { data: { message: 'Unsupported file' } }, + }); + }); + + it('maps pre-stream HTTP failures to the existing upload error shape', async () => { + authenticatedFetch.mockResolvedValue( + new Response('{"message":"File is too large"}', { + status: 413, + headers: { 'Content-Type': 'application/json' }, + }), + ); + + await expect(uploadEventStream('/api/files', createFormData())).rejects.toMatchObject({ + name: 'CustomAppError', + code: 413, + file_id: 'temporary-id', + response: { data: { message: 'File is too large' } }, + }); + }); + + it('cancels a stalled stream when heartbeats stop', async () => { + jest.useFakeTimers(); + try { + authenticatedFetch.mockResolvedValue( + new Response(new ReadableStream(), { + headers: { 'Content-Type': 'text/event-stream' }, + }), + ); + + const result = uploadEventStream('/api/files', createFormData()).catch( + (error: Error) => error, + ); + await jest.advanceTimersByTimeAsync(15_000); + + await expect(result).resolves.toMatchObject({ + message: expect.stringContaining('timed out waiting for a heartbeat'), + }); + expect(jest.getTimerCount()).toBe(0); + } finally { + jest.useRealTimers(); + } + }); + + it('preserves the canceled-upload error contract when the signal aborts', async () => { + const controller = new AbortController(); + controller.abort('User aborted upload'); + authenticatedFetch.mockRejectedValue('User aborted upload'); + + await expect( + uploadEventStream('/api/files', createFormData(), controller.signal), + ).rejects.toMatchObject({ code: 'ERR_CANCELED' }); + }); + + it('preserves the upload abort signal', async () => { + const controller = new AbortController(); + authenticatedFetch.mockResolvedValue( + new Response('{"file_id":"stored-id","temp_file_id":"temporary-id"}', { + headers: { 'Content-Type': 'application/json' }, + }), + ); + + await uploadEventStream('/api/files', createFormData(), controller.signal); + + expect(authenticatedFetch).toHaveBeenCalledWith('/api/files', { + method: 'POST', + body: expect.any(FormData), + headers: { Accept: 'text/event-stream' }, + signal: controller.signal, + }); + }); +}); diff --git a/packages/data-provider/src/upload.ts b/packages/data-provider/src/upload.ts new file mode 100644 index 0000000000..0edb86d5df --- /dev/null +++ b/packages/data-provider/src/upload.ts @@ -0,0 +1,212 @@ +import type { EToolResources } from './types/assistants'; +import type { TFileUpload } from './types/files'; +import request from './request'; + +const EVENT_STREAM_MEDIA_TYPE = 'text/event-stream'; +const HEARTBEAT_TIMEOUT_MS = 15_000; + +interface UploadErrorData { + message?: string; + code?: number; + temp_file_id?: string; + tool_resource?: EToolResources; + display_to_user?: boolean; +} + +interface ParsedEvent { + type: string; + data: string; +} + +class FileUploadError extends Error { + public code: number; + public file_id: string; + public tool_resource?: EToolResources; + public display_to_user: boolean; + public response: { data: { message: string } }; + + constructor( + message: string, + fileId: string, + toolResource?: EToolResources, + displayToUser = false, + code = 0, + ) { + super(message); + this.name = 'CustomAppError'; + this.code = code; + this.file_id = fileId; + this.tool_resource = toolResource; + this.display_to_user = displayToUser; + this.response = { data: { message: displayToUser ? message : '' } }; + } +} + +class UploadCanceledError extends Error { + public code = 'ERR_CANCELED' as const; +} + +const getFileId = (formData: FormData) => String(formData.get('file_id') ?? ''); + +const getToolResource = (formData: FormData) => + (formData.get('tool_resource') as EToolResources | null) ?? undefined; + +const parseEvent = (message: string): ParsedEvent => { + let type = 'message'; + const data: string[] = []; + + for (const line of message.split(/\r?\n/)) { + if (line.startsWith('event:')) { + type = line.slice('event:'.length).trim(); + continue; + } + if (line.startsWith('data:')) { + data.push(line.slice('data:'.length).trimStart()); + } + } + + return { type, data: data.join('\n') }; +}; + +const createHttpError = async (response: Response, formData: FormData) => { + let message = `Server responded with status: ${response.status}`; + try { + const data = (await response.json()) as { message?: string }; + message = data.message || message; + } catch { + // Preserve the status-based fallback for non-JSON responses. + } + + return new FileUploadError( + message, + getFileId(formData), + getToolResource(formData), + true, + response.status, + ); +}; + +const createStreamError = (data: string, formData: FormData) => { + let error: UploadErrorData; + try { + error = JSON.parse(data) as UploadErrorData; + } catch { + error = { message: data }; + } + + return new FileUploadError( + error.message || 'File upload failed.', + error.temp_file_id || getFileId(formData), + error.tool_resource || getToolResource(formData), + error.display_to_user ?? false, + error.code ?? 0, + ); +}; + +const readEventStream = async ( + stream: ReadableStream, + formData: FormData, +): Promise => { + const reader = stream.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let result: TFileUpload | null = null; + let streamEnded = false; + let timeoutError: Error | null = null; + let heartbeatTimer: ReturnType | undefined; + + const resetHeartbeatTimer = () => { + clearTimeout(heartbeatTimer); + heartbeatTimer = setTimeout(() => { + timeoutError = new Error('Upload connection timed out waiting for a heartbeat.'); + void reader.cancel(timeoutError); + }, HEARTBEAT_TIMEOUT_MS); + }; + + resetHeartbeatTimer(); + + try { + while (true) { + const { value, done } = await reader.read(); + if (done) { + streamEnded = true; + if (timeoutError) { + throw timeoutError; + } + if (result) { + return result; + } + throw new Error('Upload connection closed before completion.'); + } + + buffer += decoder.decode(value, { stream: true }); + const messages = buffer.split(/\r?\n\r?\n/); + buffer = messages.pop() ?? ''; + + for (const message of messages) { + const event = parseEvent(message); + if (event.type === 'heartbeat') { + resetHeartbeatTimer(); + continue; + } + if (event.type === 'error') { + throw createStreamError(event.data, formData); + } + if (event.type === 'data') { + result = JSON.parse(event.data) as TFileUpload; + continue; + } + if (event.type === 'close') { + if (result) { + return result; + } + throw new Error('Upload stream closed without a result.'); + } + } + } + } catch (error) { + if (error instanceof Error && error.name === 'AbortError') { + throw new UploadCanceledError('Upload canceled.'); + } + throw error; + } finally { + clearTimeout(heartbeatTimer); + if (!streamEnded) { + await reader.cancel().catch(() => undefined); + } + reader.releaseLock(); + } +}; + +export async function uploadEventStream( + url: string, + formData: FormData, + signal?: AbortSignal | null, +): Promise { + try { + const response = await request.authenticatedFetch(url, { + method: 'POST', + body: formData, + headers: { Accept: EVENT_STREAM_MEDIA_TYPE }, + signal: signal ?? undefined, + }); + if (!response.ok) { + throw await createHttpError(response, formData); + } + + const contentType = response.headers.get('Content-Type')?.toLowerCase() ?? ''; + if (!contentType.includes(EVENT_STREAM_MEDIA_TYPE)) { + return (await response.json()) as TFileUpload; + } + if (!response.body) { + throw new Error('No upload response body received.'); + } + + return await readEventStream(response.body, formData); + } catch (error) { + if (signal?.aborted || (error instanceof Error && error.name === 'AbortError')) { + throw new UploadCanceledError('Upload canceled.'); + } + throw error; + } +}