🚦 feat: Auto-reinitialize MCP Servers on Request (#9226)

This commit is contained in:
Danny Avila 2025-08-23 03:27:05 -04:00
parent ac608ded46
commit c827fdd10e
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
28 changed files with 871 additions and 312 deletions

View file

@ -230,15 +230,19 @@ export default function useChatFunctions({
const responseMessageId =
editedMessageId ??
(latestMessage?.messageId && isRegenerate ? latestMessage?.messageId + '_' : null) ??
(latestMessage?.messageId && isRegenerate
? latestMessage.messageId.replace(/_+$/, '') + '_'
: null) ??
null;
const initialResponseId =
responseMessageId ?? `${isRegenerate ? messageId : intermediateId}`.replace(/_+$/, '') + '_';
const initialResponse: TMessage = {
sender: responseSender,
text: '',
endpoint: endpoint ?? '',
parentMessageId: isRegenerate ? messageId : intermediateId,
messageId: responseMessageId ?? `${isRegenerate ? messageId : intermediateId}_`,
messageId: initialResponseId,
thread_id,
conversationId,
unfinished: false,

View file

@ -182,7 +182,7 @@ export default function useEventHandlers({
const { token } = useAuthContext();
const contentHandler = useContentHandler({ setMessages, getMessages });
const stepHandler = useStepHandler({
const { stepHandler, clearStepMaps } = useStepHandler({
setMessages,
getMessages,
announcePolite,
@ -806,6 +806,7 @@ export default function useEventHandlers({
);
return {
clearStepMaps,
stepHandler,
syncHandler,
finalHandler,

View file

@ -62,6 +62,7 @@ export default function useSSE(
} = chatHelpers;
const {
clearStepMaps,
stepHandler,
syncHandler,
finalHandler,
@ -101,6 +102,7 @@ export default function useSSE(
payload = removeNullishValues(payload) as TPayload;
let textIndex = null;
clearStepMaps();
const sse = new SSE(payloadData.server, {
payload: JSON.stringify(payload),

View file

@ -1,5 +1,11 @@
import { useCallback, useRef } from 'react';
import { StepTypes, ContentTypes, ToolCallTypes, getNonEmptyValue } from 'librechat-data-provider';
import {
Constants,
StepTypes,
ContentTypes,
ToolCallTypes,
getNonEmptyValue,
} from 'librechat-data-provider';
import type {
Agents,
TMessage,
@ -178,11 +184,12 @@ export default function useStepHandler({
return { ...message, content: updatedContent as TMessageContentParts[] };
};
return useCallback(
const stepHandler = useCallback(
({ event, data }: TStepEvent, submission: EventSubmission) => {
const messages = getMessages() || [];
const { userMessage } = submission;
setIsSubmitting(true);
let parentMessageId = userMessage.messageId;
const currentTime = Date.now();
if (currentTime - lastAnnouncementTimeRef.current > MESSAGE_UPDATE_INTERVAL) {
@ -197,7 +204,11 @@ export default function useStepHandler({
if (event === 'on_run_step') {
const runStep = data as Agents.RunStep;
const responseMessageId = runStep.runId ?? '';
let responseMessageId = runStep.runId ?? '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!responseMessageId) {
console.warn('No message id found in run step event');
return;
@ -211,7 +222,7 @@ export default function useStepHandler({
response = {
...responseMessage,
parentMessageId: userMessage.messageId,
parentMessageId,
conversationId: userMessage.conversationId,
messageId: responseMessageId,
content: initialContent,
@ -246,14 +257,18 @@ export default function useStepHandler({
messageMap.current.set(responseMessageId, updatedResponse);
const updatedMessages = messages.map((msg) =>
msg.messageId === runStep.runId ? updatedResponse : msg,
msg.messageId === responseMessageId ? updatedResponse : msg,
);
setMessages(updatedMessages);
}
} else if (event === 'on_agent_update') {
const { agent_update } = data as Agents.AgentUpdate;
const responseMessageId = agent_update.runId || '';
let responseMessageId = agent_update.runId || '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!responseMessageId) {
console.warn('No message id found in agent update event');
return;
@ -271,7 +286,11 @@ export default function useStepHandler({
} else if (event === 'on_message_delta') {
const messageDelta = data as Agents.MessageDeltaEvent;
const runStep = stepMap.current.get(messageDelta.id);
const responseMessageId = runStep?.runId ?? '';
let responseMessageId = runStep?.runId ?? '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!runStep || !responseMessageId) {
console.warn('No run step or runId found for message delta event');
@ -299,7 +318,11 @@ export default function useStepHandler({
} else if (event === 'on_reasoning_delta') {
const reasoningDelta = data as Agents.ReasoningDeltaEvent;
const runStep = stepMap.current.get(reasoningDelta.id);
const responseMessageId = runStep?.runId ?? '';
let responseMessageId = runStep?.runId ?? '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!runStep || !responseMessageId) {
console.warn('No run step or runId found for reasoning delta event');
@ -327,7 +350,11 @@ export default function useStepHandler({
} else if (event === 'on_run_step_delta') {
const runStepDelta = data as Agents.RunStepDeltaEvent;
const runStep = stepMap.current.get(runStepDelta.id);
const responseMessageId = runStep?.runId ?? '';
let responseMessageId = runStep?.runId ?? '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!runStep || !responseMessageId) {
console.warn('No run step or runId found for run step delta event');
@ -366,7 +393,7 @@ export default function useStepHandler({
messageMap.current.set(responseMessageId, updatedResponse);
const updatedMessages = messages.map((msg) =>
msg.messageId === runStep.runId ? updatedResponse : msg,
msg.messageId === responseMessageId ? updatedResponse : msg,
);
setMessages(updatedMessages);
@ -377,7 +404,11 @@ export default function useStepHandler({
const { id: stepId } = result;
const runStep = stepMap.current.get(stepId);
const responseMessageId = runStep?.runId ?? '';
let responseMessageId = runStep?.runId ?? '';
if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) {
responseMessageId = submission?.initialResponse?.messageId ?? '';
parentMessageId = submission?.initialResponse?.parentMessageId ?? '';
}
if (!runStep || !responseMessageId) {
console.warn('No run step or runId found for completed tool call event');
@ -399,7 +430,7 @@ export default function useStepHandler({
messageMap.current.set(responseMessageId, updatedResponse);
const updatedMessages = messages.map((msg) =>
msg.messageId === runStep.runId ? updatedResponse : msg,
msg.messageId === responseMessageId ? updatedResponse : msg,
);
setMessages(updatedMessages);
@ -414,4 +445,11 @@ export default function useStepHandler({
},
[getMessages, setIsSubmitting, lastAnnouncementTimeRef, announcePolite, setMessages],
);
const clearStepMaps = useCallback(() => {
toolCallIdMap.current.clear();
messageMap.current.clear();
stepMap.current.clear();
}, []);
return { stepHandler, clearStepMaps };
}