mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-08 22:41:38 +00:00
📻 fix: Replay MCP OAuth Prompts for Coalesced Connections (#13565)
* fix: Replay MCP OAuth URL for Joined Connections * chore: Sort MCP OAuth Imports * test: Restore MCP OAuth Registry Spies * fix: Replay pending MCP OAuth prompts * fix: Replay MCP OAuth on Stream Resume * fix: Preserve MCP OAuth Replay Context * chore: Format MCP OAuth Replay Context * test: Expect MCP OAuth Replay Expiry * fix: Render pending MCP OAuth prompts * chore: Clean MCP OAuth Replay Type Narrowing * fix: Stabilize new MCP OAuth chats * fix: Re-emit cached MCP OAuth prompts * fix: Replay pending OAuth for selected MCP tools * fix: Avoid stalling pending MCP OAuth replay * test: Clean MCP OAuth review findings * test: Restore MCP OAuth registry spy * fix: Resolve OAuth Typecheck Regressions * fix: Harden MCP OAuth replay edge cases * test: Cover MCP OAuth joined prompt expiry * test: Mark joined OAuth replay fixture * test: Use OAuth fixture for joined replay expiry * fix: Anchor resumed MCP OAuth prompts * fix: Seed resumable turn metadata before MCP init * test: Format resume metadata regression * fix: Prioritize resumable stream routes * fix: Preserve MCP OAuth resume message tree * test: Fix MCP OAuth Resume Test Types * fix: Replay MCP OAuth Regenerate Prompts * fix: Skip OAuth-only Abort Persistence * fix: Stabilize OAuth Resume Replay * fix: Target Non-Tail Regenerate Responses * fix: Scope Regenerate Step Updates * fix: Clean Up OAuth Abort State * fix: Preserve Regenerate Branch Siblings * fix: Preserve OAuth Resume Branch State * fix: Preserve OAuth Branch Resume State * chore: Sort OAuth Resume Imports * fix: Address OAuth Resume Review Findings * test: Fix Abort Fixture Typing
This commit is contained in:
Danny Avila
GitHub
parent
6950448d03
commit
cb1d536874
48 changed files with 4837 additions and 299 deletions
|
|
@ -0,0 +1,197 @@
|
|||
const mockLogger = {
|
||||
debug: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
info: jest.fn(),
|
||||
};
|
||||
|
||||
const mockGenerationJobManager = {
|
||||
createJob: jest.fn(),
|
||||
emitError: jest.fn(),
|
||||
completeJob: jest.fn(),
|
||||
getResumeState: jest.fn(),
|
||||
updateMetadata: jest.fn(),
|
||||
};
|
||||
|
||||
const mockCheckAndIncrementPendingRequest = jest.fn();
|
||||
const mockDecrementPendingRequest = jest.fn();
|
||||
const mockFilterPersistableAbortContent = jest.fn((content) =>
|
||||
content.filter((part) => part?.type !== 'tool_call'),
|
||||
);
|
||||
const mockGetConvo = jest.fn();
|
||||
const mockSaveMessage = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: mockLogger,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
sendEvent: jest.fn(),
|
||||
getViolationInfo: jest.fn(),
|
||||
buildMessageFiles: jest.fn(() => []),
|
||||
resolveTitleTiming: jest.fn(() => 'immediate'),
|
||||
GenerationJobManager: mockGenerationJobManager,
|
||||
filterPersistableAbortContent: (...args) => mockFilterPersistableAbortContent(...args),
|
||||
decrementPendingRequest: (...args) => mockDecrementPendingRequest(...args),
|
||||
sanitizeMessageForTransmit: jest.fn((message) => message),
|
||||
checkAndIncrementPendingRequest: (...args) => mockCheckAndIncrementPendingRequest(...args),
|
||||
}));
|
||||
|
||||
jest.mock('~/server/cleanup', () => ({
|
||||
disposeClient: jest.fn(),
|
||||
clientRegistry: null,
|
||||
requestDataMap: {
|
||||
set: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('~/server/middleware', () => ({
|
||||
handleAbortError: jest.fn(() => Promise.resolve()),
|
||||
}));
|
||||
|
||||
jest.mock('~/cache', () => ({
|
||||
logViolation: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/models', () => ({
|
||||
saveMessage: (...args) => mockSaveMessage(...args),
|
||||
getConvo: (...args) => mockGetConvo(...args),
|
||||
}));
|
||||
|
||||
const AgentController = require('../request');
|
||||
|
||||
describe('ResumableAgentController resume metadata', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockCheckAndIncrementPendingRequest.mockResolvedValue({ allowed: true });
|
||||
mockDecrementPendingRequest.mockResolvedValue(undefined);
|
||||
mockGetConvo.mockResolvedValue({ createdAt: '2026-06-07T00:00:00.000Z' });
|
||||
mockGenerationJobManager.createJob.mockResolvedValue({
|
||||
createdAt: 1000,
|
||||
readyPromise: Promise.resolve(),
|
||||
abortController: new AbortController(),
|
||||
emitter: { on: jest.fn() },
|
||||
});
|
||||
mockGenerationJobManager.getResumeState.mockResolvedValue(null);
|
||||
mockGenerationJobManager.updateMetadata.mockResolvedValue(undefined);
|
||||
mockGenerationJobManager.emitError.mockResolvedValue(undefined);
|
||||
mockSaveMessage.mockResolvedValue({});
|
||||
});
|
||||
|
||||
it('stores the in-flight turn before MCP initialization can emit OAuth', async () => {
|
||||
const conversationId = 'conversation-123';
|
||||
const initializeClient = jest.fn().mockRejectedValue(new Error('stop before tool loading'));
|
||||
const req = {
|
||||
user: { id: 'user-123' },
|
||||
body: {
|
||||
text: 'Check Google Workspace availability.',
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId,
|
||||
endpointOption: {
|
||||
endpoint: 'agents',
|
||||
modelOptions: { model: 'gpt-3.5-turbo' },
|
||||
},
|
||||
},
|
||||
config: {},
|
||||
};
|
||||
const res = {
|
||||
headersSent: true,
|
||||
json: jest.fn(() => {
|
||||
res.headersSent = true;
|
||||
}),
|
||||
status: jest.fn(() => res),
|
||||
};
|
||||
|
||||
await AgentController(req, res, jest.fn(), initializeClient, null);
|
||||
|
||||
expect(mockGenerationJobManager.updateMetadata).toHaveBeenCalledWith(conversationId, {
|
||||
conversationId,
|
||||
responseMessageId: 'follow-up-user_',
|
||||
userMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId,
|
||||
text: 'Check Google Workspace availability.',
|
||||
},
|
||||
});
|
||||
expect(mockGenerationJobManager.updateMetadata.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
initializeClient.mock.invocationCallOrder[0],
|
||||
);
|
||||
});
|
||||
|
||||
it('filters OAuth prompts before saving partial responses on disconnect', async () => {
|
||||
const conversationId = 'conversation-123';
|
||||
let allSubscribersLeftHandler;
|
||||
mockGenerationJobManager.createJob.mockResolvedValue({
|
||||
createdAt: 1000,
|
||||
readyPromise: Promise.resolve(),
|
||||
abortController: new AbortController(),
|
||||
emitter: {
|
||||
on: jest.fn((event, handler) => {
|
||||
if (event === 'allSubscribersLeft') {
|
||||
allSubscribersLeftHandler = handler;
|
||||
}
|
||||
}),
|
||||
},
|
||||
});
|
||||
mockGenerationJobManager.getResumeState.mockResolvedValue({
|
||||
conversationId,
|
||||
responseMessageId: 'response-message',
|
||||
userMessage: {
|
||||
messageId: 'user-message',
|
||||
parentMessageId: 'parent-message',
|
||||
conversationId,
|
||||
text: 'Use Google Workspace',
|
||||
},
|
||||
});
|
||||
|
||||
const initializeClient = jest.fn().mockRejectedValue(new Error('stop after setup'));
|
||||
const req = {
|
||||
user: { id: 'user-123' },
|
||||
body: {
|
||||
text: 'Use Google Workspace',
|
||||
messageId: 'user-message',
|
||||
parentMessageId: 'parent-message',
|
||||
conversationId,
|
||||
endpointOption: {
|
||||
endpoint: 'agents',
|
||||
modelOptions: { model: 'gpt-3.5-turbo' },
|
||||
},
|
||||
},
|
||||
config: {},
|
||||
};
|
||||
const res = {
|
||||
headersSent: true,
|
||||
json: jest.fn(() => {
|
||||
res.headersSent = true;
|
||||
}),
|
||||
status: jest.fn(() => res),
|
||||
};
|
||||
|
||||
await AgentController(req, res, jest.fn(), initializeClient, null);
|
||||
expect(allSubscribersLeftHandler).toEqual(expect.any(Function));
|
||||
|
||||
const oauthPart = {
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
},
|
||||
};
|
||||
const textPart = { type: 'text', text: 'Partial response...' };
|
||||
|
||||
await allSubscribersLeftHandler([oauthPart, textPart]);
|
||||
|
||||
expect(mockFilterPersistableAbortContent).toHaveBeenCalledWith([oauthPart, textPart]);
|
||||
expect(mockSaveMessage).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ userId: 'user-123' }),
|
||||
expect.objectContaining({
|
||||
content: [textPart],
|
||||
messageId: 'response-message',
|
||||
parentMessageId: 'user-message',
|
||||
}),
|
||||
expect.any(Object),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -6,6 +6,7 @@ const {
|
|||
buildMessageFiles,
|
||||
resolveTitleTiming,
|
||||
GenerationJobManager,
|
||||
filterPersistableAbortContent,
|
||||
decrementPendingRequest,
|
||||
sanitizeMessageForTransmit,
|
||||
checkAndIncrementPendingRequest,
|
||||
|
|
@ -75,6 +76,31 @@ async function attachConversationCreatedAt(req, { userId, conversationId, isNewC
|
|||
}
|
||||
}
|
||||
|
||||
function getPreliminaryResponseMessageId({ messageId, responseMessageId }) {
|
||||
if (typeof responseMessageId === 'string' && responseMessageId.length > 0) {
|
||||
return responseMessageId;
|
||||
}
|
||||
|
||||
if (typeof messageId !== 'string' || messageId.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return `${messageId.replace(/_+$/, '')}_`;
|
||||
}
|
||||
|
||||
function getPreliminaryUserMessage({ messageId, parentMessageId, text }, conversationId) {
|
||||
if (typeof messageId !== 'string' || messageId.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
messageId,
|
||||
parentMessageId,
|
||||
conversationId,
|
||||
text,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Resumable Agent Controller - Generation runs independently of HTTP connection.
|
||||
* Returns streamId immediately, client subscribes separately via SSE.
|
||||
|
|
@ -134,6 +160,16 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
|
||||
await attachConversationCreatedAt(req, { userId, conversationId, isNewConvo });
|
||||
|
||||
const preliminaryUserMessage = getPreliminaryUserMessage(req.body, conversationId);
|
||||
const preliminaryResponseMessageId = getPreliminaryResponseMessageId(req.body);
|
||||
if (preliminaryUserMessage || preliminaryResponseMessageId) {
|
||||
await GenerationJobManager.updateMetadata(streamId, {
|
||||
conversationId,
|
||||
responseMessageId: preliminaryResponseMessageId,
|
||||
userMessage: preliminaryUserMessage,
|
||||
});
|
||||
}
|
||||
|
||||
// Note: We no longer use res.on('close') to abort since we send JSON immediately.
|
||||
// The response closes normally after res.json(), which is not an abort condition.
|
||||
// Abort handling is done through GenerationJobManager via the SSE stream connection.
|
||||
|
|
@ -155,6 +191,12 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
return;
|
||||
}
|
||||
|
||||
const persistableContent = filterPersistableAbortContent(aggregatedContent);
|
||||
if (persistableContent.length === 0) {
|
||||
logger.debug('[ResumableAgentController] No persistable content to save partial response');
|
||||
return;
|
||||
}
|
||||
|
||||
const resumeState = await GenerationJobManager.getResumeState(streamId);
|
||||
if (!resumeState?.userMessage) {
|
||||
logger.debug('[ResumableAgentController] No user message to save partial response for');
|
||||
|
|
@ -170,7 +212,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
conversationId: responseConversationId,
|
||||
parentMessageId: resumeState.userMessage.messageId,
|
||||
sender: client?.sender ?? 'AI',
|
||||
content: aggregatedContent,
|
||||
content: persistableContent,
|
||||
unfinished: true,
|
||||
error: false,
|
||||
isCreatedByUser: false,
|
||||
|
|
@ -194,7 +236,7 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit
|
|||
);
|
||||
|
||||
logger.debug(
|
||||
`[ResumableAgentController] Saved partial response for ${streamId}, content parts: ${aggregatedContent.length}`,
|
||||
`[ResumableAgentController] Saved partial response for ${streamId}, content parts: ${persistableContent.length}`,
|
||||
);
|
||||
} catch (error) {
|
||||
logger.error('[ResumableAgentController] Error saving partial response:', error);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue