mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-16 16:30:15 +01:00
✨ feat: Implement Resumable Generation Jobs with SSE Support
- Introduced GenerationJobManager to handle resumable LLM generation jobs independently of HTTP connections. - Added support for subscribing to ongoing generation jobs via SSE, allowing clients to reconnect and receive updates without losing progress. - Enhanced existing agent controllers and routes to integrate resumable functionality, including job creation, completion, and error handling. - Updated client-side hooks to manage adaptive SSE streams, switching between standard and resumable modes based on user settings. - Added UI components and settings for enabling/disabling resumable streams, improving user experience during unstable connections.
This commit is contained in:
Danny Avila
parent
5bfebc7c9d
commit
0e850a5d5f
17 changed files with 1212 additions and 37 deletions
|
|
@ -1,5 +1,5 @@
|
|||
const { nanoid } = require('nanoid');
|
||||
const { sendEvent } = require('@librechat/api');
|
||||
const { sendEvent, GenerationJobManager } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { Tools, StepTypes, FileContext, ErrorTypes } = require('librechat-data-provider');
|
||||
const {
|
||||
|
|
@ -144,17 +144,38 @@ function checkIfLastAgent(last_agent_id, langgraph_node) {
|
|||
return langgraph_node?.endsWith(last_agent_id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper to emit events either to res (standard mode) or to job emitter (resumable mode).
|
||||
* @param {ServerResponse} res - The server response object
|
||||
* @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode
|
||||
* @param {Object} eventData - The event data to send
|
||||
*/
|
||||
function emitEvent(res, streamId, eventData) {
|
||||
if (streamId) {
|
||||
GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
sendEvent(res, eventData);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get default handlers for stream events.
|
||||
* @param {Object} options - The options object.
|
||||
* @param {ServerResponse} options.res - The options object.
|
||||
* @param {ContentAggregator} options.aggregateContent - The options object.
|
||||
* @param {ServerResponse} options.res - The server response object.
|
||||
* @param {ContentAggregator} options.aggregateContent - Content aggregator function.
|
||||
* @param {ToolEndCallback} options.toolEndCallback - Callback to use when tool ends.
|
||||
* @param {Array<UsageMetadata>} options.collectedUsage - The list of collected usage metadata.
|
||||
* @param {string | null} [options.streamId] - The stream ID for resumable mode, or null for standard mode.
|
||||
* @returns {Record<string, t.EventHandler>} The default handlers.
|
||||
* @throws {Error} If the request is not found.
|
||||
*/
|
||||
function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedUsage }) {
|
||||
function getDefaultHandlers({
|
||||
res,
|
||||
aggregateContent,
|
||||
toolEndCallback,
|
||||
collectedUsage,
|
||||
streamId = null,
|
||||
}) {
|
||||
if (!res || !aggregateContent) {
|
||||
throw new Error(
|
||||
`[getDefaultHandlers] Missing required options: res: ${!res}, aggregateContent: ${!aggregateContent}`,
|
||||
|
|
@ -173,16 +194,16 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (data?.stepDetails.type === StepTypes.TOOL_CALLS) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else {
|
||||
const agentName = metadata?.name ?? 'Agent';
|
||||
const isToolCall = data?.stepDetails.type === StepTypes.TOOL_CALLS;
|
||||
const action = isToolCall ? 'performing a task...' : 'thinking...';
|
||||
sendEvent(res, {
|
||||
emitEvent(res, streamId, {
|
||||
event: 'on_agent_update',
|
||||
data: {
|
||||
runId: metadata?.run_id,
|
||||
|
|
@ -202,11 +223,11 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (data?.delta.type === StepTypes.TOOL_CALLS) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
aggregateContent({ event, data });
|
||||
},
|
||||
|
|
@ -220,11 +241,11 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (data?.result != null) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
aggregateContent({ event, data });
|
||||
},
|
||||
|
|
@ -238,9 +259,9 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
aggregateContent({ event, data });
|
||||
},
|
||||
|
|
@ -254,9 +275,9 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
*/
|
||||
handle: (event, data, metadata) => {
|
||||
if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
} else if (!metadata?.hide_sequential_outputs) {
|
||||
sendEvent(res, { event, data });
|
||||
emitEvent(res, streamId, { event, data });
|
||||
}
|
||||
aggregateContent({ event, data });
|
||||
},
|
||||
|
|
@ -266,15 +287,30 @@ function getDefaultHandlers({ res, aggregateContent, toolEndCallback, collectedU
|
|||
return handlers;
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper to write attachment events either to res or to job emitter.
|
||||
* @param {ServerResponse} res - The server response object
|
||||
* @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode
|
||||
* @param {Object} attachment - The attachment data
|
||||
*/
|
||||
function writeAttachment(res, streamId, attachment) {
|
||||
if (streamId) {
|
||||
GenerationJobManager.emitChunk(streamId, { event: 'attachment', data: attachment });
|
||||
} else {
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {Object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {ServerResponse} params.res
|
||||
* @param {Promise<MongoFile | { filename: string; filepath: string; expires: number;} | null>[]} params.artifactPromises
|
||||
* @param {string | null} [params.streamId] - The stream ID for resumable mode, or null for standard mode.
|
||||
* @returns {ToolEndCallback} The tool end callback.
|
||||
*/
|
||||
function createToolEndCallback({ req, res, artifactPromises }) {
|
||||
function createToolEndCallback({ req, res, artifactPromises, streamId = null }) {
|
||||
/**
|
||||
* @type {ToolEndCallback}
|
||||
*/
|
||||
|
|
@ -302,10 +338,10 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
if (!attachment) {
|
||||
return null;
|
||||
}
|
||||
if (!res.headersSent) {
|
||||
if (!streamId && !res.headersSent) {
|
||||
return attachment;
|
||||
}
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`);
|
||||
writeAttachment(res, streamId, attachment);
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing file citations:', error);
|
||||
|
|
@ -314,8 +350,6 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
);
|
||||
}
|
||||
|
||||
// TODO: a lot of duplicated code in createToolEndCallback
|
||||
// we should refactor this to use a helper function in a follow-up PR
|
||||
if (output.artifact[Tools.ui_resources]) {
|
||||
artifactPromises.push(
|
||||
(async () => {
|
||||
|
|
@ -326,10 +360,10 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
conversationId: metadata.thread_id,
|
||||
[Tools.ui_resources]: output.artifact[Tools.ui_resources].data,
|
||||
};
|
||||
if (!res.headersSent) {
|
||||
if (!streamId && !res.headersSent) {
|
||||
return attachment;
|
||||
}
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`);
|
||||
writeAttachment(res, streamId, attachment);
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
|
|
@ -348,10 +382,10 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
conversationId: metadata.thread_id,
|
||||
[Tools.web_search]: { ...output.artifact[Tools.web_search] },
|
||||
};
|
||||
if (!res.headersSent) {
|
||||
if (!streamId && !res.headersSent) {
|
||||
return attachment;
|
||||
}
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`);
|
||||
writeAttachment(res, streamId, attachment);
|
||||
return attachment;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
|
|
@ -388,7 +422,7 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
toolCallId: output.tool_call_id,
|
||||
conversationId: metadata.thread_id,
|
||||
});
|
||||
if (!res.headersSent) {
|
||||
if (!streamId && !res.headersSent) {
|
||||
return fileMetadata;
|
||||
}
|
||||
|
||||
|
|
@ -396,7 +430,7 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
return null;
|
||||
}
|
||||
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(fileMetadata)}\n\n`);
|
||||
writeAttachment(res, streamId, fileMetadata);
|
||||
return fileMetadata;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing artifact content:', error);
|
||||
|
|
@ -435,7 +469,7 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
conversationId: metadata.thread_id,
|
||||
session_id: output.artifact.session_id,
|
||||
});
|
||||
if (!res.headersSent) {
|
||||
if (!streamId && !res.headersSent) {
|
||||
return fileMetadata;
|
||||
}
|
||||
|
||||
|
|
@ -443,7 +477,7 @@ function createToolEndCallback({ req, res, artifactPromises }) {
|
|||
return null;
|
||||
}
|
||||
|
||||
res.write(`event: attachment\ndata: ${JSON.stringify(fileMetadata)}\n\n`);
|
||||
writeAttachment(res, streamId, fileMetadata);
|
||||
return fileMetadata;
|
||||
})().catch((error) => {
|
||||
logger.error('Error processing code output:', error);
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { Constants } = require('librechat-data-provider');
|
||||
const {
|
||||
sendEvent,
|
||||
GenerationJobManager,
|
||||
sanitizeFileForTransmit,
|
||||
sanitizeMessageForTransmit,
|
||||
} = require('@librechat/api');
|
||||
|
|
@ -31,7 +32,232 @@ function createCloseHandler(abortController) {
|
|||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Resumable Agent Controller - Generation runs independently of HTTP connection.
|
||||
* Returns streamId immediately, client subscribes separately via SSE.
|
||||
*/
|
||||
const ResumableAgentController = async (req, res, next, initializeClient, addTitle) => {
|
||||
const {
|
||||
text,
|
||||
isRegenerate,
|
||||
endpointOption,
|
||||
conversationId: reqConversationId,
|
||||
isContinued = false,
|
||||
editedContent = null,
|
||||
parentMessageId = null,
|
||||
overrideParentMessageId = null,
|
||||
responseMessageId: editedResponseMessageId = null,
|
||||
} = req.body;
|
||||
|
||||
const userId = req.user.id;
|
||||
const streamId =
|
||||
reqConversationId || `stream_${Date.now()}_${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
let client = null;
|
||||
|
||||
try {
|
||||
const prelimAbortController = new AbortController();
|
||||
res.on('close', () => {
|
||||
if (!prelimAbortController.signal.aborted) {
|
||||
prelimAbortController.abort();
|
||||
}
|
||||
});
|
||||
|
||||
const job = GenerationJobManager.createJob(streamId, userId, reqConversationId);
|
||||
req._resumableStreamId = streamId;
|
||||
|
||||
/** @type {{ client: TAgentClient; userMCPAuthMap?: Record<string, Record<string, string>> }} */
|
||||
const result = await initializeClient({
|
||||
req,
|
||||
res,
|
||||
endpointOption,
|
||||
signal: prelimAbortController.signal,
|
||||
});
|
||||
|
||||
if (prelimAbortController.signal.aborted) {
|
||||
GenerationJobManager.completeJob(streamId, 'Request aborted during initialization');
|
||||
return res.status(400).json({ error: 'Request aborted during initialization' });
|
||||
}
|
||||
|
||||
client = result.client;
|
||||
|
||||
res.json({ streamId, status: 'started' });
|
||||
|
||||
let conversationId = reqConversationId;
|
||||
let userMessage;
|
||||
|
||||
const getReqData = (data = {}) => {
|
||||
if (data.userMessage) {
|
||||
userMessage = data.userMessage;
|
||||
}
|
||||
if (!conversationId && data.conversationId) {
|
||||
conversationId = data.conversationId;
|
||||
}
|
||||
};
|
||||
|
||||
// Start background generation - wait for subscriber with timeout fallback
|
||||
const startGeneration = async () => {
|
||||
try {
|
||||
await Promise.race([job.readyPromise, new Promise((resolve) => setTimeout(resolve, 5000))]);
|
||||
} catch (waitError) {
|
||||
logger.warn(
|
||||
`[ResumableAgentController] Error waiting for subscriber: ${waitError.message}`,
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
const onStart = (userMsg, _respMsgId, _isNewConvo) => {
|
||||
userMessage = userMsg;
|
||||
|
||||
GenerationJobManager.emitChunk(streamId, {
|
||||
created: true,
|
||||
message: userMessage,
|
||||
streamId,
|
||||
});
|
||||
};
|
||||
|
||||
const messageOptions = {
|
||||
user: userId,
|
||||
onStart,
|
||||
getReqData,
|
||||
isContinued,
|
||||
isRegenerate,
|
||||
editedContent,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
abortController: job.abortController,
|
||||
overrideParentMessageId,
|
||||
isEdited: !!editedContent,
|
||||
userMCPAuthMap: result.userMCPAuthMap,
|
||||
responseMessageId: editedResponseMessageId,
|
||||
progressOptions: {
|
||||
res: {
|
||||
write: () => true,
|
||||
end: () => {},
|
||||
headersSent: false,
|
||||
writableEnded: false,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const response = await client.sendMessage(text, messageOptions);
|
||||
|
||||
const messageId = response.messageId;
|
||||
const endpoint = endpointOption.endpoint;
|
||||
response.endpoint = endpoint;
|
||||
|
||||
const databasePromise = response.databasePromise;
|
||||
delete response.databasePromise;
|
||||
|
||||
const { conversation: convoData = {} } = await databasePromise;
|
||||
const conversation = { ...convoData };
|
||||
conversation.title =
|
||||
conversation && !conversation.title ? null : conversation?.title || 'New Chat';
|
||||
|
||||
if (req.body.files && client.options?.attachments) {
|
||||
userMessage.files = [];
|
||||
const messageFiles = new Set(req.body.files.map((file) => file.file_id));
|
||||
for (const attachment of client.options.attachments) {
|
||||
if (messageFiles.has(attachment.file_id)) {
|
||||
userMessage.files.push(sanitizeFileForTransmit(attachment));
|
||||
}
|
||||
}
|
||||
delete userMessage.image_urls;
|
||||
}
|
||||
|
||||
if (!job.abortController.signal.aborted) {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
conversation,
|
||||
title: conversation.title,
|
||||
requestMessage: sanitizeMessageForTransmit(userMessage),
|
||||
responseMessage: { ...response },
|
||||
};
|
||||
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId);
|
||||
|
||||
if (client.savedMessageIds && !client.savedMessageIds.has(messageId)) {
|
||||
await saveMessage(
|
||||
req,
|
||||
{ ...response, user: userId },
|
||||
{ context: 'api/server/controllers/agents/request.js - resumable response end' },
|
||||
);
|
||||
}
|
||||
} else {
|
||||
const finalEvent = {
|
||||
final: true,
|
||||
conversation,
|
||||
title: conversation.title,
|
||||
requestMessage: sanitizeMessageForTransmit(userMessage),
|
||||
responseMessage: { ...response, error: true },
|
||||
error: { message: 'Request was aborted' },
|
||||
};
|
||||
GenerationJobManager.emitDone(streamId, finalEvent);
|
||||
GenerationJobManager.completeJob(streamId, 'Request aborted');
|
||||
}
|
||||
|
||||
if (!client.skipSaveUserMessage && userMessage) {
|
||||
await saveMessage(req, userMessage, {
|
||||
context: 'api/server/controllers/agents/request.js - resumable user message',
|
||||
});
|
||||
}
|
||||
|
||||
const newConvo = !reqConversationId;
|
||||
if (addTitle && parentMessageId === Constants.NO_PARENT && newConvo) {
|
||||
addTitle(req, {
|
||||
text,
|
||||
response: { ...response },
|
||||
client,
|
||||
})
|
||||
.catch((err) => {
|
||||
logger.error('[ResumableAgentController] Error in title generation', err);
|
||||
})
|
||||
.finally(() => {
|
||||
if (client) {
|
||||
disposeClient(client);
|
||||
}
|
||||
});
|
||||
} else {
|
||||
if (client) {
|
||||
disposeClient(client);
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
logger.error(`[ResumableAgentController] Generation error for ${streamId}:`, error);
|
||||
GenerationJobManager.emitError(streamId, error.message || 'Generation failed');
|
||||
GenerationJobManager.completeJob(streamId, error.message);
|
||||
if (client) {
|
||||
disposeClient(client);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Start generation and handle any unhandled errors
|
||||
startGeneration().catch((err) => {
|
||||
logger.error(
|
||||
`[ResumableAgentController] Unhandled error in background generation: ${err.message}`,
|
||||
);
|
||||
GenerationJobManager.completeJob(streamId, err.message);
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error('[ResumableAgentController] Initialization error:', error);
|
||||
if (!res.headersSent) {
|
||||
res.status(500).json({ error: error.message || 'Failed to start generation' });
|
||||
}
|
||||
GenerationJobManager.completeJob(streamId, error.message);
|
||||
if (client) {
|
||||
disposeClient(client);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const AgentController = async (req, res, next, initializeClient, addTitle) => {
|
||||
const isResumable = req.query.resumable === 'true';
|
||||
if (isResumable) {
|
||||
return ResumableAgentController(req, res, next, initializeClient, addTitle);
|
||||
}
|
||||
|
||||
let {
|
||||
text,
|
||||
isRegenerate,
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ const {
|
|||
performStartupChecks,
|
||||
handleJsonParseError,
|
||||
initializeFileStorage,
|
||||
GenerationJobManager,
|
||||
} = require('@librechat/api');
|
||||
const { connectDb, indexSync } = require('~/db');
|
||||
const initializeOAuthReconnectManager = require('./services/initializeOAuthReconnectManager');
|
||||
|
|
@ -192,6 +193,7 @@ const startServer = async () => {
|
|||
await initializeMCPs();
|
||||
await initializeOAuthReconnectManager();
|
||||
await checkMigrations();
|
||||
GenerationJobManager.initialize();
|
||||
});
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,9 @@
|
|||
function setHeaders(req, res, next) {
|
||||
// Skip SSE headers for resumable mode - it returns JSON first, then client subscribes separately
|
||||
if (req.query.resumable === 'true') {
|
||||
return next();
|
||||
}
|
||||
|
||||
res.writeHead(200, {
|
||||
Connection: 'keep-alive',
|
||||
'Content-Type': 'text/event-stream',
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
const express = require('express');
|
||||
const { generateCheckAccess, skipAgentCheck } = require('@librechat/api');
|
||||
const { generateCheckAccess, skipAgentCheck, GenerationJobManager } = require('@librechat/api');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { PermissionTypes, Permissions, PermissionBits } = require('librechat-data-provider');
|
||||
const {
|
||||
setHeaders,
|
||||
moderateText,
|
||||
requireJwtAuth,
|
||||
// validateModel,
|
||||
validateConvoAccess,
|
||||
buildEndpointOption,
|
||||
|
|
@ -28,6 +30,97 @@ const checkAgentResourceAccess = canAccessAgentFromBody({
|
|||
requiredPermission: PermissionBits.VIEW,
|
||||
});
|
||||
|
||||
/**
|
||||
* @route GET /stream/:streamId
|
||||
* @desc Subscribe to an ongoing generation job's SSE stream
|
||||
* @access Private
|
||||
*/
|
||||
router.get('/stream/:streamId', requireJwtAuth, (req, res) => {
|
||||
const { streamId } = req.params;
|
||||
|
||||
const job = GenerationJobManager.getJob(streamId);
|
||||
if (!job) {
|
||||
return res.status(404).json({
|
||||
error: 'Stream not found',
|
||||
message: 'The generation job does not exist or has expired.',
|
||||
});
|
||||
}
|
||||
|
||||
// Disable compression for SSE
|
||||
res.setHeader('Content-Encoding', 'identity');
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-cache, no-transform');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
res.setHeader('X-Accel-Buffering', 'no');
|
||||
res.flushHeaders();
|
||||
|
||||
logger.debug(`[AgentStream] Client subscribed to ${streamId}`);
|
||||
|
||||
const unsubscribe = GenerationJobManager.subscribe(
|
||||
streamId,
|
||||
(event) => {
|
||||
if (!res.writableEnded) {
|
||||
res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
|
||||
if (typeof res.flush === 'function') {
|
||||
res.flush();
|
||||
}
|
||||
}
|
||||
},
|
||||
(event) => {
|
||||
if (!res.writableEnded) {
|
||||
res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
|
||||
if (typeof res.flush === 'function') {
|
||||
res.flush();
|
||||
}
|
||||
res.end();
|
||||
}
|
||||
},
|
||||
(error) => {
|
||||
if (!res.writableEnded) {
|
||||
res.write(`event: error\ndata: ${JSON.stringify({ error })}\n\n`);
|
||||
if (typeof res.flush === 'function') {
|
||||
res.flush();
|
||||
}
|
||||
res.end();
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
if (!unsubscribe) {
|
||||
return res.status(404).json({ error: 'Failed to subscribe to stream' });
|
||||
}
|
||||
|
||||
if (job.status === 'complete' || job.status === 'error' || job.status === 'aborted') {
|
||||
res.write(`event: message\ndata: ${JSON.stringify({ final: true, status: job.status })}\n\n`);
|
||||
res.end();
|
||||
return;
|
||||
}
|
||||
|
||||
req.on('close', () => {
|
||||
logger.debug(`[AgentStream] Client disconnected from ${streamId}`);
|
||||
unsubscribe();
|
||||
});
|
||||
});
|
||||
|
||||
/**
|
||||
* @route POST /abort
|
||||
* @desc Abort an ongoing generation job
|
||||
* @access Private
|
||||
*/
|
||||
router.post('/abort', (req, res) => {
|
||||
const { streamId, abortKey } = req.body;
|
||||
|
||||
const jobStreamId = streamId || abortKey?.split(':')?.[0];
|
||||
|
||||
if (jobStreamId && GenerationJobManager.hasJob(jobStreamId)) {
|
||||
GenerationJobManager.abortJob(jobStreamId);
|
||||
logger.debug(`[AgentStream] Job aborted: ${jobStreamId}`);
|
||||
return res.json({ success: true, aborted: jobStreamId });
|
||||
}
|
||||
|
||||
res.status(404).json({ error: 'Job not found' });
|
||||
});
|
||||
|
||||
router.use(checkAgentAccess);
|
||||
router.use(checkAgentResourceAccess);
|
||||
router.use(validateConvoAccess);
|
||||
|
|
|
|||
|
|
@ -65,18 +65,21 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
}
|
||||
const appConfig = req.config;
|
||||
|
||||
// TODO: use endpointOption to determine options/modelOptions
|
||||
/** @type {string | null} */
|
||||
const streamId = req._resumableStreamId || null;
|
||||
|
||||
/** @type {Array<UsageMetadata>} */
|
||||
const collectedUsage = [];
|
||||
/** @type {ArtifactPromises} */
|
||||
const artifactPromises = [];
|
||||
const { contentParts, aggregateContent } = createContentAggregator();
|
||||
const toolEndCallback = createToolEndCallback({ req, res, artifactPromises });
|
||||
const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId });
|
||||
const eventHandlers = getDefaultHandlers({
|
||||
res,
|
||||
aggregateContent,
|
||||
toolEndCallback,
|
||||
collectedUsage,
|
||||
streamId,
|
||||
});
|
||||
|
||||
if (!endpointOption.agent) {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import { Constants, buildTree } from 'librechat-data-provider';
|
|||
import type { TMessage } from 'librechat-data-provider';
|
||||
import type { ChatFormValues } from '~/common';
|
||||
import { ChatContext, AddedChatContext, useFileMapContext, ChatFormProvider } from '~/Providers';
|
||||
import { useChatHelpers, useAddedResponse, useSSE } from '~/hooks';
|
||||
import { useChatHelpers, useAddedResponse, useAdaptiveSSE } from '~/hooks';
|
||||
import ConversationStarters from './Input/ConversationStarters';
|
||||
import { useGetMessagesByConvoId } from '~/data-provider';
|
||||
import MessagesView from './Messages/MessagesView';
|
||||
|
|
@ -51,8 +51,8 @@ function ChatView({ index = 0 }: { index?: number }) {
|
|||
const chatHelpers = useChatHelpers(index, conversationId);
|
||||
const addedChatHelpers = useAddedResponse({ rootIndex: index });
|
||||
|
||||
useSSE(rootSubmission, chatHelpers, false);
|
||||
useSSE(addedSubmission, addedChatHelpers, true);
|
||||
useAdaptiveSSE(rootSubmission, chatHelpers, false, index);
|
||||
useAdaptiveSSE(addedSubmission, addedChatHelpers, true, index + 1);
|
||||
|
||||
const methods = useForm<ChatFormValues>({
|
||||
defaultValues: { text: '' },
|
||||
|
|
|
|||
|
|
@ -84,6 +84,13 @@ const toggleSwitchConfigs = [
|
|||
hoverCardText: 'com_nav_info_default_temporary_chat',
|
||||
key: 'defaultTemporaryChat',
|
||||
},
|
||||
{
|
||||
stateAtom: store.resumableStreams,
|
||||
localizationKey: 'com_nav_resumable_streams',
|
||||
switchId: 'resumableStreams',
|
||||
hoverCardText: 'com_nav_info_resumable_streams',
|
||||
key: 'resumableStreams',
|
||||
},
|
||||
];
|
||||
|
||||
function Chat() {
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
export { default as useSSE } from './useSSE';
|
||||
export { default as useResumableSSE } from './useResumableSSE';
|
||||
export { default as useAdaptiveSSE } from './useAdaptiveSSE';
|
||||
export { default as useStepHandler } from './useStepHandler';
|
||||
export { default as useContentHandler } from './useContentHandler';
|
||||
export { default as useAttachmentHandler } from './useAttachmentHandler';
|
||||
|
|
|
|||
43
client/src/hooks/SSE/useAdaptiveSSE.ts
Normal file
43
client/src/hooks/SSE/useAdaptiveSSE.ts
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
import { useRecoilValue } from 'recoil';
|
||||
import type { TSubmission } from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import useSSE from './useSSE';
|
||||
import useResumableSSE from './useResumableSSE';
|
||||
import store from '~/store';
|
||||
|
||||
type ChatHelpers = Pick<
|
||||
EventHandlerParams,
|
||||
| 'setMessages'
|
||||
| 'getMessages'
|
||||
| 'setConversation'
|
||||
| 'setIsSubmitting'
|
||||
| 'newConversation'
|
||||
| 'resetLatestMessage'
|
||||
>;
|
||||
|
||||
/**
|
||||
* Adaptive SSE hook that switches between standard and resumable modes.
|
||||
* Uses Recoil state to determine which mode to use.
|
||||
*
|
||||
* Note: Both hooks are always called to comply with React's Rules of Hooks.
|
||||
* We pass null submission to the inactive one.
|
||||
*/
|
||||
export default function useAdaptiveSSE(
|
||||
submission: TSubmission | null,
|
||||
chatHelpers: ChatHelpers,
|
||||
isAddedRequest = false,
|
||||
runIndex = 0,
|
||||
) {
|
||||
const resumableEnabled = useRecoilValue(store.resumableStreams);
|
||||
|
||||
useSSE(resumableEnabled ? null : submission, chatHelpers, isAddedRequest, runIndex);
|
||||
|
||||
const { streamId } = useResumableSSE(
|
||||
resumableEnabled ? submission : null,
|
||||
chatHelpers,
|
||||
isAddedRequest,
|
||||
runIndex,
|
||||
);
|
||||
|
||||
return { streamId, resumableEnabled };
|
||||
}
|
||||
406
client/src/hooks/SSE/useResumableSSE.ts
Normal file
406
client/src/hooks/SSE/useResumableSSE.ts
Normal file
|
|
@ -0,0 +1,406 @@
|
|||
import { useEffect, useState, useRef, useCallback } from 'react';
|
||||
import { v4 } from 'uuid';
|
||||
import { SSE } from 'sse.js';
|
||||
import { useSetRecoilState } from 'recoil';
|
||||
import {
|
||||
request,
|
||||
Constants,
|
||||
createPayload,
|
||||
LocalStorageKeys,
|
||||
removeNullishValues,
|
||||
} from 'librechat-data-provider';
|
||||
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
|
||||
import type { EventHandlerParams } from './useEventHandlers';
|
||||
import type { TResData } from '~/common';
|
||||
import { useGenTitleMutation, useGetStartupConfig, useGetUserBalance } from '~/data-provider';
|
||||
import { useAuthContext } from '~/hooks/AuthContext';
|
||||
import useEventHandlers from './useEventHandlers';
|
||||
import store from '~/store';
|
||||
|
||||
const clearDraft = (conversationId?: string | null) => {
|
||||
if (conversationId) {
|
||||
localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`);
|
||||
localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${conversationId}`);
|
||||
} else {
|
||||
localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${Constants.NEW_CONVO}`);
|
||||
localStorage.removeItem(`${LocalStorageKeys.FILES_DRAFT}${Constants.NEW_CONVO}`);
|
||||
}
|
||||
};
|
||||
|
||||
type ChatHelpers = Pick<
|
||||
EventHandlerParams,
|
||||
| 'setMessages'
|
||||
| 'getMessages'
|
||||
| 'setConversation'
|
||||
| 'setIsSubmitting'
|
||||
| 'newConversation'
|
||||
| 'resetLatestMessage'
|
||||
>;
|
||||
|
||||
const MAX_RETRIES = 5;
|
||||
|
||||
/**
|
||||
* Hook for resumable SSE streams.
|
||||
* Separates generation start (POST) from stream subscription (GET EventSource).
|
||||
* Supports auto-reconnection with exponential backoff.
|
||||
*/
|
||||
export default function useResumableSSE(
|
||||
submission: TSubmission | null,
|
||||
chatHelpers: ChatHelpers,
|
||||
isAddedRequest = false,
|
||||
runIndex = 0,
|
||||
) {
|
||||
const genTitle = useGenTitleMutation();
|
||||
const setActiveRunId = useSetRecoilState(store.activeRunFamily(runIndex));
|
||||
|
||||
const { token, isAuthenticated } = useAuthContext();
|
||||
const [completed, setCompleted] = useState(new Set());
|
||||
const [streamId, setStreamId] = useState<string | null>(null);
|
||||
const setAbortScroll = useSetRecoilState(store.abortScrollFamily(runIndex));
|
||||
const setShowStopButton = useSetRecoilState(store.showStopButtonByIndex(runIndex));
|
||||
|
||||
const sseRef = useRef<SSE | null>(null);
|
||||
const reconnectAttemptRef = useRef(0);
|
||||
const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
|
||||
const submissionRef = useRef<TSubmission | null>(null);
|
||||
|
||||
const {
|
||||
setMessages,
|
||||
getMessages,
|
||||
setConversation,
|
||||
setIsSubmitting,
|
||||
newConversation,
|
||||
resetLatestMessage,
|
||||
} = chatHelpers;
|
||||
|
||||
const {
|
||||
clearStepMaps,
|
||||
stepHandler,
|
||||
syncHandler,
|
||||
finalHandler,
|
||||
errorHandler,
|
||||
messageHandler,
|
||||
contentHandler,
|
||||
createdHandler,
|
||||
attachmentHandler,
|
||||
abortConversation,
|
||||
} = useEventHandlers({
|
||||
genTitle,
|
||||
setMessages,
|
||||
getMessages,
|
||||
setCompleted,
|
||||
isAddedRequest,
|
||||
setConversation,
|
||||
setIsSubmitting,
|
||||
newConversation,
|
||||
setShowStopButton,
|
||||
resetLatestMessage,
|
||||
});
|
||||
|
||||
const { data: startupConfig } = useGetStartupConfig();
|
||||
const balanceQuery = useGetUserBalance({
|
||||
enabled: !!isAuthenticated && startupConfig?.balance?.enabled,
|
||||
});
|
||||
|
||||
/**
|
||||
* Subscribe to stream via SSE library (supports custom headers)
|
||||
*/
|
||||
const subscribeToStream = useCallback(
|
||||
(currentStreamId: string, currentSubmission: TSubmission) => {
|
||||
let { userMessage } = currentSubmission;
|
||||
let textIndex: number | null = null;
|
||||
|
||||
const url = `/api/agents/chat/stream/${encodeURIComponent(currentStreamId)}`;
|
||||
console.log('[ResumableSSE] Subscribing to stream:', url);
|
||||
|
||||
const sse = new SSE(url, {
|
||||
headers: { Authorization: `Bearer ${token}` },
|
||||
method: 'GET',
|
||||
});
|
||||
sseRef.current = sse;
|
||||
|
||||
sse.addEventListener('open', () => {
|
||||
console.log('[ResumableSSE] Stream connected');
|
||||
setAbortScroll(false);
|
||||
setShowStopButton(true);
|
||||
reconnectAttemptRef.current = 0;
|
||||
});
|
||||
|
||||
sse.addEventListener('message', (e: MessageEvent) => {
|
||||
try {
|
||||
const data = JSON.parse(e.data);
|
||||
|
||||
if (data.final != null) {
|
||||
clearDraft(currentSubmission.conversation?.conversationId);
|
||||
try {
|
||||
finalHandler(data, currentSubmission as EventSubmission);
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error in finalHandler:', error);
|
||||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
}
|
||||
(startupConfig?.balance?.enabled ?? false) && balanceQuery.refetch();
|
||||
sse.close();
|
||||
setStreamId(null);
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.created != null) {
|
||||
const runId = v4();
|
||||
setActiveRunId(runId);
|
||||
userMessage = {
|
||||
...userMessage,
|
||||
...data.message,
|
||||
overrideParentMessageId: userMessage.overrideParentMessageId,
|
||||
};
|
||||
createdHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.event === 'attachment' && data.data) {
|
||||
attachmentHandler({
|
||||
data: data.data,
|
||||
submission: currentSubmission as EventSubmission,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.event != null) {
|
||||
stepHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.sync != null) {
|
||||
const runId = v4();
|
||||
setActiveRunId(runId);
|
||||
syncHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.type != null) {
|
||||
const { text, index } = data;
|
||||
if (text != null && index !== textIndex) {
|
||||
textIndex = index;
|
||||
}
|
||||
contentHandler({ data, submission: currentSubmission as EventSubmission });
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.message != null) {
|
||||
const text = data.text ?? data.response;
|
||||
const initialResponse = {
|
||||
...(currentSubmission.initialResponse as TMessage),
|
||||
parentMessageId: data.parentMessageId,
|
||||
messageId: data.messageId,
|
||||
};
|
||||
messageHandler(text, { ...currentSubmission, userMessage, initialResponse });
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error processing message:', error);
|
||||
}
|
||||
});
|
||||
|
||||
// Handle cancel event (triggered when stop button is clicked)
|
||||
sse.addEventListener('cancel', async () => {
|
||||
console.log('[ResumableSSE] Cancel requested, aborting job');
|
||||
sse.close();
|
||||
|
||||
// Call abort endpoint to stop backend generation
|
||||
try {
|
||||
await fetch('/api/agents/chat/abort', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${token}`,
|
||||
},
|
||||
body: JSON.stringify({ streamId: currentStreamId }),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error aborting job:', error);
|
||||
}
|
||||
|
||||
// Handle UI cleanup via abortConversation
|
||||
const latestMessages = getMessages();
|
||||
const conversationId = latestMessages?.[latestMessages.length - 1]?.conversationId;
|
||||
try {
|
||||
await abortConversation(
|
||||
conversationId ??
|
||||
userMessage.conversationId ??
|
||||
currentSubmission.conversation?.conversationId ??
|
||||
'',
|
||||
currentSubmission as EventSubmission,
|
||||
latestMessages,
|
||||
);
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error during abort:', error);
|
||||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
}
|
||||
setStreamId(null);
|
||||
});
|
||||
|
||||
sse.addEventListener('error', async (e: MessageEvent) => {
|
||||
console.log('[ResumableSSE] Stream error, connection closed');
|
||||
sse.close();
|
||||
|
||||
// Check for 401 and try to refresh token
|
||||
/* @ts-ignore */
|
||||
if (e.responseCode === 401) {
|
||||
try {
|
||||
const refreshResponse = await request.refreshToken();
|
||||
const newToken = refreshResponse?.token ?? '';
|
||||
if (newToken) {
|
||||
request.dispatchTokenUpdatedEvent(newToken);
|
||||
// Retry with new token
|
||||
if (submissionRef.current) {
|
||||
subscribeToStream(currentStreamId, submissionRef.current);
|
||||
}
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
console.log('[ResumableSSE] Token refresh failed:', error);
|
||||
}
|
||||
}
|
||||
|
||||
if (reconnectAttemptRef.current < MAX_RETRIES) {
|
||||
reconnectAttemptRef.current++;
|
||||
const delay = Math.min(1000 * Math.pow(2, reconnectAttemptRef.current - 1), 30000);
|
||||
|
||||
console.log(
|
||||
`[ResumableSSE] Reconnecting in ${delay}ms (attempt ${reconnectAttemptRef.current}/${MAX_RETRIES})`,
|
||||
);
|
||||
|
||||
reconnectTimeoutRef.current = setTimeout(() => {
|
||||
if (submissionRef.current) {
|
||||
subscribeToStream(currentStreamId, submissionRef.current);
|
||||
}
|
||||
}, delay);
|
||||
} else {
|
||||
console.error('[ResumableSSE] Max reconnect attempts reached');
|
||||
errorHandler({ data: undefined, submission: currentSubmission as EventSubmission });
|
||||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
setStreamId(null);
|
||||
}
|
||||
});
|
||||
|
||||
// Start the SSE connection
|
||||
sse.stream();
|
||||
},
|
||||
[
|
||||
token,
|
||||
setAbortScroll,
|
||||
setActiveRunId,
|
||||
setShowStopButton,
|
||||
finalHandler,
|
||||
createdHandler,
|
||||
attachmentHandler,
|
||||
stepHandler,
|
||||
syncHandler,
|
||||
contentHandler,
|
||||
messageHandler,
|
||||
errorHandler,
|
||||
setIsSubmitting,
|
||||
startupConfig?.balance?.enabled,
|
||||
balanceQuery,
|
||||
abortConversation,
|
||||
getMessages,
|
||||
],
|
||||
);
|
||||
|
||||
/**
|
||||
* Start generation (POST request that returns streamId)
|
||||
*/
|
||||
const startGeneration = useCallback(
|
||||
async (currentSubmission: TSubmission): Promise<string | null> => {
|
||||
const payloadData = createPayload(currentSubmission);
|
||||
let { payload } = payloadData;
|
||||
payload = removeNullishValues(payload) as TPayload;
|
||||
|
||||
clearStepMaps();
|
||||
|
||||
const url = payloadData.server.includes('?')
|
||||
? `${payloadData.server}&resumable=true`
|
||||
: `${payloadData.server}?resumable=true`;
|
||||
|
||||
try {
|
||||
const response = await fetch(url, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
Authorization: `Bearer ${token}`,
|
||||
},
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json().catch(() => ({}));
|
||||
throw new Error(errorData.error || `Failed to start generation: ${response.statusText}`);
|
||||
}
|
||||
|
||||
const { streamId: newStreamId } = await response.json();
|
||||
console.log('[ResumableSSE] Generation started:', { streamId: newStreamId });
|
||||
|
||||
return newStreamId;
|
||||
} catch (error) {
|
||||
console.error('[ResumableSSE] Error starting generation:', error);
|
||||
errorHandler({ data: undefined, submission: currentSubmission as EventSubmission });
|
||||
setIsSubmitting(false);
|
||||
return null;
|
||||
}
|
||||
},
|
||||
[token, clearStepMaps, errorHandler, setIsSubmitting],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (!submission || Object.keys(submission).length === 0) {
|
||||
if (reconnectTimeoutRef.current) {
|
||||
clearTimeout(reconnectTimeoutRef.current);
|
||||
reconnectTimeoutRef.current = null;
|
||||
}
|
||||
if (sseRef.current) {
|
||||
sseRef.current.close();
|
||||
sseRef.current = null;
|
||||
}
|
||||
setStreamId(null);
|
||||
reconnectAttemptRef.current = 0;
|
||||
submissionRef.current = null;
|
||||
return;
|
||||
}
|
||||
|
||||
submissionRef.current = submission;
|
||||
|
||||
const initStream = async () => {
|
||||
setIsSubmitting(true);
|
||||
|
||||
const newStreamId = await startGeneration(submission);
|
||||
if (newStreamId) {
|
||||
setStreamId(newStreamId);
|
||||
subscribeToStream(newStreamId, submission);
|
||||
}
|
||||
};
|
||||
|
||||
initStream();
|
||||
|
||||
return () => {
|
||||
if (reconnectTimeoutRef.current) {
|
||||
clearTimeout(reconnectTimeoutRef.current);
|
||||
reconnectTimeoutRef.current = null;
|
||||
}
|
||||
if (sseRef.current) {
|
||||
const isCancelled = sseRef.current.readyState <= 1;
|
||||
sseRef.current.close();
|
||||
if (isCancelled) {
|
||||
// Dispatch cancel event to trigger abort
|
||||
const e = new Event('cancel');
|
||||
/* @ts-ignore */
|
||||
sseRef.current.dispatchEvent(e);
|
||||
}
|
||||
sseRef.current = null;
|
||||
}
|
||||
};
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [submission]);
|
||||
|
||||
return { streamId };
|
||||
}
|
||||
|
|
@ -490,6 +490,7 @@
|
|||
"com_nav_info_save_draft": "When enabled, the text and attachments you enter in the chat form will be automatically saved locally as drafts. These drafts will be available even if you reload the page or switch to a different conversation. Drafts are stored locally on your device and are deleted once the message is sent.",
|
||||
"com_nav_info_show_thinking": "When enabled, the chat will display the thinking dropdowns open by default, allowing you to view the AI's reasoning in real-time. When disabled, the thinking dropdowns will remain closed by default for a cleaner and more streamlined interface",
|
||||
"com_nav_info_user_name_display": "When enabled, the username of the sender will be shown above each message you send. When disabled, you will only see \"You\" above your messages.",
|
||||
"com_nav_info_resumable_streams": "When enabled, LLM generation continues in the background even if your connection drops. You can reconnect and resume receiving the response without losing progress. This is useful for unstable connections or long responses.",
|
||||
"com_nav_keep_screen_awake": "Keep screen awake during response generation",
|
||||
"com_nav_lang_arabic": "العربية",
|
||||
"com_nav_lang_armenian": "Հայերեն",
|
||||
|
|
@ -548,6 +549,7 @@
|
|||
"com_nav_plus_command": "+-Command",
|
||||
"com_nav_plus_command_description": "Toggle command \"+\" for adding a multi-response setting",
|
||||
"com_nav_profile_picture": "Profile Picture",
|
||||
"com_nav_resumable_streams": "Resumable Streams (Beta)",
|
||||
"com_nav_save_badges_state": "Save badges state",
|
||||
"com_nav_save_drafts": "Save drafts locally",
|
||||
"com_nav_scroll_button": "Scroll to the end button",
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ const localStorageAtoms = {
|
|||
LaTeXParsing: atomWithLocalStorage('LaTeXParsing', true),
|
||||
centerFormOnLanding: atomWithLocalStorage('centerFormOnLanding', true),
|
||||
showFooter: atomWithLocalStorage('showFooter', true),
|
||||
resumableStreams: atomWithLocalStorage('resumableStreams', true),
|
||||
|
||||
// Commands settings
|
||||
atCommand: atomWithLocalStorage('atCommand', true),
|
||||
|
|
|
|||
|
|
@ -38,6 +38,8 @@ export * from './tools';
|
|||
export * from './web';
|
||||
/* Cache */
|
||||
export * from './cache';
|
||||
/* Stream */
|
||||
export * from './stream';
|
||||
/* types */
|
||||
export type * from './mcp/types';
|
||||
export type * from './flow/types';
|
||||
|
|
|
|||
320
packages/api/src/stream/GenerationJobManager.ts
Normal file
320
packages/api/src/stream/GenerationJobManager.ts
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
import { EventEmitter } from 'events';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import type { ServerSentEvent } from '~/types';
|
||||
import type {
|
||||
GenerationJob,
|
||||
GenerationJobStatus,
|
||||
ChunkHandler,
|
||||
DoneHandler,
|
||||
ErrorHandler,
|
||||
UnsubscribeFn,
|
||||
} from './types';
|
||||
|
||||
/**
|
||||
* Manages generation jobs for resumable LLM streams.
|
||||
* Generation runs independently of HTTP connections via EventEmitter.
|
||||
* Clients can subscribe/unsubscribe to job events without affecting generation.
|
||||
*/
|
||||
class GenerationJobManagerClass {
|
||||
private jobs = new Map<string, GenerationJob>();
|
||||
private cleanupInterval: NodeJS.Timeout | null = null;
|
||||
/** Time to keep completed jobs before cleanup (1 hour) */
|
||||
private ttlAfterComplete = 3600000;
|
||||
/** Maximum number of concurrent jobs */
|
||||
private maxJobs = 1000;
|
||||
|
||||
/**
|
||||
* Initialize the job manager with periodic cleanup.
|
||||
*/
|
||||
initialize(): void {
|
||||
if (this.cleanupInterval) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.cleanupInterval = setInterval(() => {
|
||||
this.cleanup();
|
||||
}, 60000);
|
||||
|
||||
if (this.cleanupInterval.unref) {
|
||||
this.cleanupInterval.unref();
|
||||
}
|
||||
|
||||
logger.debug('[GenerationJobManager] Initialized with cleanup interval');
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a new generation job.
|
||||
* @param streamId - Unique identifier for the stream
|
||||
* @param userId - User ID who initiated the generation
|
||||
* @param conversationId - Optional conversation ID
|
||||
* @returns The created job
|
||||
*/
|
||||
createJob(streamId: string, userId: string, conversationId?: string): GenerationJob {
|
||||
if (this.jobs.size >= this.maxJobs) {
|
||||
this.evictOldest();
|
||||
}
|
||||
|
||||
let resolveReady: () => void;
|
||||
const readyPromise = new Promise<void>((resolve) => {
|
||||
resolveReady = resolve;
|
||||
});
|
||||
|
||||
const job: GenerationJob = {
|
||||
streamId,
|
||||
emitter: new EventEmitter(),
|
||||
status: 'running',
|
||||
createdAt: Date.now(),
|
||||
abortController: new AbortController(),
|
||||
metadata: { userId, conversationId },
|
||||
readyPromise,
|
||||
resolveReady: resolveReady!,
|
||||
};
|
||||
|
||||
job.emitter.setMaxListeners(100);
|
||||
|
||||
this.jobs.set(streamId, job);
|
||||
logger.debug(`[GenerationJobManager] Created job: ${streamId}`);
|
||||
|
||||
return job;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get a job by streamId.
|
||||
* @param streamId - The stream identifier
|
||||
* @returns The job if found, undefined otherwise
|
||||
*/
|
||||
getJob(streamId: string): GenerationJob | undefined {
|
||||
return this.jobs.get(streamId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a job exists.
|
||||
* @param streamId - The stream identifier
|
||||
* @returns True if job exists
|
||||
*/
|
||||
hasJob(streamId: string): boolean {
|
||||
return this.jobs.has(streamId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get job status.
|
||||
* @param streamId - The stream identifier
|
||||
* @returns The job status or undefined if not found
|
||||
*/
|
||||
getJobStatus(streamId: string): GenerationJobStatus | undefined {
|
||||
return this.jobs.get(streamId)?.status;
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark job as complete.
|
||||
* @param streamId - The stream identifier
|
||||
* @param error - Optional error message if job failed
|
||||
*/
|
||||
completeJob(streamId: string, error?: string): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job) {
|
||||
return;
|
||||
}
|
||||
|
||||
job.status = error ? 'error' : 'complete';
|
||||
job.completedAt = Date.now();
|
||||
if (error) {
|
||||
job.error = error;
|
||||
}
|
||||
|
||||
logger.debug(`[GenerationJobManager] Job completed: ${streamId}, status: ${job.status}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Abort a job (user-initiated).
|
||||
* @param streamId - The stream identifier
|
||||
*/
|
||||
abortJob(streamId: string): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job) {
|
||||
return;
|
||||
}
|
||||
|
||||
job.abortController.abort();
|
||||
job.status = 'aborted';
|
||||
job.completedAt = Date.now();
|
||||
job.emitter.emit('error', 'Request aborted by user');
|
||||
|
||||
logger.debug(`[GenerationJobManager] Job aborted: ${streamId}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribe to a job's event stream.
|
||||
* @param streamId - The stream identifier
|
||||
* @param onChunk - Handler for chunk events
|
||||
* @param onDone - Optional handler for completion
|
||||
* @param onError - Optional handler for errors
|
||||
* @returns Unsubscribe function, or null if job not found
|
||||
*/
|
||||
subscribe(
|
||||
streamId: string,
|
||||
onChunk: ChunkHandler,
|
||||
onDone?: DoneHandler,
|
||||
onError?: ErrorHandler,
|
||||
): UnsubscribeFn | null {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const chunkHandler = (event: ServerSentEvent) => onChunk(event);
|
||||
const doneHandler = (event: ServerSentEvent) => onDone?.(event);
|
||||
const errorHandler = (error: string) => onError?.(error);
|
||||
|
||||
job.emitter.on('chunk', chunkHandler);
|
||||
job.emitter.on('done', doneHandler);
|
||||
job.emitter.on('error', errorHandler);
|
||||
|
||||
// Signal that we're ready to receive events (first subscriber)
|
||||
if (job.emitter.listenerCount('chunk') === 1) {
|
||||
job.resolveReady();
|
||||
logger.debug(`[GenerationJobManager] First subscriber ready for ${streamId}`);
|
||||
}
|
||||
|
||||
return () => {
|
||||
const currentJob = this.jobs.get(streamId);
|
||||
if (currentJob) {
|
||||
currentJob.emitter.off('chunk', chunkHandler);
|
||||
currentJob.emitter.off('done', doneHandler);
|
||||
currentJob.emitter.off('error', errorHandler);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Emit a chunk event to all subscribers.
|
||||
* @param streamId - The stream identifier
|
||||
* @param event - The event data to emit
|
||||
*/
|
||||
emitChunk(streamId: string, event: ServerSentEvent): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job || job.status !== 'running') {
|
||||
return;
|
||||
}
|
||||
job.emitter.emit('chunk', event);
|
||||
}
|
||||
|
||||
/**
|
||||
* Emit a done event to all subscribers.
|
||||
* @param streamId - The stream identifier
|
||||
* @param event - The final event data
|
||||
*/
|
||||
emitDone(streamId: string, event: ServerSentEvent): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job) {
|
||||
return;
|
||||
}
|
||||
job.emitter.emit('done', event);
|
||||
}
|
||||
|
||||
/**
|
||||
* Emit an error event to all subscribers.
|
||||
* @param streamId - The stream identifier
|
||||
* @param error - The error message
|
||||
*/
|
||||
emitError(streamId: string, error: string): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (!job) {
|
||||
return;
|
||||
}
|
||||
job.emitter.emit('error', error);
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleanup completed jobs after TTL.
|
||||
*/
|
||||
private cleanup(): void {
|
||||
const now = Date.now();
|
||||
const toDelete: string[] = [];
|
||||
|
||||
for (const [streamId, job] of this.jobs) {
|
||||
const isFinished = ['complete', 'error', 'aborted'].includes(job.status);
|
||||
if (isFinished && job.completedAt && now - job.completedAt > this.ttlAfterComplete) {
|
||||
toDelete.push(streamId);
|
||||
}
|
||||
}
|
||||
|
||||
toDelete.forEach((id) => this.deleteJob(id));
|
||||
|
||||
if (toDelete.length > 0) {
|
||||
logger.debug(`[GenerationJobManager] Cleaned up ${toDelete.length} expired jobs`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete a job and cleanup listeners.
|
||||
* @param streamId - The stream identifier
|
||||
*/
|
||||
private deleteJob(streamId: string): void {
|
||||
const job = this.jobs.get(streamId);
|
||||
if (job) {
|
||||
job.emitter.removeAllListeners();
|
||||
this.jobs.delete(streamId);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Evict oldest job (LRU).
|
||||
*/
|
||||
private evictOldest(): void {
|
||||
let oldestId: string | null = null;
|
||||
let oldestTime = Infinity;
|
||||
|
||||
for (const [streamId, job] of this.jobs) {
|
||||
if (job.createdAt < oldestTime) {
|
||||
oldestTime = job.createdAt;
|
||||
oldestId = streamId;
|
||||
}
|
||||
}
|
||||
|
||||
if (oldestId) {
|
||||
logger.warn(`[GenerationJobManager] Evicting oldest job: ${oldestId}`);
|
||||
this.deleteJob(oldestId);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get total number of active jobs.
|
||||
*/
|
||||
getJobCount(): number {
|
||||
return this.jobs.size;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get count of jobs by status.
|
||||
*/
|
||||
getJobCountByStatus(): Record<GenerationJobStatus, number> {
|
||||
const counts: Record<GenerationJobStatus, number> = {
|
||||
running: 0,
|
||||
complete: 0,
|
||||
error: 0,
|
||||
aborted: 0,
|
||||
};
|
||||
|
||||
for (const job of this.jobs.values()) {
|
||||
counts[job.status]++;
|
||||
}
|
||||
|
||||
return counts;
|
||||
}
|
||||
|
||||
/**
|
||||
* Destroy the manager and cleanup all jobs.
|
||||
*/
|
||||
destroy(): void {
|
||||
if (this.cleanupInterval) {
|
||||
clearInterval(this.cleanupInterval);
|
||||
this.cleanupInterval = null;
|
||||
}
|
||||
this.jobs.forEach((_, streamId) => this.deleteJob(streamId));
|
||||
logger.debug('[GenerationJobManager] Destroyed');
|
||||
}
|
||||
}
|
||||
|
||||
export const GenerationJobManager = new GenerationJobManagerClass();
|
||||
export { GenerationJobManagerClass };
|
||||
2
packages/api/src/stream/index.ts
Normal file
2
packages/api/src/stream/index.ts
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
export { GenerationJobManager, GenerationJobManagerClass } from './GenerationJobManager';
|
||||
export type * from './types';
|
||||
27
packages/api/src/stream/types.ts
Normal file
27
packages/api/src/stream/types.ts
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
import type { EventEmitter } from 'events';
|
||||
import type { ServerSentEvent } from '~/types';
|
||||
|
||||
export interface GenerationJobMetadata {
|
||||
userId: string;
|
||||
conversationId?: string;
|
||||
}
|
||||
|
||||
export type GenerationJobStatus = 'running' | 'complete' | 'error' | 'aborted';
|
||||
|
||||
export interface GenerationJob {
|
||||
streamId: string;
|
||||
emitter: EventEmitter;
|
||||
status: GenerationJobStatus;
|
||||
createdAt: number;
|
||||
completedAt?: number;
|
||||
abortController: AbortController;
|
||||
error?: string;
|
||||
metadata: GenerationJobMetadata;
|
||||
readyPromise: Promise<void>;
|
||||
resolveReady: () => void;
|
||||
}
|
||||
|
||||
export type ChunkHandler = (event: ServerSentEvent) => void;
|
||||
export type DoneHandler = (event: ServerSentEvent) => void;
|
||||
export type ErrorHandler = (error: string) => void;
|
||||
export type UnsubscribeFn = () => void;
|
||||
Loading…
Add table
Add a link
Reference in a new issue