mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-28 12:44:28 +00:00
🔐 fix: Handle Multiple Concurrent MCP OAuth Login Prompts (#13200)
* fix: handle multiple MCP OAuth prompts * fix: address MCP OAuth review feedback * fix: address MCP OAuth prompt lifecycle review * fix: narrow OAuth prompt slot cleanup * fix: format OAuth prompt test --------- Co-authored-by: Danny Avila <danny@librechat.ai>
This commit is contained in:
janluedemann-esome
GitHub
parent
aeb5adff34
commit
2ed59ac98a
5 changed files with 436 additions and 32 deletions
|
|
@ -234,6 +234,182 @@ describe('useStepHandler', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('should preserve multiple tool call steps for the same preliminary response', () => {
|
||||
const responseMessage = createResponseMessage();
|
||||
let currentMessages = [responseMessage];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
const submission = createSubmission({ initialResponse: responseMessage });
|
||||
|
||||
const firstRunStep = createToolCallRunStep({
|
||||
id: 'step-oauth-eli',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-eli',
|
||||
name: `oauth${Constants.mcp_delimiter}ELI`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
const secondRunStep = createToolCallRunStep({
|
||||
id: 'step-oauth-vespa',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 1,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-vespa',
|
||||
name: `oauth${Constants.mcp_delimiter}Vespa`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler(
|
||||
{ event: StepEvents.ON_RUN_STEP, data: firstRunStep },
|
||||
submission,
|
||||
);
|
||||
result.current.stepHandler(
|
||||
{ event: StepEvents.ON_RUN_STEP, data: secondRunStep },
|
||||
submission,
|
||||
);
|
||||
});
|
||||
|
||||
const responseMsg = currentMessages.find((m) => !m.isCreatedByUser);
|
||||
expect(responseMsg?.content).toHaveLength(2);
|
||||
expect(responseMsg?.content?.[0]?.tool_call?.name).toBe(`oauth${Constants.mcp_delimiter}ELI`);
|
||||
expect(responseMsg?.content?.[1]?.tool_call?.name).toBe(
|
||||
`oauth${Constants.mcp_delimiter}Vespa`,
|
||||
);
|
||||
});
|
||||
|
||||
it('should clear OAuth prompt slots when one occupies the real response slot', () => {
|
||||
const responseMessage = createResponseMessage();
|
||||
let currentMessages = [responseMessage];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
const submission = createSubmission({ initialResponse: responseMessage });
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler(
|
||||
{
|
||||
event: StepEvents.ON_RUN_STEP,
|
||||
data: createToolCallRunStep({
|
||||
id: 'step-oauth-eli',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-eli',
|
||||
name: `oauth${Constants.mcp_delimiter}ELI`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
},
|
||||
submission,
|
||||
);
|
||||
result.current.stepHandler(
|
||||
{
|
||||
event: StepEvents.ON_RUN_STEP,
|
||||
data: createToolCallRunStep({
|
||||
id: 'step-oauth-vespa',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 1,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-vespa',
|
||||
name: `oauth${Constants.mcp_delimiter}Vespa`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
}),
|
||||
},
|
||||
submission,
|
||||
);
|
||||
result.current.stepHandler(
|
||||
{
|
||||
event: StepEvents.ON_RUN_STEP,
|
||||
data: createRunStep({
|
||||
id: 'step-message',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 0,
|
||||
}),
|
||||
},
|
||||
submission,
|
||||
);
|
||||
result.current.stepHandler(
|
||||
{ event: StepEvents.ON_MESSAGE_DELTA, data: createMessageDelta('step-message', 'Ready') },
|
||||
submission,
|
||||
);
|
||||
});
|
||||
|
||||
const responseMsg = currentMessages.find((m) => !m.isCreatedByUser);
|
||||
expect(responseMsg?.content).toEqual([{ type: ContentTypes.TEXT, text: 'Ready' }]);
|
||||
});
|
||||
|
||||
it('should not replace the message list from a shorter refresh during tool call steps', () => {
|
||||
const userMessage = createUserMessage();
|
||||
const responseMessage = createResponseMessage();
|
||||
mockGetMessages.mockReturnValueOnce([userMessage, responseMessage]).mockReturnValueOnce([]);
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createToolCallRunStep({
|
||||
runId: responseMessage.messageId,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-eli',
|
||||
name: `oauth${Constants.mcp_delimiter}ELI`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler(
|
||||
{ event: StepEvents.ON_RUN_STEP, data: runStep },
|
||||
createSubmission({ userMessage, initialResponse: responseMessage }),
|
||||
);
|
||||
});
|
||||
|
||||
const lastCall = mockSetMessages.mock.calls[mockSetMessages.mock.calls.length - 1][0];
|
||||
expect(lastCall.map((message: TMessage) => message.messageId)).toEqual([
|
||||
userMessage.messageId,
|
||||
responseMessage.messageId,
|
||||
]);
|
||||
});
|
||||
|
||||
it('should replay buffered deltas after registering step', () => {
|
||||
const responseMessage = createResponseMessage();
|
||||
mockGetMessages.mockReturnValue([responseMessage]);
|
||||
|
|
@ -756,6 +932,59 @@ describe('useStepHandler', () => {
|
|||
);
|
||||
consoleSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('should mark completed OAuth prompts as finished', () => {
|
||||
const responseMessage = createResponseMessage();
|
||||
mockGetMessages.mockReturnValue([responseMessage]);
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createToolCallRunStep({
|
||||
id: 'step-oauth-eli',
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-eli',
|
||||
name: `oauth${Constants.mcp_delimiter}ELI`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
const submission = createSubmission();
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler(
|
||||
{
|
||||
event: StepEvents.ON_RUN_STEP_COMPLETED,
|
||||
data: {
|
||||
result: {
|
||||
id: 'step-oauth-eli',
|
||||
index: 0,
|
||||
tool_call: {
|
||||
id: 'tool-call-eli',
|
||||
name: `oauth${Constants.mcp_delimiter}ELI`,
|
||||
args: '',
|
||||
output: 'OAuth authentication completed',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
submission,
|
||||
);
|
||||
});
|
||||
|
||||
const lastCall = mockSetMessages.mock.calls[mockSetMessages.mock.calls.length - 1][0];
|
||||
const responseMsg = lastCall.find((m: TMessage) => !m.isCreatedByUser);
|
||||
expect(responseMsg?.content?.[0]?.tool_call?.progress).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('clearStepMaps', () => {
|
||||
|
|
|
|||
|
|
@ -63,6 +63,14 @@ type AllContentTypes =
|
|||
| ContentTypes.SUMMARY
|
||||
| ContentTypes.ERROR;
|
||||
|
||||
const isOAuthToolCallName = (name?: string) =>
|
||||
typeof name === 'string' && name.startsWith(`oauth${Constants.mcp_delimiter}`);
|
||||
|
||||
const isOAuthToolCallContent = (part?: Partial<TMessageContentParts>) =>
|
||||
part?.type === ContentTypes.TOOL_CALL &&
|
||||
'tool_call' in part &&
|
||||
isOAuthToolCallName(part.tool_call?.name);
|
||||
|
||||
export default function useStepHandler({
|
||||
setMessages,
|
||||
getMessages,
|
||||
|
|
@ -100,6 +108,14 @@ export default function useStepHandler({
|
|||
*/
|
||||
const knownSubagentAtomKeys = useRef(new Set<string>());
|
||||
|
||||
const getCurrentMessages = useCallback(
|
||||
(messages: TMessage[]) => {
|
||||
const freshMessages = getMessages();
|
||||
return freshMessages && freshMessages.length >= messages.length ? freshMessages : messages;
|
||||
},
|
||||
[getMessages],
|
||||
);
|
||||
|
||||
/** Both content parts and ticker lines are aggregated incrementally
|
||||
* into the atom as each `ON_SUBAGENT_UPDATE` arrives — we never
|
||||
* retain the raw event array, so no rolling window is needed. A
|
||||
|
|
@ -270,9 +286,20 @@ export default function useStepHandler({
|
|||
return message;
|
||||
}
|
||||
|
||||
const updatedContent = [...(message.content || [])] as Array<
|
||||
const incomingOAuthToolCall =
|
||||
contentType === ContentTypes.TOOL_CALL &&
|
||||
'tool_call' in contentPart &&
|
||||
isOAuthToolCallName(contentPart.tool_call?.name);
|
||||
|
||||
let updatedContent = [...(message.content || [])] as Array<
|
||||
Partial<TMessageContentParts> | undefined
|
||||
>;
|
||||
|
||||
const oauthPromptOccupiesSlot = isOAuthToolCallContent(updatedContent[index]);
|
||||
if (!incomingOAuthToolCall && oauthPromptOccupiesSlot) {
|
||||
updatedContent = updatedContent.filter((part) => !isOAuthToolCallContent(part));
|
||||
}
|
||||
|
||||
if (!updatedContent[index] && contentType !== ContentTypes.TOOL_CALL) {
|
||||
updatedContent[index] = { type: contentPart.type as AllContentTypes };
|
||||
}
|
||||
|
|
@ -517,9 +544,15 @@ export default function useStepHandler({
|
|||
});
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const updatedMessages = messages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
|
|
@ -724,9 +757,15 @@ export default function useStepHandler({
|
|||
});
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const updatedMessages = messages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
|
|
@ -767,9 +806,15 @@ export default function useStepHandler({
|
|||
);
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const updatedMessages = messages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
|
|
@ -883,6 +928,7 @@ export default function useStepHandler({
|
|||
announcePolite,
|
||||
setMessages,
|
||||
calculateContentIndex,
|
||||
getCurrentMessages,
|
||||
applySubagentUpdate,
|
||||
],
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue