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; + } +}