mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 06:52:47 +00:00
🛟 fix: Keep File Uploads Alive With SSE Heartbeats (#14295)
Some checks are pending
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Waiting to run
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Waiting to run
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Blocked by required conditions
Sync Helm Chart Tags / Ignore non-main push (push) Waiting to run
Sync Helm Chart Tags / Sync chart tags (push) Waiting to run
Some checks are pending
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Waiting to run
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Waiting to run
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Blocked by required conditions
Sync Helm Chart Tags / Ignore non-main push (push) Waiting to run
Sync Helm Chart Tags / Sync chart tags (push) Waiting to run
* 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 <MarcAmick@jhu.edu> Co-authored-by: Danny Avila <danny@librechat.ai>
This commit is contained in:
parent
7406f5d79e
commit
ade02054c8
17 changed files with 919 additions and 37 deletions
|
|
@ -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 #
|
||||
#===================================================#
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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<void>}
|
||||
*/
|
||||
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<void>}
|
||||
*/
|
||||
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<void>}
|
||||
*/
|
||||
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);
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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 }) => {
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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';
|
||||
|
|
|
|||
218
packages/api/src/files/sse.spec.ts
Normal file
218
packages/api/src/files/sse.spec.ts
Normal file
|
|
@ -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<Response> => {
|
||||
const res = new EventEmitter() as unknown as jest.Mocked<Response>;
|
||||
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<Response>): 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' });
|
||||
});
|
||||
});
|
||||
});
|
||||
100
packages/api/src/files/sse.ts
Normal file
100
packages/api/src/files/sse.ts
Normal file
|
|
@ -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<T>(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: <T>(data: T) => void;
|
||||
/** Emits a failure payload as an `event:error` message. */
|
||||
sendError: <T>(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<T extends object>(
|
||||
res: Response,
|
||||
sseStream: UploadSseStream | null | undefined,
|
||||
message: string,
|
||||
result: T,
|
||||
): void {
|
||||
if (sseStream) {
|
||||
sseStream.sendData({ message, ...result });
|
||||
return;
|
||||
}
|
||||
res.status(200).json({ message, ...result });
|
||||
}
|
||||
|
|
@ -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));
|
||||
|
|
|
|||
|
|
@ -1578,6 +1578,7 @@ export type TStartupConfig = {
|
|||
branch?: string | null;
|
||||
buildDate?: string | null;
|
||||
};
|
||||
fileUploadSseEnabled?: boolean;
|
||||
};
|
||||
|
||||
export type TSharedLinkStartupInterface = Pick<
|
||||
|
|
|
|||
|
|
@ -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<TFileConfig> => {
|
|||
export const uploadImage = (
|
||||
data: FormData,
|
||||
signal?: AbortSignal | null,
|
||||
sseEnabled = false,
|
||||
): Promise<f.TFileUpload> => {
|
||||
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<f.TFileUpload> => {
|
||||
export const uploadFile = (
|
||||
data: FormData,
|
||||
signal?: AbortSignal | null,
|
||||
sseEnabled = false,
|
||||
): Promise<f.TFileUpload> => {
|
||||
const requestConfig = signal ? { signal } : undefined;
|
||||
if (sseEnabled) {
|
||||
return uploadEventStream(endpoints.files(), data, signal);
|
||||
}
|
||||
return request.postMultiPart(endpoints.files(), data, requestConfig);
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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<Response> {
|
||||
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,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
152
packages/data-provider/src/upload.spec.ts
Normal file
152
packages/data-provider/src/upload.spec.ts
Normal file
|
|
@ -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<Uint8Array>({
|
||||
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<Uint8Array>(), {
|
||||
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,
|
||||
});
|
||||
});
|
||||
});
|
||||
212
packages/data-provider/src/upload.ts
Normal file
212
packages/data-provider/src/upload.ts
Normal file
|
|
@ -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<Uint8Array>,
|
||||
formData: FormData,
|
||||
): Promise<TFileUpload> => {
|
||||
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<typeof setTimeout> | 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<TFileUpload> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue