mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-23 10:36:09 +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);
|
||||
|
|
|
|||
|
|
@ -195,6 +195,43 @@ describe('Agent Abort Endpoint', () => {
|
|||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should skip message saving when abort content is only an OAuth prompt', async () => {
|
||||
const jobStreamId = 'test-stream-123';
|
||||
|
||||
mockGenerationJobManager.getJob.mockResolvedValue({
|
||||
metadata: { userId: 'test-user-123' },
|
||||
});
|
||||
|
||||
mockGenerationJobManager.abortJob.mockResolvedValue({
|
||||
success: true,
|
||||
jobData: {
|
||||
userMessage: { messageId: 'user-msg-123' },
|
||||
responseMessageId: 'response-msg-456',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [
|
||||
{
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
type: 'tool_call',
|
||||
id: 'oauth-call-1',
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
args: '',
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
},
|
||||
},
|
||||
],
|
||||
text: '',
|
||||
});
|
||||
|
||||
const response = await request(app)
|
||||
.post('/api/agents/chat/abort')
|
||||
.send({ conversationId: jobStreamId });
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockSaveMessage).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Partial Response Saving', () => {
|
||||
|
|
@ -264,8 +301,8 @@ describe('Agent Abort Endpoint', () => {
|
|||
responseMessageId: 'response-msg-456',
|
||||
conversationId: jobStreamId,
|
||||
},
|
||||
content: [],
|
||||
text: '',
|
||||
content: [{ type: 'text', text: 'Partial response...' }],
|
||||
text: 'Partial response...',
|
||||
});
|
||||
|
||||
mockSaveMessage.mockRejectedValue(new Error('Database error'));
|
||||
|
|
|
|||
|
|
@ -45,9 +45,11 @@ jest.mock('~/server/middleware', () => ({
|
|||
}));
|
||||
|
||||
jest.mock('~/server/routes/agents/chat', () => require('express').Router());
|
||||
jest.mock('~/server/routes/agents/v1', () => ({
|
||||
v1: require('express').Router(),
|
||||
}));
|
||||
jest.mock('~/server/routes/agents/v1', () => {
|
||||
const router = require('express').Router();
|
||||
router.use((req, res) => res.status(418).json({ error: 'v1 caught stream route' }));
|
||||
return { v1: router };
|
||||
});
|
||||
jest.mock('~/server/routes/agents/openai', () => require('express').Router());
|
||||
jest.mock('~/server/routes/agents/responses', () => require('express').Router());
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
const express = require('express');
|
||||
const { isEnabled, GenerationJobManager } = require('@librechat/api');
|
||||
const { isEnabled, GenerationJobManager, hasPersistableAbortContent } = require('@librechat/api');
|
||||
const { createSseStreamTelemetry } = require('@librechat/api/telemetry');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const {
|
||||
|
|
@ -43,8 +43,6 @@ router.use(requireJwtAuth);
|
|||
router.use(checkBan);
|
||||
router.use(uaParser);
|
||||
|
||||
router.use('/', v1);
|
||||
|
||||
/**
|
||||
* Stream endpoints - mounted before chatRouter to bypass rate limiters
|
||||
* These are GET requests and don't need message body validation or rate limiting
|
||||
|
|
@ -273,7 +271,8 @@ router.post('/chat/abort', async (req, res) => {
|
|||
if (
|
||||
abortResult.success &&
|
||||
abortResult.jobData?.userMessage?.messageId &&
|
||||
abortResult.jobData?.responseMessageId
|
||||
abortResult.jobData?.responseMessageId &&
|
||||
hasPersistableAbortContent(abortResult.content)
|
||||
) {
|
||||
const { jobData, content, text } = abortResult;
|
||||
const responseMessage = {
|
||||
|
|
@ -314,6 +313,8 @@ router.post('/chat/abort', async (req, res) => {
|
|||
return res.status(404).json({ error: 'Job not found', streamId: jobStreamId });
|
||||
});
|
||||
|
||||
router.use('/', v1);
|
||||
|
||||
const chatRouter = express.Router();
|
||||
chatRouter.use(configMiddleware);
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,6 @@
|
|||
const { tool } = require('@librechat/agents/langchain/tools');
|
||||
const { logger, getTenantId } = require('@librechat/data-schemas');
|
||||
const {
|
||||
Providers,
|
||||
StepTypes,
|
||||
GraphEvents,
|
||||
Constants: AgentConstants,
|
||||
} = require('@librechat/agents');
|
||||
const { Providers, Constants: AgentConstants } = require('@librechat/agents');
|
||||
const {
|
||||
sendEvent,
|
||||
MCPOAuthHandler,
|
||||
|
|
@ -14,7 +9,11 @@ const {
|
|||
normalizeJsonSchema,
|
||||
GenerationJobManager,
|
||||
resolveJsonSchemaRefs,
|
||||
buildOAuthToolCallName,
|
||||
buildMCPAuthStepId,
|
||||
buildMCPAuthToolCall,
|
||||
buildMCPAuthRunStepEvent,
|
||||
buildMCPAuthRunStepDeltaEvent,
|
||||
buildMCPAuthRunStepEndDeltaEvent,
|
||||
checkAccessWithRequestCache,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
|
|
@ -201,20 +200,11 @@ function isEmptyObjectSchema(jsonSchema) {
|
|||
function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null }) {
|
||||
/**
|
||||
* @param {string} authURL - The URL to redirect the user for OAuth authentication.
|
||||
* @param {{ expiresAt?: number }} [options]
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
return async function (authURL) {
|
||||
/** @type {{ id: string; delta: AgentToolCallDelta }} */
|
||||
const data = {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [{ ...toolCall, args: '' }],
|
||||
auth: authURL,
|
||||
expires_at: Date.now() + Time.TWO_MINUTES,
|
||||
},
|
||||
};
|
||||
const eventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data };
|
||||
return async function (authURL, options) {
|
||||
const eventData = buildMCPAuthRunStepDeltaEvent({ authURL, stepId, toolCall, options });
|
||||
if (streamId) {
|
||||
await GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
|
|
@ -235,18 +225,7 @@ function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null }) {
|
|||
*/
|
||||
function createRunStepEmitter({ res, runId, stepId, toolCall, index, streamId = null }) {
|
||||
return async function () {
|
||||
/** @type {import('@librechat/agents').RunStep} */
|
||||
const data = {
|
||||
runId: runId ?? Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
id: stepId,
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
index: index ?? 0,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [toolCall],
|
||||
},
|
||||
};
|
||||
const eventData = { event: GraphEvents.ON_RUN_STEP, data };
|
||||
const eventData = buildMCPAuthRunStepEvent({ runId, stepId, toolCall, index });
|
||||
if (streamId) {
|
||||
await GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
|
|
@ -260,20 +239,43 @@ function createRunStepEmitter({ res, runId, stepId, toolCall, index, streamId =
|
|||
* @param {object} params
|
||||
* @param {string} params.flowId - The ID of the login flow.
|
||||
* @param {FlowStateManager<any>} params.flowManager - The flow manager instance.
|
||||
* @param {(authURL: string) => void} [params.callback]
|
||||
* @param {(authURL: string, options?: { expiresAt?: number }) => void | Promise<void>} [params.callback]
|
||||
*/
|
||||
function createOAuthStart({ flowId, flowManager, callback }) {
|
||||
/**
|
||||
* Creates a function to handle OAuth login requests.
|
||||
* @param {string} authURL - The URL to redirect the user for OAuth authentication.
|
||||
* @param {{ expiresAt?: number }} [options]
|
||||
* @returns {Promise<boolean>} Returns true to indicate the event was sent successfully.
|
||||
*/
|
||||
return async function (authURL) {
|
||||
return async function (authURL, options) {
|
||||
let emitted = false;
|
||||
const emitOAuthStart = async (message) => {
|
||||
if (options) {
|
||||
await callback?.(authURL, options);
|
||||
} else {
|
||||
await callback?.(authURL);
|
||||
}
|
||||
emitted = true;
|
||||
logger.debug(message);
|
||||
};
|
||||
|
||||
const existingFlow = await flowManager.getFlowState(flowId, 'oauth_login');
|
||||
if (existingFlow) {
|
||||
await emitOAuthStart('Re-sent OAuth login request to client');
|
||||
return true;
|
||||
}
|
||||
|
||||
await flowManager.createFlowWithHandler(flowId, 'oauth_login', async () => {
|
||||
callback?.(authURL);
|
||||
logger.debug('Sent OAuth login request to client');
|
||||
await emitOAuthStart('Sent OAuth login request to client');
|
||||
return true;
|
||||
});
|
||||
|
||||
if (!emitted) {
|
||||
await emitOAuthStart('Re-sent OAuth login request to client');
|
||||
}
|
||||
|
||||
return true;
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -286,15 +288,7 @@ function createOAuthStart({ flowId, flowManager, callback }) {
|
|||
*/
|
||||
function createOAuthEnd({ res, stepId, toolCall, streamId = null }) {
|
||||
return async function () {
|
||||
/** @type {{ id: string; delta: AgentToolCallDelta }} */
|
||||
const data = {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [{ ...toolCall }],
|
||||
},
|
||||
};
|
||||
const eventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data };
|
||||
const eventData = buildMCPAuthRunStepEndDeltaEvent({ stepId, toolCall });
|
||||
if (streamId) {
|
||||
await GenerationJobManager.emitChunk(streamId, eventData);
|
||||
} else {
|
||||
|
|
@ -323,14 +317,14 @@ function createAbortHandler({ userId, serverName, toolName, flowManager }) {
|
|||
|
||||
/**
|
||||
* @param {Object} params
|
||||
* @param {() => void} params.runStepEmitter
|
||||
* @param {(authURL: string) => void} params.runStepDeltaEmitter
|
||||
* @returns {(authURL: string) => void}
|
||||
* @param {() => Promise<void>} params.runStepEmitter
|
||||
* @param {(authURL: string, options?: { expiresAt?: number }) => Promise<void>} params.runStepDeltaEmitter
|
||||
* @returns {(authURL: string, options?: { expiresAt?: number }) => Promise<void>}
|
||||
*/
|
||||
function createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }) {
|
||||
return function (authURL) {
|
||||
runStepEmitter();
|
||||
runStepDeltaEmitter(authURL);
|
||||
return async function (authURL, options) {
|
||||
await runStepEmitter();
|
||||
await runStepDeltaEmitter(authURL, options);
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -373,12 +367,11 @@ async function reconnectServer({
|
|||
const runId = Constants.USE_PRELIM_RESPONSE_MESSAGE_ID;
|
||||
const flowId = `${user.id}:${serverName}:${Date.now()}`;
|
||||
const flowManager = getFlowStateManager(getLogStores(CacheKeys.FLOWS));
|
||||
const stepId = 'step_oauth_login_' + serverName;
|
||||
const toolCall = {
|
||||
const stepId = buildMCPAuthStepId(serverName);
|
||||
const toolCall = buildMCPAuthToolCall({
|
||||
id: flowId,
|
||||
name: buildOAuthToolCallName(serverName),
|
||||
type: 'tool_call_chunk',
|
||||
};
|
||||
serverName,
|
||||
});
|
||||
|
||||
// Set up abort handler to clean up OAuth flows if request is aborted
|
||||
const oauthFlowId = MCPOAuthHandler.generateFlowId(user.id, serverName);
|
||||
|
|
@ -949,6 +942,7 @@ module.exports = {
|
|||
resolveConfigServers,
|
||||
resolveMcpConfigNames,
|
||||
resolveAllMcpConfigs,
|
||||
createOAuthStart,
|
||||
checkOAuthFlowStatus,
|
||||
getServerConnectionStatus,
|
||||
createUnavailableToolStub,
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ const {
|
|||
createMCPTools,
|
||||
createMCPPermissionContext,
|
||||
getMCPSetupData,
|
||||
createOAuthStart,
|
||||
checkOAuthFlowStatus,
|
||||
getServerConnectionStatus,
|
||||
createUnavailableToolStub,
|
||||
|
|
@ -100,6 +101,85 @@ describe('tests for the new helper functions used by the MCP connection status e
|
|||
mockGetOAuthReconnectionManager = require('~/config').getOAuthReconnectionManager;
|
||||
});
|
||||
|
||||
describe('createOAuthStart', () => {
|
||||
const flowId = 'test-server:oauth_login:thread-1:run-1';
|
||||
const authUrl = 'https://auth.example.com/oauth?state=test';
|
||||
|
||||
it('should create a login flow and emit the OAuth URL for the first request', async () => {
|
||||
const callback = jest.fn();
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue(null),
|
||||
createFlowWithHandler: jest.fn(async (_flowId, _type, handler) => handler()),
|
||||
};
|
||||
|
||||
const oauthStart = createOAuthStart({
|
||||
flowId,
|
||||
flowManager: mockFlowManager,
|
||||
callback,
|
||||
});
|
||||
|
||||
await expect(oauthStart(authUrl)).resolves.toBe(true);
|
||||
|
||||
expect(mockFlowManager.getFlowState).toHaveBeenCalledWith(flowId, 'oauth_login');
|
||||
expect(mockFlowManager.createFlowWithHandler).toHaveBeenCalledWith(
|
||||
flowId,
|
||||
'oauth_login',
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(callback).toHaveBeenCalledWith(authUrl);
|
||||
expect(logger.debug).toHaveBeenCalledWith('Sent OAuth login request to client');
|
||||
});
|
||||
|
||||
it('should replay the OAuth URL when the login flow already exists', async () => {
|
||||
const callback = jest.fn();
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({
|
||||
status: 'COMPLETED',
|
||||
result: true,
|
||||
}),
|
||||
createFlowWithHandler: jest.fn(),
|
||||
};
|
||||
|
||||
const oauthStart = createOAuthStart({
|
||||
flowId,
|
||||
flowManager: mockFlowManager,
|
||||
callback,
|
||||
});
|
||||
|
||||
await expect(oauthStart(authUrl)).resolves.toBe(true);
|
||||
|
||||
expect(mockFlowManager.getFlowState).toHaveBeenCalledWith(flowId, 'oauth_login');
|
||||
expect(mockFlowManager.createFlowWithHandler).not.toHaveBeenCalled();
|
||||
expect(callback).toHaveBeenCalledWith(authUrl);
|
||||
expect(logger.debug).toHaveBeenCalledWith('Re-sent OAuth login request to client');
|
||||
});
|
||||
|
||||
it('should replay the OAuth URL when flow creation is deduped internally', async () => {
|
||||
const callback = jest.fn();
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue(null),
|
||||
createFlowWithHandler: jest.fn().mockResolvedValue(true),
|
||||
};
|
||||
|
||||
const oauthStart = createOAuthStart({
|
||||
flowId,
|
||||
flowManager: mockFlowManager,
|
||||
callback,
|
||||
});
|
||||
|
||||
await expect(oauthStart(authUrl)).resolves.toBe(true);
|
||||
|
||||
expect(mockFlowManager.getFlowState).toHaveBeenCalledWith(flowId, 'oauth_login');
|
||||
expect(mockFlowManager.createFlowWithHandler).toHaveBeenCalledWith(
|
||||
flowId,
|
||||
'oauth_login',
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(callback).toHaveBeenCalledWith(authUrl);
|
||||
expect(logger.debug).toHaveBeenCalledWith('Re-sent OAuth login request to client');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getMCPSetupData', () => {
|
||||
const mockUserId = 'user-123';
|
||||
const mockConfig = {
|
||||
|
|
|
|||
|
|
@ -2,8 +2,6 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { tool: toolFn, DynamicStructuredTool } = require('@librechat/agents/langchain/tools');
|
||||
const {
|
||||
sleep,
|
||||
StepTypes,
|
||||
GraphEvents,
|
||||
createToolSearch,
|
||||
createBashExecutionTool,
|
||||
Constants: AgentConstants,
|
||||
|
|
@ -18,11 +16,17 @@ const {
|
|||
isActionDomainAllowed,
|
||||
buildWebSearchContext,
|
||||
buildImageToolContext,
|
||||
buildOAuthToolCallName,
|
||||
buildToolClassification,
|
||||
getMissingCustomUserVars,
|
||||
buildWebSearchDynamicContext,
|
||||
getCodeApiAuthHeaders,
|
||||
getReplayablePendingMCPOAuthStart,
|
||||
getMCPServerNamesFromTools,
|
||||
buildMCPAuthToolCall,
|
||||
buildMCPAuthStepId,
|
||||
buildMCPAuthRunStepEvent,
|
||||
buildMCPAuthRunStepDeltaEvent,
|
||||
buildMCPAuthRunStepCompletedEvent,
|
||||
isFileAuthoringToolDefinition,
|
||||
} = require('@librechat/api');
|
||||
const {
|
||||
|
|
@ -596,44 +600,35 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
const flowManager = getFlowStateManager(flowsCache);
|
||||
const configServers = await resolveConfigServers(req);
|
||||
const pendingOAuthServers = new Set();
|
||||
const pendingOAuthStarts = new Map();
|
||||
const emittedOAuthStarts = new Map();
|
||||
const oauthToolCallIds = new Map();
|
||||
const oauthStepIndexes = new Map();
|
||||
|
||||
const createOAuthEmitter = (serverName, index) => {
|
||||
return async (authURL) => {
|
||||
const flowId = `${req.user.id}:${serverName}:${Date.now()}`;
|
||||
const stepId = 'step_oauth_login_' + serverName;
|
||||
return async (authURL, options) => {
|
||||
if (emittedOAuthStarts.get(serverName) === authURL) {
|
||||
return;
|
||||
}
|
||||
emittedOAuthStarts.set(serverName, authURL);
|
||||
|
||||
const flowId =
|
||||
oauthToolCallIds.get(serverName) ?? `${req.user.id}:${serverName}:${Date.now()}`;
|
||||
const stepId = buildMCPAuthStepId(serverName);
|
||||
oauthToolCallIds.set(serverName, flowId);
|
||||
oauthStepIndexes.set(serverName, index);
|
||||
const toolCall = {
|
||||
const toolCall = buildMCPAuthToolCall({
|
||||
id: flowId,
|
||||
name: buildOAuthToolCallName(serverName),
|
||||
type: 'tool_call_chunk',
|
||||
};
|
||||
serverName,
|
||||
});
|
||||
|
||||
const runStepData = {
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
id: stepId,
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
index,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [toolCall],
|
||||
},
|
||||
};
|
||||
|
||||
const runStepDeltaData = {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [{ ...toolCall, args: '' }],
|
||||
auth: authURL,
|
||||
expires_at: Date.now() + Time.TWO_MINUTES,
|
||||
},
|
||||
};
|
||||
|
||||
const runStepEvent = { event: GraphEvents.ON_RUN_STEP, data: runStepData };
|
||||
const runStepDeltaEvent = { event: GraphEvents.ON_RUN_STEP_DELTA, data: runStepDeltaData };
|
||||
const runStepEvent = buildMCPAuthRunStepEvent({ stepId, toolCall, index });
|
||||
const runStepDeltaEvent = buildMCPAuthRunStepDeltaEvent({
|
||||
authURL,
|
||||
stepId,
|
||||
toolCall,
|
||||
options,
|
||||
});
|
||||
|
||||
if (streamId) {
|
||||
await GenerationJobManager.emitChunk(streamId, runStepEvent);
|
||||
|
|
@ -651,25 +646,19 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
|
||||
const createOAuthEndEmitter = (serverName) => {
|
||||
return async () => {
|
||||
const stepId = 'step_oauth_login_' + serverName;
|
||||
const toolCall = {
|
||||
const stepId = buildMCPAuthStepId(serverName);
|
||||
const toolCall = buildMCPAuthToolCall({
|
||||
id: oauthToolCallIds.get(serverName),
|
||||
name: buildOAuthToolCallName(serverName),
|
||||
args: '',
|
||||
output: 'OAuth authentication completed',
|
||||
serverName,
|
||||
type: 'tool_call',
|
||||
};
|
||||
|
||||
const runStepCompletedEvent = {
|
||||
event: GraphEvents.ON_RUN_STEP_COMPLETED,
|
||||
data: {
|
||||
result: {
|
||||
id: stepId,
|
||||
index: oauthStepIndexes.get(serverName) ?? 0,
|
||||
tool_call: toolCall,
|
||||
},
|
||||
},
|
||||
};
|
||||
});
|
||||
const runStepCompletedEvent = buildMCPAuthRunStepCompletedEvent({
|
||||
stepId,
|
||||
toolCall,
|
||||
index: oauthStepIndexes.get(serverName) ?? 0,
|
||||
});
|
||||
|
||||
if (streamId) {
|
||||
await GenerationJobManager.emitChunk(streamId, runStepCompletedEvent);
|
||||
|
|
@ -683,7 +672,45 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
};
|
||||
};
|
||||
|
||||
const getPendingOAuthStartForEmit = async (serverName) => {
|
||||
const cachedOAuthStart = pendingOAuthStarts.get(serverName);
|
||||
if (cachedOAuthStart?.options?.expiresAt != null) {
|
||||
return cachedOAuthStart;
|
||||
}
|
||||
|
||||
const pendingOAuthStart = await getReplayablePendingMCPOAuthStart({
|
||||
flowManager,
|
||||
userId: req.user.id,
|
||||
serverName,
|
||||
});
|
||||
if (!pendingOAuthStart) {
|
||||
return cachedOAuthStart;
|
||||
}
|
||||
|
||||
if (!cachedOAuthStart || pendingOAuthStart.authURL === cachedOAuthStart.authURL) {
|
||||
pendingOAuthStarts.set(serverName, pendingOAuthStart);
|
||||
return pendingOAuthStart;
|
||||
}
|
||||
|
||||
return cachedOAuthStart;
|
||||
};
|
||||
|
||||
const getOrFetchMCPServerTools = async (userId, serverName) => {
|
||||
const addPendingOAuthServer = async () => {
|
||||
const pendingOAuthStart = await getReplayablePendingMCPOAuthStart({
|
||||
flowManager,
|
||||
userId,
|
||||
serverName,
|
||||
});
|
||||
if (!pendingOAuthStart) {
|
||||
return false;
|
||||
}
|
||||
|
||||
pendingOAuthServers.add(serverName);
|
||||
pendingOAuthStarts.set(serverName, pendingOAuthStart);
|
||||
return true;
|
||||
};
|
||||
|
||||
let serverConfig;
|
||||
try {
|
||||
serverConfig =
|
||||
|
|
@ -718,11 +745,19 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
|
||||
const cached = await getMCPServerTools(userId, serverName);
|
||||
if (cached) {
|
||||
await addPendingOAuthServer();
|
||||
return cached;
|
||||
}
|
||||
|
||||
const oauthStart = async () => {
|
||||
if (await addPendingOAuthServer()) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const oauthStart = async (authURL, options) => {
|
||||
pendingOAuthServers.add(serverName);
|
||||
if (typeof authURL === 'string' && authURL.length > 0) {
|
||||
pendingOAuthStarts.set(serverName, { authURL, options });
|
||||
}
|
||||
};
|
||||
|
||||
const result = await reinitMCPServer({
|
||||
|
|
@ -813,6 +848,22 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
},
|
||||
);
|
||||
|
||||
for (const serverName of getMCPServerNamesFromTools(filteredTools)) {
|
||||
if (pendingOAuthServers.has(serverName)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const pendingOAuthStart = await getReplayablePendingMCPOAuthStart({
|
||||
flowManager,
|
||||
userId: req.user.id,
|
||||
serverName,
|
||||
});
|
||||
if (pendingOAuthStart) {
|
||||
pendingOAuthServers.add(serverName);
|
||||
pendingOAuthStarts.set(serverName, pendingOAuthStart);
|
||||
}
|
||||
}
|
||||
|
||||
if (pendingOAuthServers.size > 0 && (res || streamId)) {
|
||||
const serverNames = Array.from(pendingOAuthServers);
|
||||
logger.info(
|
||||
|
|
@ -821,6 +872,12 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
|
||||
const oauthWaitPromises = serverNames.map(async (serverName, index) => {
|
||||
try {
|
||||
const pendingOAuthStart = await getPendingOAuthStartForEmit(serverName);
|
||||
const oauthStart = createOAuthEmitter(serverName, index);
|
||||
if (pendingOAuthStart) {
|
||||
await oauthStart(pendingOAuthStart.authURL, pendingOAuthStart.options);
|
||||
}
|
||||
|
||||
const result = await reinitMCPServer({
|
||||
user: req.user,
|
||||
serverName,
|
||||
|
|
@ -828,7 +885,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to
|
|||
userMCPAuthMap,
|
||||
flowManager,
|
||||
returnOnOAuth: false,
|
||||
oauthStart: createOAuthEmitter(serverName, index),
|
||||
oauthStart,
|
||||
oauthEnd: createOAuthEndEmitter(serverName),
|
||||
connectionTimeout: Time.TWO_MINUTES,
|
||||
});
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ const { getLogStores } = require('~/cache');
|
|||
* @param {boolean} [params.forceNew]
|
||||
* @param {number} [params.connectionTimeout]
|
||||
* @param {FlowStateManager<any>} [params.flowManager]
|
||||
* @param {(authURL: string) => Promise<void>} [params.oauthStart]
|
||||
* @param {(authURL: string, options?: { expiresAt?: number }) => Promise<void>} [params.oauthStart]
|
||||
* @param {() => Promise<void>} [params.oauthEnd]
|
||||
* @param {Record<string, Record<string, string>>} [params.userMCPAuthMap]
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ const mockLegacyDomainEncode = jest.fn();
|
|||
const mockDecryptMetadata = jest.fn();
|
||||
const mockCreateActionTool = jest.fn();
|
||||
const mockGetServerConfig = jest.fn();
|
||||
const mockFlowManager = { getFlowState: jest.fn() };
|
||||
const mockResolveConfigServers = jest.fn();
|
||||
const mockUserCanUseMCPServers = jest.fn().mockResolvedValue(true);
|
||||
jest.mock('~/server/services/Tools/credentials', () => ({
|
||||
|
|
@ -77,7 +78,7 @@ jest.mock('~/models', () => ({
|
|||
findPluginAuthsByKeys: jest.fn(),
|
||||
}));
|
||||
jest.mock('~/config', () => ({
|
||||
getFlowStateManager: jest.fn(() => ({})),
|
||||
getFlowStateManager: jest.fn(() => mockFlowManager),
|
||||
getMCPServersRegistry: jest.fn(() => ({
|
||||
getServerConfig: (...args) => mockGetServerConfig(...args),
|
||||
})),
|
||||
|
|
@ -100,6 +101,7 @@ const {
|
|||
resolveAgentCapabilities,
|
||||
} = require('../ToolService');
|
||||
const { reinitMCPServer } = require('~/server/services/Tools/mcp');
|
||||
const { PENDING_STALE_MS } = require('@librechat/api');
|
||||
|
||||
function createMockReq(capabilities) {
|
||||
return {
|
||||
|
|
@ -134,6 +136,7 @@ describe('ToolService - Action Capability Gating', () => {
|
|||
mockGetCachedTools.mockResolvedValue(null);
|
||||
mockGetUserMCPAuthMap.mockResolvedValue({});
|
||||
mockGetServerConfig.mockResolvedValue(undefined);
|
||||
mockFlowManager.getFlowState.mockResolvedValue(undefined);
|
||||
mockResolveConfigServers.mockResolvedValue({});
|
||||
});
|
||||
|
||||
|
|
@ -421,6 +424,349 @@ describe('ToolService - Action Capability Gating', () => {
|
|||
expect(mockGetMCPServerTools).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should re-emit pending MCP OAuth prompts when cached tool definitions exist', async () => {
|
||||
const serverName = 'Google-Workspace';
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockGetServerConfig.mockResolvedValue({
|
||||
type: 'streamable-http',
|
||||
url: 'https://demo.librechat.ai/mcp',
|
||||
requiresOAuth: true,
|
||||
});
|
||||
mockGetMCPServerTools.mockResolvedValue({
|
||||
[mcpTool]: {
|
||||
function: {
|
||||
name: mcpTool,
|
||||
description: 'Cached search',
|
||||
parameters: {},
|
||||
},
|
||||
},
|
||||
});
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
createdAt: Date.now(),
|
||||
metadata: { authorizationUrl },
|
||||
});
|
||||
mockLoadToolDefinitions.mockImplementation(async (params, deps) => {
|
||||
const serverTools = await deps.getOrFetchMCPServerTools(params.userId, serverName);
|
||||
return {
|
||||
toolDefinitions: serverTools ? Object.keys(serverTools) : [],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
};
|
||||
});
|
||||
reinitMCPServer.mockImplementation(async ({ oauthStart }) => {
|
||||
await oauthStart(authorizationUrl);
|
||||
return { availableTools: { [mcpTool]: {} } };
|
||||
});
|
||||
|
||||
const result = await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
expect(result.toolDefinitions).toEqual([mcpTool]);
|
||||
expect(mockGetMCPServerTools).toHaveBeenCalledWith(req.user.id, serverName);
|
||||
expect(reinitMCPServer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
serverName,
|
||||
returnOnOAuth: false,
|
||||
oauthStart: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step_delta',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
delta: expect.objectContaining({
|
||||
auth: authorizationUrl,
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not join in-flight MCP initialization before replaying pending OAuth prompts', async () => {
|
||||
const serverName = 'Google-Workspace';
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `${Constants.mcp_all}${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockGetServerConfig.mockResolvedValue({
|
||||
type: 'streamable-http',
|
||||
url: 'https://demo.librechat.ai/mcp',
|
||||
requiresOAuth: true,
|
||||
});
|
||||
mockGetMCPServerTools.mockResolvedValue(null);
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
createdAt: Date.now(),
|
||||
metadata: { authorizationUrl },
|
||||
});
|
||||
mockLoadToolDefinitions.mockImplementation(async (params, deps) => {
|
||||
await deps.getOrFetchMCPServerTools(params.userId, serverName);
|
||||
return {
|
||||
toolDefinitions: [],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
};
|
||||
});
|
||||
reinitMCPServer.mockImplementation(async ({ oauthStart }) => {
|
||||
await oauthStart(authorizationUrl);
|
||||
return { availableTools: null };
|
||||
});
|
||||
|
||||
await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
expect(mockGetMCPServerTools).toHaveBeenCalledWith(req.user.id, serverName);
|
||||
expect(reinitMCPServer).toHaveBeenCalledTimes(1);
|
||||
expect(reinitMCPServer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
serverName,
|
||||
returnOnOAuth: false,
|
||||
oauthStart: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step_delta',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
delta: expect.objectContaining({
|
||||
auth: authorizationUrl,
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should re-emit pending MCP OAuth prompts when selected MCP tools are already concrete', async () => {
|
||||
const serverName = `Google${Constants.mcp_delimiter}Workspace`;
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
createdAt: Date.now(),
|
||||
metadata: { authorizationUrl },
|
||||
});
|
||||
mockLoadToolDefinitions.mockResolvedValue({
|
||||
toolDefinitions: [mcpTool],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
});
|
||||
reinitMCPServer.mockImplementation(async ({ oauthStart }) => {
|
||||
await oauthStart(authorizationUrl);
|
||||
return { availableTools: { [mcpTool]: {} } };
|
||||
});
|
||||
|
||||
const result = await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
expect(result.toolDefinitions).toEqual([mcpTool]);
|
||||
expect(mockGetMCPServerTools).not.toHaveBeenCalled();
|
||||
expect(reinitMCPServer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
serverName,
|
||||
returnOnOAuth: false,
|
||||
oauthStart: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step_delta',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
delta: expect.objectContaining({
|
||||
auth: authorizationUrl,
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should emit stored pending MCP OAuth prompts before waiting on a silent in-flight join', async () => {
|
||||
const serverName = 'Google-Workspace';
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
createdAt: Date.now(),
|
||||
metadata: { authorizationUrl },
|
||||
});
|
||||
mockLoadToolDefinitions.mockResolvedValue({
|
||||
toolDefinitions: [mcpTool],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
});
|
||||
reinitMCPServer.mockResolvedValue({ availableTools: null });
|
||||
|
||||
const result = await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
expect(result.toolDefinitions).toEqual([mcpTool]);
|
||||
expect(reinitMCPServer).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
serverName,
|
||||
returnOnOAuth: false,
|
||||
oauthStart: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step_delta',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
delta: expect.objectContaining({
|
||||
auth: authorizationUrl,
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should preserve OAuth URLs emitted while discovering MCP tools before a silent wait join', async () => {
|
||||
const serverName = 'Google-Workspace';
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockGetServerConfig.mockResolvedValue({
|
||||
type: 'streamable-http',
|
||||
url: 'https://demo.librechat.ai/mcp',
|
||||
requiresOAuth: true,
|
||||
});
|
||||
mockGetMCPServerTools.mockResolvedValue(null);
|
||||
mockFlowManager.getFlowState.mockResolvedValue(null);
|
||||
mockLoadToolDefinitions.mockImplementation(async (params, deps) => {
|
||||
await deps.getOrFetchMCPServerTools(params.userId, serverName);
|
||||
return {
|
||||
toolDefinitions: [],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
};
|
||||
});
|
||||
reinitMCPServer
|
||||
.mockImplementationOnce(async ({ oauthStart }) => {
|
||||
await oauthStart(authorizationUrl, { expiresAt: Date.now() + 60_000 });
|
||||
return { availableTools: null };
|
||||
})
|
||||
.mockResolvedValue({ availableTools: null });
|
||||
|
||||
await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
expect(reinitMCPServer).toHaveBeenCalledTimes(2);
|
||||
expect(mockSendEvent).toHaveBeenCalledWith(
|
||||
res,
|
||||
expect.objectContaining({
|
||||
event: 'on_run_step_delta',
|
||||
data: expect.objectContaining({
|
||||
id: `step_oauth_login_${serverName}`,
|
||||
delta: expect.objectContaining({
|
||||
auth: authorizationUrl,
|
||||
}),
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should preserve pending-flow expiry for OAuth URLs captured during discovery', async () => {
|
||||
const serverName = 'Google-Workspace';
|
||||
const authorizationUrl = 'https://auth.example.com/Google-Workspace';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
const capabilities = [AgentCapabilities.tools];
|
||||
const req = createMockReq(capabilities);
|
||||
const res = { writableEnded: false };
|
||||
const createdAt = Date.now() - 45_000;
|
||||
mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities));
|
||||
mockGetServerConfig.mockResolvedValue({
|
||||
type: 'streamable-http',
|
||||
url: 'https://demo.librechat.ai/mcp',
|
||||
requiresOAuth: true,
|
||||
});
|
||||
mockGetMCPServerTools.mockResolvedValue(null);
|
||||
mockFlowManager.getFlowState.mockResolvedValueOnce(null).mockResolvedValueOnce({
|
||||
status: 'PENDING',
|
||||
createdAt,
|
||||
metadata: { authorizationUrl },
|
||||
});
|
||||
mockLoadToolDefinitions.mockImplementation(async (params, deps) => {
|
||||
await deps.getOrFetchMCPServerTools(params.userId, serverName);
|
||||
return {
|
||||
toolDefinitions: [],
|
||||
toolRegistry: new Map(),
|
||||
hasDeferredTools: false,
|
||||
};
|
||||
});
|
||||
reinitMCPServer
|
||||
.mockImplementationOnce(async ({ oauthStart }) => {
|
||||
await oauthStart(authorizationUrl);
|
||||
return { availableTools: null };
|
||||
})
|
||||
.mockResolvedValue({ availableTools: null });
|
||||
|
||||
await loadAgentTools({
|
||||
req,
|
||||
res,
|
||||
agent: { id: 'agent_123', tools: [mcpTool] },
|
||||
definitionsOnly: true,
|
||||
});
|
||||
|
||||
const authDeltaEvent = mockSendEvent.mock.calls
|
||||
.map(([, event]) => event)
|
||||
.find((event) => event.data?.delta?.auth === authorizationUrl);
|
||||
expect(authDeltaEvent?.data.delta.expires_at).toBe(createdAt + PENDING_STALE_MS);
|
||||
});
|
||||
|
||||
it('should use request-scoped MCP config before falling back to the registry', async () => {
|
||||
const serverName = 'config-server';
|
||||
const mcpTool = `search${Constants.mcp_delimiter}${serverName}`;
|
||||
|
|
|
|||
|
|
@ -351,6 +351,12 @@ export type TOptions = {
|
|||
isResubmission?: boolean;
|
||||
/** Currently only utilized when `isResubmission === true`, uses that message's currently attached files */
|
||||
overrideFiles?: t.TMessage['files'];
|
||||
/**
|
||||
* Assistant message being regenerated. Used to derive the optimistic response
|
||||
* id for non-tail regenerations without accidentally keying the stream to the
|
||||
* conversation tail.
|
||||
*/
|
||||
targetResponseMessageId?: string | null;
|
||||
/**
|
||||
* Carry forward a user message's manually-invoked skills when the caller
|
||||
* is resubmitting / regenerating that same message — the compose-time
|
||||
|
|
|
|||
|
|
@ -73,8 +73,19 @@ export default function MCPServerStatusIcon({
|
|||
|
||||
const { connectionState } = serverStatus;
|
||||
|
||||
// Connecting: show spinner only (no action available)
|
||||
// Connecting: show spinner, with cancel when an OAuth flow is pending.
|
||||
if (connectionState === 'connecting') {
|
||||
if (canCancel) {
|
||||
return (
|
||||
<LoadingStatusIcon
|
||||
serverName={serverName}
|
||||
onConfigClick={onConfigClick}
|
||||
onCancel={onCancel}
|
||||
canCancel={canCancel}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return <ConnectingSpinner serverName={serverName} />;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -38,15 +38,33 @@ export default function ServerInitializationSection({
|
|||
|
||||
const serverStatus = connectionStatus?.[serverName];
|
||||
const isConnected = serverStatus?.connectionState === 'connected';
|
||||
const canCancel = isCancellable(serverName);
|
||||
const hasPendingOAuth =
|
||||
requiresOAuth &&
|
||||
serverStatus?.requiresOAuth === true &&
|
||||
serverStatus.connectionState === 'connecting';
|
||||
const canCancel = isCancellable(serverName) || hasPendingOAuth;
|
||||
const isServerInitializing = isInitializing(serverName);
|
||||
const serverOAuthUrl = getOAuthUrl(serverName);
|
||||
|
||||
const shouldShowReinit = isConnected && (requiresOAuth || hasCustomUserVars);
|
||||
const shouldShowInit = !isConnected && !serverOAuthUrl;
|
||||
const shouldShowInit = !isConnected && !serverOAuthUrl && !hasPendingOAuth;
|
||||
|
||||
if (!shouldShowReinit && !shouldShowInit && !serverOAuthUrl) {
|
||||
return null;
|
||||
if (!hasPendingOAuth) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<Button
|
||||
onClick={() => cancelOAuthFlow(serverName)}
|
||||
disabled={!canCancel}
|
||||
variant="outline"
|
||||
size={sidePanel ? 'sm' : 'default'}
|
||||
className="w-full"
|
||||
>
|
||||
{localize('com_ui_cancel')}
|
||||
</Button>
|
||||
);
|
||||
}
|
||||
|
||||
if (serverOAuthUrl) {
|
||||
|
|
|
|||
|
|
@ -187,7 +187,9 @@ export function shouldShowActionButton(statusIconProps?: MCPServerStatusIconProp
|
|||
|
||||
// Show for disconnected/error (can reconnect/configure)
|
||||
if (connectionState === 'disconnected' || connectionState === 'error') return true;
|
||||
// Don't show connecting spinner (no action)
|
||||
// Show a cancel action for pending OAuth connections.
|
||||
if (connectionState === 'connecting' && requiresOAuth && canCancel) return true;
|
||||
// Don't show connecting spinner when no action is available.
|
||||
if (connectionState === 'connecting') return false;
|
||||
// Connected: only show if there's config to manage
|
||||
if (connectionState === 'connected') return hasCustomUserVars || requiresOAuth;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,143 @@
|
|||
import { renderHook, act } from '@testing-library/react';
|
||||
import { EModelEndpoint } from 'librechat-data-provider';
|
||||
import type { TConversation, TMessage, TSubmission } from 'librechat-data-provider';
|
||||
import useChatFunctions from '../useChatFunctions';
|
||||
|
||||
const mockNavigate = jest.fn();
|
||||
const mockSetShowStopButton = jest.fn();
|
||||
const mockSetIsSubmitting = jest.fn();
|
||||
const mockGetEphemeralAgent = jest.fn(() => null);
|
||||
const mockSetFilesToDelete = jest.fn();
|
||||
const mockGetSender = jest.fn(() => 'Assistant');
|
||||
const mockGetExpiry = jest.fn(() => 'expiry-key');
|
||||
const mockGetQueryData = jest.fn(() => ({}));
|
||||
|
||||
jest.mock('react-router-dom', () => ({
|
||||
useNavigate: () => mockNavigate,
|
||||
}));
|
||||
|
||||
jest.mock('@tanstack/react-query', () => ({
|
||||
useQueryClient: () => ({
|
||||
getQueryData: mockGetQueryData,
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('recoil', () => ({
|
||||
useRecoilValue: () => false,
|
||||
useSetRecoilState: (atom: unknown) =>
|
||||
String(atom).includes('isSubmitting') ? mockSetIsSubmitting : mockSetShowStopButton,
|
||||
useRecoilCallback: (factory: any) =>
|
||||
factory({
|
||||
snapshot: {
|
||||
getLoadable: () => ({ state: 'hasValue', contents: [] }),
|
||||
},
|
||||
reset: jest.fn(),
|
||||
}),
|
||||
}));
|
||||
|
||||
jest.mock('~/hooks/Files/useSetFilesToDelete', () => () => mockSetFilesToDelete);
|
||||
jest.mock('~/hooks/Conversations/useGetSender', () => () => mockGetSender);
|
||||
jest.mock('~/hooks/Input/useUserKey', () => () => ({ getExpiry: mockGetExpiry }));
|
||||
jest.mock('~/hooks', () => ({
|
||||
useAuthContext: () => ({ user: null }),
|
||||
}));
|
||||
jest.mock('~/store', () => ({
|
||||
__esModule: true,
|
||||
default: {
|
||||
isTemporary: 'isTemporary',
|
||||
isSubmittingFamily: () => 'isSubmitting',
|
||||
showStopButtonByIndex: () => 'showStopButton',
|
||||
pendingManualSkillsByConvoId: () => 'pendingManualSkills',
|
||||
},
|
||||
useGetEphemeralAgent: () => mockGetEphemeralAgent,
|
||||
}));
|
||||
jest.mock('~/utils', () => ({
|
||||
logger: {
|
||||
log: jest.fn(),
|
||||
dir: jest.fn(),
|
||||
},
|
||||
createDualMessageContent: jest.fn(() => []),
|
||||
getRouteChatProjectId: jest.fn(() => null),
|
||||
}));
|
||||
|
||||
const userMessage = (messageId: string, parentMessageId = '00000000-0000-0000-0000-000000000000') =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
const assistantMessage = (messageId: string, parentMessageId: string) =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
describe('useChatFunctions regenerate', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockGetQueryData.mockReturnValue({});
|
||||
});
|
||||
|
||||
it('keys a non-tail regenerate to the selected assistant response', () => {
|
||||
let messages = [
|
||||
userMessage('user-1'),
|
||||
assistantMessage('assistant-1', 'user-1'),
|
||||
userMessage('user-2', 'assistant-1'),
|
||||
assistantMessage('assistant-2', 'user-2'),
|
||||
userMessage('user-3', 'assistant-2'),
|
||||
assistantMessage('assistant-3', 'user-3'),
|
||||
];
|
||||
const setMessages = jest.fn((nextMessages: TMessage[]) => {
|
||||
messages = nextMessages;
|
||||
});
|
||||
const setSubmission = jest.fn();
|
||||
const conversation = {
|
||||
conversationId: 'conversation-1',
|
||||
endpoint: EModelEndpoint.agents,
|
||||
model: 'gpt-4o',
|
||||
agent_id: 'agent-1',
|
||||
} as TConversation;
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useChatFunctions({
|
||||
isSubmitting: false,
|
||||
latestMessage: messages[5],
|
||||
conversation,
|
||||
getMessages: () => messages,
|
||||
setMessages,
|
||||
setSubmission,
|
||||
}),
|
||||
);
|
||||
|
||||
act(() => {
|
||||
result.current.regenerate(messages[1]);
|
||||
});
|
||||
|
||||
const submission = setSubmission.mock.calls.at(-1)?.[0] as TSubmission;
|
||||
expect(submission.userMessage.overrideParentMessageId).toBe('user-1');
|
||||
expect(submission.userMessage.responseMessageId).toBe('assistant-1_');
|
||||
expect(submission.initialResponse?.messageId).toBe('assistant-1_');
|
||||
expect(submission.initialResponse?.parentMessageId).toBe('user-1');
|
||||
expect(submission.messages.map((message) => message.messageId)).toEqual(['user-1']);
|
||||
expect(submission.regenerateMessages?.map((message) => message.messageId)).toEqual([
|
||||
'user-1',
|
||||
'assistant-1',
|
||||
'user-2',
|
||||
'assistant-2',
|
||||
'user-3',
|
||||
'assistant-3',
|
||||
]);
|
||||
expect(
|
||||
setMessages.mock.calls.at(-1)?.[0].map((message: TMessage) => message.messageId),
|
||||
).toEqual(['user-1', 'assistant-1_']);
|
||||
expect(messages.at(-1)?.messageId).toBe('assistant-1_');
|
||||
});
|
||||
});
|
||||
93
client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts
Normal file
93
client/src/hooks/Chat/__tests__/useChatFunctions.spec.ts
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
import type { TMessage } from 'librechat-data-provider';
|
||||
import {
|
||||
getRegenerateSubmissionMessages,
|
||||
getPreliminaryRegenerateResponseMessageId,
|
||||
getRegenerateTargetResponseMessage,
|
||||
} from '../useChatFunctions';
|
||||
|
||||
const userMessage = (messageId: string, parentMessageId = '00000000-0000-0000-0000-000000000000') =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
const assistantMessage = (messageId: string, parentMessageId: string) =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
describe('regenerate response targeting', () => {
|
||||
it('uses the clicked assistant response instead of the conversation tail', () => {
|
||||
const messages = [
|
||||
userMessage('user-1'),
|
||||
assistantMessage('assistant-1', 'user-1'),
|
||||
userMessage('user-2', 'assistant-1'),
|
||||
assistantMessage('assistant-2', 'user-2'),
|
||||
userMessage('user-3', 'assistant-2'),
|
||||
assistantMessage('assistant-3', 'user-3'),
|
||||
];
|
||||
|
||||
const targetResponse = getRegenerateTargetResponseMessage({
|
||||
messages,
|
||||
parentMessageId: 'user-1',
|
||||
targetResponseMessageId: 'assistant-1',
|
||||
latestMessage: messages[5],
|
||||
});
|
||||
|
||||
expect(targetResponse?.messageId).toBe('assistant-1');
|
||||
expect(getPreliminaryRegenerateResponseMessageId(targetResponse?.messageId)).toBe(
|
||||
'assistant-1_',
|
||||
);
|
||||
});
|
||||
|
||||
it('only falls back to latestMessage when it belongs to the regenerated user turn', () => {
|
||||
const messages = [
|
||||
userMessage('user-1'),
|
||||
assistantMessage('assistant-1', 'user-1'),
|
||||
userMessage('user-2', 'assistant-1'),
|
||||
assistantMessage('assistant-2', 'user-2'),
|
||||
];
|
||||
|
||||
expect(
|
||||
getRegenerateTargetResponseMessage({
|
||||
messages,
|
||||
parentMessageId: 'user-1',
|
||||
latestMessage: messages[3],
|
||||
})?.messageId,
|
||||
).toBe('assistant-1');
|
||||
|
||||
expect(
|
||||
getRegenerateTargetResponseMessage({
|
||||
messages,
|
||||
parentMessageId: 'user-2',
|
||||
latestMessage: messages[3],
|
||||
})?.messageId,
|
||||
).toBe('assistant-2');
|
||||
});
|
||||
|
||||
it('truncates regenerate history before the targeted assistant response', () => {
|
||||
const messages = [
|
||||
userMessage('user-1'),
|
||||
assistantMessage('assistant-1', 'user-1'),
|
||||
userMessage('user-2', 'assistant-1'),
|
||||
assistantMessage('assistant-2', 'user-2'),
|
||||
userMessage('user-3', 'assistant-2'),
|
||||
assistantMessage('assistant-3', 'user-3'),
|
||||
];
|
||||
|
||||
expect(
|
||||
getRegenerateSubmissionMessages({
|
||||
messages,
|
||||
targetResponseMessage: messages[1],
|
||||
initialResponseId: 'assistant-1_',
|
||||
}).map((message) => message.messageId),
|
||||
).toEqual(['user-1']);
|
||||
});
|
||||
});
|
||||
|
|
@ -70,6 +70,88 @@ const getAppendParentMessageId = ({
|
|||
return failedUserMessage.parentMessageId ?? Constants.NO_PARENT;
|
||||
};
|
||||
|
||||
type RegenerateTargetResponseArgs = {
|
||||
messages: TMessage[];
|
||||
parentMessageId?: string | null;
|
||||
targetResponseMessageId?: string | null;
|
||||
latestMessage?: TMessage | null;
|
||||
};
|
||||
|
||||
const isAssistantResponseForParent = (
|
||||
message: TMessage | null | undefined,
|
||||
parentMessageId?: string | null,
|
||||
): message is TMessage =>
|
||||
!!message &&
|
||||
!message.isCreatedByUser &&
|
||||
!!parentMessageId &&
|
||||
message.parentMessageId === parentMessageId;
|
||||
|
||||
export function getPreliminaryRegenerateResponseMessageId(
|
||||
responseMessageId?: string | null,
|
||||
): string | null {
|
||||
if (typeof responseMessageId !== 'string' || responseMessageId.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return `${responseMessageId.replace(/_+$/, '')}_`;
|
||||
}
|
||||
|
||||
export function getRegenerateTargetResponseMessage({
|
||||
messages,
|
||||
parentMessageId,
|
||||
targetResponseMessageId,
|
||||
latestMessage,
|
||||
}: RegenerateTargetResponseArgs): TMessage | null {
|
||||
if (!parentMessageId) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (targetResponseMessageId) {
|
||||
const targetResponse = messages.find(
|
||||
(message) =>
|
||||
message.messageId === targetResponseMessageId &&
|
||||
isAssistantResponseForParent(message, parentMessageId),
|
||||
);
|
||||
if (targetResponse) {
|
||||
return targetResponse;
|
||||
}
|
||||
}
|
||||
|
||||
if (isAssistantResponseForParent(latestMessage, parentMessageId)) {
|
||||
return latestMessage;
|
||||
}
|
||||
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const message = messages[i];
|
||||
if (isAssistantResponseForParent(message, parentMessageId)) {
|
||||
return message;
|
||||
}
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
export function getRegenerateSubmissionMessages({
|
||||
messages,
|
||||
targetResponseMessage,
|
||||
initialResponseId,
|
||||
}: {
|
||||
messages: TMessage[];
|
||||
targetResponseMessage?: TMessage | null;
|
||||
initialResponseId?: string | null;
|
||||
}): TMessage[] {
|
||||
if (targetResponseMessage?.messageId) {
|
||||
const targetIndex = messages.findIndex(
|
||||
(message) => message.messageId === targetResponseMessage.messageId,
|
||||
);
|
||||
if (targetIndex >= 0) {
|
||||
return messages.slice(0, targetIndex);
|
||||
}
|
||||
}
|
||||
|
||||
return messages.filter((msg) => msg.messageId !== initialResponseId);
|
||||
}
|
||||
|
||||
export default function useChatFunctions({
|
||||
index = 0,
|
||||
files,
|
||||
|
|
@ -144,6 +226,7 @@ export default function useChatFunctions({
|
|||
isEdited = false,
|
||||
overrideMessages,
|
||||
overrideFiles,
|
||||
targetResponseMessageId,
|
||||
overrideManualSkills,
|
||||
addedConvo,
|
||||
} = {},
|
||||
|
|
@ -244,6 +327,14 @@ export default function useChatFunctions({
|
|||
const targetParentMessage = currentMessages.find(
|
||||
(msg) => msg.messageId === targetParentMessageId,
|
||||
);
|
||||
const targetResponseMessage = isRegenerate
|
||||
? getRegenerateTargetResponseMessage({
|
||||
messages: currentMessages,
|
||||
parentMessageId: messageId,
|
||||
targetResponseMessageId,
|
||||
latestMessage,
|
||||
})
|
||||
: null;
|
||||
|
||||
let thread_id = targetParentMessage?.thread_id ?? latestMessage?.thread_id;
|
||||
if (thread_id == null) {
|
||||
|
|
@ -327,8 +418,10 @@ export default function useChatFunctions({
|
|||
|
||||
const responseMessageId =
|
||||
editedMessageId ??
|
||||
(latestMessage?.messageId && isRegenerate
|
||||
? latestMessage.messageId.replace(/_+$/, '') + '_'
|
||||
(isRegenerate
|
||||
? getPreliminaryRegenerateResponseMessageId(
|
||||
targetResponseMessage?.messageId ?? targetResponseMessageId,
|
||||
)
|
||||
: null) ??
|
||||
null;
|
||||
const initialResponseId =
|
||||
|
|
@ -407,6 +500,15 @@ export default function useChatFunctions({
|
|||
currentMessages = currentMessages.filter((msg) => msg.messageId !== responseMessageId);
|
||||
}
|
||||
|
||||
const submissionMessages = isRegenerate
|
||||
? getRegenerateSubmissionMessages({
|
||||
messages: currentMessages,
|
||||
targetResponseMessage,
|
||||
initialResponseId: initialResponse.messageId,
|
||||
})
|
||||
: currentMessages;
|
||||
const regenerateMessages = isRegenerate ? [...currentMessages] : undefined;
|
||||
|
||||
logger.log('message_state', initialResponse);
|
||||
const submission: TSubmission = {
|
||||
conversation: {
|
||||
|
|
@ -420,7 +522,8 @@ export default function useChatFunctions({
|
|||
responseMessageId,
|
||||
overrideParentMessageId: isRegenerate ? messageId : null,
|
||||
},
|
||||
messages: currentMessages,
|
||||
messages: submissionMessages,
|
||||
regenerateMessages,
|
||||
isEdited: isEditOrContinue,
|
||||
isContinued,
|
||||
isRegenerate,
|
||||
|
|
@ -433,17 +536,24 @@ export default function useChatFunctions({
|
|||
};
|
||||
|
||||
if (isRegenerate) {
|
||||
setMessages([...submission.messages, initialResponse]);
|
||||
setMessages([...submissionMessages, initialResponse]);
|
||||
} else {
|
||||
setMessages([...submission.messages, currentMsg, initialResponse]);
|
||||
setMessages([...submissionMessages, currentMsg, initialResponse]);
|
||||
}
|
||||
|
||||
setSubmission(submission);
|
||||
logger.dir('message_stream', submission, { depth: null });
|
||||
};
|
||||
|
||||
const regenerate = ({ parentMessageId }, options?: { addedConvo?: TConversation | null }) => {
|
||||
const regenerate = (
|
||||
message: Partial<Pick<TMessage, 'messageId' | 'parentMessageId' | 'isCreatedByUser'>>,
|
||||
options?: { addedConvo?: TConversation | null },
|
||||
) => {
|
||||
const messages = getMessages();
|
||||
const parentMessageId =
|
||||
message.isCreatedByUser === true ? message.messageId : message.parentMessageId;
|
||||
const targetResponseMessageId =
|
||||
message.isCreatedByUser === true ? undefined : message.messageId;
|
||||
const parentMessage = messages?.find((element) => element.messageId == parentMessageId);
|
||||
|
||||
if (parentMessage && parentMessage.isCreatedByUser) {
|
||||
|
|
@ -452,6 +562,7 @@ export default function useChatFunctions({
|
|||
{
|
||||
isRegenerate: true,
|
||||
addedConvo: options?.addedConvo ?? undefined,
|
||||
targetResponseMessageId,
|
||||
/** Carry the original user message's manual skill picks forward
|
||||
* so the regenerated response is primed with the same skills.
|
||||
* The compose-time atom was drained on the first submit; without
|
||||
|
|
|
|||
|
|
@ -579,6 +579,9 @@ export function useMCPServerManager({
|
|||
|
||||
const hasCustomUserVars =
|
||||
serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0;
|
||||
const hasPendingOAuth =
|
||||
serverStatus?.requiresOAuth === true && serverStatus.connectionState === 'connecting';
|
||||
const canCancelOAuth = isCancellable(serverName) || hasPendingOAuth;
|
||||
|
||||
return {
|
||||
serverName,
|
||||
|
|
@ -592,8 +595,8 @@ export function useMCPServerManager({
|
|||
} as TPlugin)
|
||||
: undefined,
|
||||
onConfigClick: handleConfigClick,
|
||||
isInitializing: isInitializing(serverName),
|
||||
canCancel: isCancellable(serverName),
|
||||
isInitializing: isInitializing(serverName) || hasPendingOAuth,
|
||||
canCancel: canCancelOAuth,
|
||||
onCancel: handleCancelClick,
|
||||
hasCustomUserVars,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
import React from 'react';
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import { RecoilRoot, type MutableSnapshot } from 'recoil';
|
||||
import { MemoryRouter, Route, Routes } from 'react-router-dom';
|
||||
import { act, renderHook } from '@testing-library/react';
|
||||
import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
|
||||
import { QueryKeys, type TConversation, type TMessage } from 'librechat-data-provider';
|
||||
import { getBranchSiblingIndexesForTarget, getMessageBranchSiblingParentIds } from '~/utils';
|
||||
import { useLatestMessage, useLatestMessageId } from '~/hooks/Messages/useLatestMessage';
|
||||
import { getMessageBranchSiblingParentIds } from '~/utils';
|
||||
import store from '~/store';
|
||||
|
||||
function createQueryClient() {
|
||||
|
|
@ -254,3 +254,26 @@ describe('getMessageBranchSiblingParentIds', () => {
|
|||
).toEqual([userMessage.messageId]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getBranchSiblingIndexesForTarget', () => {
|
||||
it('returns sibling indexes that select the branch containing the target message', () => {
|
||||
expect(
|
||||
getBranchSiblingIndexesForTarget(
|
||||
[
|
||||
userMessage,
|
||||
olderAssistantMessage,
|
||||
olderFollowUpUserMessage,
|
||||
olderFollowUpAssistantMessage,
|
||||
assistantMessage,
|
||||
],
|
||||
olderFollowUpAssistantMessage.messageId,
|
||||
conversation.conversationId,
|
||||
),
|
||||
).toEqual([
|
||||
{
|
||||
parentMessageId: userMessage.messageId,
|
||||
siblingIdx: 1,
|
||||
},
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
175
client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts
Normal file
175
client/src/hooks/SSE/__tests__/useEventHandlers.spec.ts
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
import { Constants } from 'librechat-data-provider';
|
||||
import type { EventSubmission, TMessage } from 'librechat-data-provider';
|
||||
import {
|
||||
buildCreatedInitialResponse,
|
||||
getExistingConversationAbortMessages,
|
||||
isInitialNewConversationSubmission,
|
||||
mergeRegenerateFinalMessages,
|
||||
} from '~/hooks/SSE/useEventHandlers';
|
||||
|
||||
describe('buildCreatedInitialResponse', () => {
|
||||
const userMessage = {
|
||||
messageId: 'server-user-message',
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: true,
|
||||
text: 'Hello',
|
||||
sender: 'User',
|
||||
} as TMessage;
|
||||
|
||||
const initialResponse = {
|
||||
messageId: 'prelim-response',
|
||||
parentMessageId: 'original-user-message',
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: false,
|
||||
text: '',
|
||||
sender: 'Assistant',
|
||||
} as TMessage;
|
||||
|
||||
it('uses the created user message id for new turns', () => {
|
||||
expect(
|
||||
buildCreatedInitialResponse({
|
||||
initialResponse,
|
||||
userMessage,
|
||||
isRegenerate: false,
|
||||
}),
|
||||
).toEqual(
|
||||
expect.objectContaining({
|
||||
messageId: 'server-user-message_',
|
||||
parentMessageId: 'server-user-message',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('preserves the regenerated prelim response id and parent', () => {
|
||||
expect(
|
||||
buildCreatedInitialResponse({
|
||||
initialResponse,
|
||||
userMessage,
|
||||
isRegenerate: true,
|
||||
}),
|
||||
).toEqual(
|
||||
expect.objectContaining({
|
||||
messageId: 'prelim-response',
|
||||
parentMessageId: 'original-user-message',
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isInitialNewConversationSubmission', () => {
|
||||
it('treats a root user message as an optimistic new chat', () => {
|
||||
expect(
|
||||
isInitialNewConversationSubmission({
|
||||
userMessage: {
|
||||
messageId: 'user-1',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
} as TMessage,
|
||||
} as EventSubmission),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('does not treat follow-up turns as optimistic new chats', () => {
|
||||
expect(
|
||||
isInitialNewConversationSubmission({
|
||||
userMessage: {
|
||||
messageId: 'user-2',
|
||||
parentMessageId: 'assistant-1',
|
||||
} as TMessage,
|
||||
} as EventSubmission),
|
||||
).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('mergeRegenerateFinalMessages', () => {
|
||||
const userMessage = (messageId: string, parentMessageId: string = Constants.NO_PARENT) =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
const assistantMessage = (messageId: string, parentMessageId: string) =>
|
||||
({
|
||||
messageId,
|
||||
parentMessageId,
|
||||
conversationId: 'conversation-1',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
it('keeps the original branch siblings when a non-tail regenerate finalizes', () => {
|
||||
const rootUser = userMessage('user-1');
|
||||
const originalResponse = assistantMessage('assistant-1', rootUser.messageId);
|
||||
const followUpUser = userMessage('user-2', originalResponse.messageId);
|
||||
const followUpResponse = assistantMessage('assistant-2', followUpUser.messageId);
|
||||
const finalResponse = assistantMessage('assistant-3', rootUser.messageId);
|
||||
|
||||
expect(
|
||||
mergeRegenerateFinalMessages({
|
||||
messages: [rootUser, originalResponse, followUpUser, followUpResponse],
|
||||
responseMessage: finalResponse,
|
||||
initialResponseId: 'assistant-1_',
|
||||
}).map((message) => message.messageId),
|
||||
).toEqual([
|
||||
rootUser.messageId,
|
||||
originalResponse.messageId,
|
||||
followUpUser.messageId,
|
||||
followUpResponse.messageId,
|
||||
finalResponse.messageId,
|
||||
]);
|
||||
});
|
||||
|
||||
it('replaces the streamed preliminary response when it is present', () => {
|
||||
const rootUser = userMessage('user-1');
|
||||
const preliminaryResponse = assistantMessage('assistant-1_', rootUser.messageId);
|
||||
const finalResponse = assistantMessage('assistant-3', rootUser.messageId);
|
||||
|
||||
expect(
|
||||
mergeRegenerateFinalMessages({
|
||||
messages: [rootUser, preliminaryResponse],
|
||||
responseMessage: finalResponse,
|
||||
initialResponseId: preliminaryResponse.messageId,
|
||||
}).map((message) => message.messageId),
|
||||
).toEqual([rootUser.messageId, finalResponse.messageId]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('getExistingConversationAbortMessages', () => {
|
||||
const message = (messageId: string) =>
|
||||
({
|
||||
messageId,
|
||||
conversationId: 'conversation-1',
|
||||
text: messageId,
|
||||
}) as TMessage;
|
||||
|
||||
it('restores the full pre-regenerate branch on early abort', () => {
|
||||
const originalMessages = [message('user-1'), message('assistant-1'), message('user-2')];
|
||||
const scopedRegenerateMessages = [message('user-1')];
|
||||
const currentStreamMessages = [message('user-1'), message('assistant-1_')];
|
||||
|
||||
expect(
|
||||
getExistingConversationAbortMessages({
|
||||
messages: scopedRegenerateMessages,
|
||||
currentMessages: currentStreamMessages,
|
||||
regenerateMessages: originalMessages,
|
||||
isRegenerate: true,
|
||||
}).map(({ messageId }) => messageId),
|
||||
).toEqual(['user-1', 'assistant-1', 'user-2']);
|
||||
});
|
||||
|
||||
it('keeps the existing non-regenerate abort rollback behavior', () => {
|
||||
const submissionMessages = [message('user-1')];
|
||||
const currentMessages = [message('user-1'), message('assistant-1')];
|
||||
|
||||
expect(
|
||||
getExistingConversationAbortMessages({
|
||||
messages: submissionMessages,
|
||||
currentMessages,
|
||||
}).map(({ messageId }) => messageId),
|
||||
).toEqual(['user-1']);
|
||||
});
|
||||
});
|
||||
|
|
@ -1,6 +1,12 @@
|
|||
import { renderHook, act, waitFor } from '@testing-library/react';
|
||||
import { Constants, LocalStorageKeys, QueryKeys, request } from 'librechat-data-provider';
|
||||
import type { TSubmission } from 'librechat-data-provider';
|
||||
import {
|
||||
Constants,
|
||||
LocalStorageKeys,
|
||||
QueryKeys,
|
||||
StepEvents,
|
||||
request,
|
||||
} from 'librechat-data-provider';
|
||||
import type { TMessage, TSubmission } from 'librechat-data-provider';
|
||||
|
||||
type SSEEventListener = (e: Partial<MessageEvent> & { responseCode?: number }) => void;
|
||||
|
||||
|
|
@ -107,6 +113,7 @@ jest.mock('~/data-provider', () => ({
|
|||
}));
|
||||
|
||||
const mockErrorHandler = jest.fn();
|
||||
const mockFinalHandler = jest.fn();
|
||||
const mockCreatedHandler = jest.fn();
|
||||
const mockStepHandler = jest.fn();
|
||||
const mockTitleHandler = jest.fn();
|
||||
|
|
@ -116,7 +123,7 @@ const mockClearStepMaps = jest.fn();
|
|||
jest.mock('~/hooks/SSE/useEventHandlers', () =>
|
||||
jest.fn(() => ({
|
||||
errorHandler: mockErrorHandler,
|
||||
finalHandler: jest.fn(),
|
||||
finalHandler: mockFinalHandler,
|
||||
createdHandler: mockCreatedHandler,
|
||||
attachmentHandler: jest.fn(),
|
||||
stepHandler: mockStepHandler,
|
||||
|
|
@ -156,7 +163,7 @@ const CONV_ID = 'conv-abc-123';
|
|||
type PartialSubmission = {
|
||||
conversation: { conversationId?: string };
|
||||
userMessage: Record<string, unknown>;
|
||||
messages: never[];
|
||||
messages: TMessage[];
|
||||
isTemporary: boolean;
|
||||
initialResponse: Record<string, unknown>;
|
||||
endpointOption: { endpoint: string };
|
||||
|
|
@ -190,7 +197,7 @@ const buildSubmission = (overrides: Partial<PartialSubmission> = {}): TSubmissio
|
|||
|
||||
const buildChatHelpers = () => ({
|
||||
setMessages: jest.fn(),
|
||||
getMessages: jest.fn(() => []),
|
||||
getMessages: jest.fn<TMessage[], []>(() => []),
|
||||
setConversation: jest.fn(),
|
||||
setIsSubmitting: mockSetIsSubmitting,
|
||||
newConversation: jest.fn(),
|
||||
|
|
@ -229,6 +236,7 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
mockSSEInstances.length = 0;
|
||||
localStorage.clear();
|
||||
mockErrorHandler.mockClear();
|
||||
mockFinalHandler.mockClear();
|
||||
mockCreatedHandler.mockClear();
|
||||
mockStepHandler.mockClear();
|
||||
mockTitleHandler.mockClear();
|
||||
|
|
@ -446,6 +454,46 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
unmount();
|
||||
});
|
||||
|
||||
it('replaces the new-chat URL when the stream id is known despite a stale parent id', async () => {
|
||||
window.history.pushState({}, '', '/c/new');
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: 'stale-parent-message',
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(window.location.pathname).toBe('/c/stream-123');
|
||||
expect(mockSetQueryData).toHaveBeenCalledWith(
|
||||
[QueryKeys.messages, 'stream-123'],
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ messageId: 'msg-1', conversationId: 'stream-123' }),
|
||||
]),
|
||||
);
|
||||
|
||||
unmount();
|
||||
window.history.pushState({}, '', '/');
|
||||
});
|
||||
|
||||
it('hydrates the submission conversation id before created handlers run', async () => {
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
|
|
@ -496,6 +544,193 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
unmount();
|
||||
});
|
||||
|
||||
it('queues run step events until created hydrates the submission', async () => {
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const runStepEvent = {
|
||||
event: StepEvents.ON_RUN_STEP,
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
runId: 'msg-1_',
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'Google-Workspace', args: '' }],
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', { data: JSON.stringify(runStepEvent) });
|
||||
});
|
||||
|
||||
expect(mockStepHandler).not.toHaveBeenCalled();
|
||||
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
created: true,
|
||||
message: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockStepHandler).toHaveBeenCalledWith(
|
||||
runStepEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('renders OAuth run step events before created while retaining replay after hydration', async () => {
|
||||
const submission = buildSubmission({
|
||||
conversation: {},
|
||||
userMessage: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: null,
|
||||
text: 'Hello',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'msg-1_',
|
||||
conversationId: null,
|
||||
text: '',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const runStepEvent = {
|
||||
event: StepEvents.ON_RUN_STEP,
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
},
|
||||
},
|
||||
};
|
||||
const runStepDeltaEvent = {
|
||||
event: StepEvents.ON_RUN_STEP_DELTA,
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', { data: JSON.stringify(runStepEvent) });
|
||||
sse._emit('message', { data: JSON.stringify(runStepDeltaEvent) });
|
||||
});
|
||||
|
||||
expect(mockStepHandler).toHaveBeenCalledTimes(2);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
runStepEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
runStepDeltaEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
created: true,
|
||||
message: {
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockStepHandler).toHaveBeenCalledTimes(4);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
3,
|
||||
runStepEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
4,
|
||||
runStepDeltaEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({
|
||||
messageId: 'msg-1',
|
||||
conversationId: 'stream-123',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('routes title stream events to the title handler', async () => {
|
||||
const submission = buildSubmission();
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
|
@ -637,6 +872,447 @@ describe('useResumableSSE - 404 error path', () => {
|
|||
unmount();
|
||||
});
|
||||
|
||||
it('replays OAuth run step delta events from resume state sync', async () => {
|
||||
const submission = buildSubmission();
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const runStep = {
|
||||
id: 'step-oauth',
|
||||
runId: 'resp-1',
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'Google-Workspace', args: '' }],
|
||||
},
|
||||
};
|
||||
const replayEvent = {
|
||||
event: StepEvents.ON_RUN_STEP_DELTA,
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ name: 'Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
sync: true,
|
||||
resumeState: {
|
||||
runSteps: [runStep],
|
||||
replayEvents: [replayEvent],
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ event: StepEvents.ON_RUN_STEP, data: runStep },
|
||||
expect.objectContaining({ userMessage: expect.objectContaining({ messageId: 'msg-1' }) }),
|
||||
);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
replayEvent,
|
||||
expect.objectContaining({ userMessage: expect.objectContaining({ messageId: 'msg-1' }) }),
|
||||
);
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('anchors preliminary OAuth replay events to the response message from resume state sync', async () => {
|
||||
const submission = buildSubmission({
|
||||
userMessage: {
|
||||
messageId: 'original-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original prompt',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: Constants.NO_PARENT,
|
||||
},
|
||||
initialResponse: {
|
||||
messageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original response',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
parentMessageId: 'original-user',
|
||||
},
|
||||
});
|
||||
const chatHelpers = buildChatHelpers();
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const runStep = {
|
||||
id: 'step-oauth',
|
||||
runId: Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
},
|
||||
};
|
||||
const replayEvent = {
|
||||
event: StepEvents.ON_RUN_STEP_DELTA,
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
sync: true,
|
||||
resumeState: {
|
||||
runSteps: [runStep],
|
||||
replayEvents: [replayEvent],
|
||||
responseMessageId: 'follow-up-response',
|
||||
conversationId: CONV_ID,
|
||||
userMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
{ event: StepEvents.ON_RUN_STEP, data: runStep },
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({ messageId: 'follow-up-user' }),
|
||||
initialResponse: expect.objectContaining({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
expect(mockStepHandler).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
replayEvent,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({ messageId: 'follow-up-user' }),
|
||||
initialResponse: expect.objectContaining({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('merges resumed user and response messages into loaded conversation history', async () => {
|
||||
const originalUser = {
|
||||
messageId: 'original-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original prompt',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: '00000000-0000-0000-0000-000000000000',
|
||||
};
|
||||
const originalResponse = {
|
||||
messageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original response',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
parentMessageId: 'original-user',
|
||||
};
|
||||
const submission = {
|
||||
...buildSubmission({
|
||||
conversation: { conversationId: CONV_ID },
|
||||
userMessage: originalUser,
|
||||
initialResponse: originalResponse,
|
||||
}),
|
||||
resumeStreamId: CONV_ID,
|
||||
} as TSubmission & { resumeStreamId: string };
|
||||
const chatHelpers = buildChatHelpers();
|
||||
chatHelpers.getMessages.mockReturnValue([originalUser, originalResponse]);
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
sync: true,
|
||||
resumeState: {
|
||||
runSteps: [],
|
||||
replayEvents: [],
|
||||
aggregatedContent: [
|
||||
{
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
id: 'call-oauth',
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
args: '',
|
||||
},
|
||||
},
|
||||
],
|
||||
responseMessageId: 'follow-up-response',
|
||||
conversationId: CONV_ID,
|
||||
userMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
const lastMessages = chatHelpers.setMessages.mock.calls.at(-1)?.[0];
|
||||
expect(lastMessages.map((message: { messageId: string }) => message.messageId)).toEqual([
|
||||
'original-user',
|
||||
'original-response',
|
||||
'follow-up-user',
|
||||
'follow-up-response',
|
||||
]);
|
||||
expect(lastMessages[2]).toEqual(
|
||||
expect.objectContaining({
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
text: 'Follow-up prompt',
|
||||
}),
|
||||
);
|
||||
expect(lastMessages[3]).toEqual(
|
||||
expect.objectContaining({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
content: expect.any(Array),
|
||||
}),
|
||||
);
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('uses the resumed submission for final events after sync', async () => {
|
||||
const originalUser = {
|
||||
messageId: 'original-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original prompt',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: String(Constants.NO_PARENT),
|
||||
};
|
||||
const originalResponse = {
|
||||
messageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original response',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
parentMessageId: 'original-user',
|
||||
};
|
||||
const submission = {
|
||||
...buildSubmission({
|
||||
conversation: { conversationId: CONV_ID },
|
||||
userMessage: originalUser,
|
||||
initialResponse: originalResponse,
|
||||
}),
|
||||
resumeStreamId: CONV_ID,
|
||||
} as TSubmission & { resumeStreamId: string };
|
||||
const chatHelpers = buildChatHelpers();
|
||||
chatHelpers.getMessages.mockReturnValue([originalUser, originalResponse]);
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const sse = getLastSSE();
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify({
|
||||
sync: true,
|
||||
resumeState: {
|
||||
runSteps: [],
|
||||
replayEvents: [],
|
||||
responseMessageId: 'follow-up-response',
|
||||
conversationId: CONV_ID,
|
||||
userMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
const finalPayload = {
|
||||
final: true,
|
||||
conversation: { conversationId: CONV_ID },
|
||||
requestMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
responseMessage: {
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Done',
|
||||
isCreatedByUser: false,
|
||||
},
|
||||
};
|
||||
|
||||
await act(async () => {
|
||||
sse._emit('message', {
|
||||
data: JSON.stringify(finalPayload),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockFinalHandler).toHaveBeenCalledWith(
|
||||
finalPayload,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({ messageId: 'follow-up-user' }),
|
||||
initialResponse: expect.objectContaining({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it('uses the resumed submission for final events after reconnecting from sync', async () => {
|
||||
jest.useFakeTimers();
|
||||
const originalUser = {
|
||||
messageId: 'original-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original prompt',
|
||||
isCreatedByUser: true,
|
||||
sender: 'User',
|
||||
parentMessageId: String(Constants.NO_PARENT),
|
||||
};
|
||||
const originalResponse = {
|
||||
messageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Original response',
|
||||
isCreatedByUser: false,
|
||||
sender: 'Assistant',
|
||||
parentMessageId: 'original-user',
|
||||
};
|
||||
const submission = {
|
||||
...buildSubmission({
|
||||
conversation: { conversationId: CONV_ID },
|
||||
userMessage: originalUser,
|
||||
initialResponse: originalResponse,
|
||||
}),
|
||||
resumeStreamId: CONV_ID,
|
||||
} as TSubmission & { resumeStreamId: string };
|
||||
const chatHelpers = buildChatHelpers();
|
||||
chatHelpers.getMessages.mockReturnValue([originalUser, originalResponse]);
|
||||
|
||||
const { unmount } = renderHook(() => useResumableSSE(submission, chatHelpers));
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
const initialSSE = getLastSSE();
|
||||
await act(async () => {
|
||||
initialSSE._emit('message', {
|
||||
data: JSON.stringify({
|
||||
sync: true,
|
||||
resumeState: {
|
||||
runSteps: [],
|
||||
replayEvents: [],
|
||||
responseMessageId: 'follow-up-response',
|
||||
conversationId: CONV_ID,
|
||||
userMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
},
|
||||
},
|
||||
}),
|
||||
});
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
initialSSE._emit('error');
|
||||
});
|
||||
await advanceRetryTimer(1000);
|
||||
|
||||
expect(mockSSEInstances).toHaveLength(2);
|
||||
const reconnectedSSE = getLastSSE();
|
||||
const finalPayload = {
|
||||
final: true,
|
||||
conversation: { conversationId: CONV_ID },
|
||||
requestMessage: {
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: 'original-response',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Follow-up prompt',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
responseMessage: {
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
conversationId: CONV_ID,
|
||||
text: 'Done',
|
||||
isCreatedByUser: false,
|
||||
},
|
||||
};
|
||||
|
||||
await act(async () => {
|
||||
reconnectedSSE._emit('message', {
|
||||
data: JSON.stringify(finalPayload),
|
||||
});
|
||||
});
|
||||
|
||||
expect(mockFinalHandler).toHaveBeenLastCalledWith(
|
||||
finalPayload,
|
||||
expect.objectContaining({
|
||||
userMessage: expect.objectContaining({ messageId: 'follow-up-user' }),
|
||||
initialResponse: expect.objectContaining({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: 'follow-up-user',
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
unmount();
|
||||
});
|
||||
|
||||
it.each([undefined, 500, 503])(
|
||||
'does not call errorHandler for responseCode %s (reconnect path)',
|
||||
async (responseCode) => {
|
||||
|
|
|
|||
|
|
@ -67,6 +67,8 @@ function renderUseResumeOnLoad({
|
|||
conversationId = CONVERSATION_ID,
|
||||
messagesLoaded = true,
|
||||
onSubmission,
|
||||
siblingIndexParentId,
|
||||
onSiblingIndex,
|
||||
}: {
|
||||
messages?: TMessage[];
|
||||
getMessages?: () => TMessage[] | undefined;
|
||||
|
|
@ -74,6 +76,8 @@ function renderUseResumeOnLoad({
|
|||
conversationId?: string;
|
||||
messagesLoaded?: boolean;
|
||||
onSubmission?: (submission: TSubmission | null) => void;
|
||||
siblingIndexParentId?: string;
|
||||
onSiblingIndex?: (siblingIndex: number) => void;
|
||||
}) {
|
||||
const getMessages = jest.fn(getMessagesOverride ?? (() => messages));
|
||||
const initializeState = (snapshot: MutableSnapshot) => {
|
||||
|
|
@ -86,10 +90,18 @@ function renderUseResumeOnLoad({
|
|||
onSubmission?.(currentSubmission);
|
||||
return null;
|
||||
};
|
||||
const SiblingIndexProbe = () => {
|
||||
const siblingIndex = useRecoilValue(store.messagesSiblingIdxFamily(siblingIndexParentId));
|
||||
if (siblingIndexParentId) {
|
||||
onSiblingIndex?.(siblingIndex);
|
||||
}
|
||||
return null;
|
||||
};
|
||||
|
||||
const wrapper = ({ children }: { children: ReactNode }) => (
|
||||
<RecoilRoot initializeState={initializeState}>
|
||||
<SubmissionProbe />
|
||||
<SiblingIndexProbe />
|
||||
{children}
|
||||
</RecoilRoot>
|
||||
);
|
||||
|
|
@ -191,4 +203,121 @@ describe('useResumeOnLoad', () => {
|
|||
expect(mockUseStreamStatus).toHaveBeenCalledWith(CONVERSATION_ID, true);
|
||||
expect(observedSubmissions[observedSubmissions.length - 1]).toBe(submission);
|
||||
});
|
||||
|
||||
it('restores the branch that owns a pending OAuth resume user message', async () => {
|
||||
const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user');
|
||||
const branchOneResponse = {
|
||||
messageId: 'branch-one-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Branch one response',
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
const branchOneFollowUp = buildUserMessage(CONVERSATION_ID, 'branch-one-follow-up');
|
||||
branchOneFollowUp.parentMessageId = branchOneResponse.messageId;
|
||||
const branchOneTail = {
|
||||
messageId: 'branch-one-tail',
|
||||
parentMessageId: branchOneFollowUp.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Branch one tail',
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
const branchTwoResponse = {
|
||||
messageId: 'branch-two-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Branch two response',
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
const observedSiblingIndexes: number[] = [];
|
||||
|
||||
mockUseStreamStatus.mockReturnValue({
|
||||
isSuccess: true,
|
||||
isFetching: false,
|
||||
data: {
|
||||
active: true,
|
||||
status: 'running',
|
||||
streamId: CONVERSATION_ID,
|
||||
resumeState: {
|
||||
runSteps: [],
|
||||
aggregatedContent: [],
|
||||
replayEvents: [],
|
||||
responseMessageId: 'pending-user_',
|
||||
conversationId: CONVERSATION_ID,
|
||||
userMessage: {
|
||||
messageId: 'pending-user',
|
||||
parentMessageId: branchOneTail.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Use OAuth tool on branch one',
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
renderUseResumeOnLoad({
|
||||
messages: [rootUser, branchOneResponse, branchOneFollowUp, branchOneTail, branchTwoResponse],
|
||||
siblingIndexParentId: rootUser.messageId,
|
||||
onSiblingIndex: (siblingIndex) => observedSiblingIndexes.push(siblingIndex),
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(1);
|
||||
});
|
||||
|
||||
it('restores the assistant sibling selected by a pending regenerate response', async () => {
|
||||
const rootUser = buildUserMessage(CONVERSATION_ID, 'root-user');
|
||||
const olderResponse = {
|
||||
messageId: 'older-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Older response',
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
const newerResponse = {
|
||||
messageId: 'newer-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: 'Newer response',
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
const observedSiblingIndexes: number[] = [];
|
||||
|
||||
mockUseStreamStatus.mockReturnValue({
|
||||
isSuccess: true,
|
||||
isFetching: false,
|
||||
data: {
|
||||
active: true,
|
||||
status: 'running',
|
||||
streamId: CONVERSATION_ID,
|
||||
resumeState: {
|
||||
runSteps: [],
|
||||
aggregatedContent: [],
|
||||
replayEvents: [],
|
||||
responseMessageId: `${olderResponse.messageId}_`,
|
||||
conversationId: CONVERSATION_ID,
|
||||
userMessage: {
|
||||
messageId: rootUser.messageId,
|
||||
parentMessageId: rootUser.parentMessageId,
|
||||
conversationId: CONVERSATION_ID,
|
||||
text: rootUser.text,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
renderUseResumeOnLoad({
|
||||
messages: [rootUser, olderResponse, newerResponse],
|
||||
siblingIndexParentId: rootUser.messageId,
|
||||
onSiblingIndex: (siblingIndex) => observedSiblingIndexes.push(siblingIndex),
|
||||
});
|
||||
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
});
|
||||
|
||||
expect(observedSiblingIndexes[observedSiblingIndexes.length - 1]).toBe(1);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -474,6 +474,369 @@ describe('useStepHandler', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('keeps the pending user message when replayed OAuth tool calls merge immediately', () => {
|
||||
const rootUser = createUserMessage({ messageId: 'root-user' });
|
||||
const selectedResponse = createResponseMessage({
|
||||
messageId: 'selected-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
});
|
||||
const followUpUser = createUserMessage({
|
||||
messageId: 'follow-up-user',
|
||||
parentMessageId: selectedResponse.messageId,
|
||||
});
|
||||
const followUpResponse = createResponseMessage({
|
||||
messageId: 'follow-up-response',
|
||||
parentMessageId: followUpUser.messageId,
|
||||
});
|
||||
const siblingResponse = createResponseMessage({
|
||||
messageId: 'sibling-response',
|
||||
parentMessageId: rootUser.messageId,
|
||||
});
|
||||
const pendingUser = createUserMessage({
|
||||
messageId: 'pending-user',
|
||||
parentMessageId: followUpResponse.messageId,
|
||||
});
|
||||
const initialResponse = createResponseMessage({
|
||||
messageId: 'pending-user_',
|
||||
parentMessageId: pendingUser.messageId,
|
||||
});
|
||||
|
||||
let currentMessages = [
|
||||
rootUser,
|
||||
selectedResponse,
|
||||
followUpUser,
|
||||
followUpResponse,
|
||||
siblingResponse,
|
||||
];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createToolCallRunStep({
|
||||
id: 'step-oauth-login',
|
||||
runId: USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-oauth',
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
const submission = createSubmission({
|
||||
userMessage: pendingUser,
|
||||
messages: currentMessages,
|
||||
initialResponse,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
expect(currentMessages.map((message) => message.messageId)).toEqual([
|
||||
rootUser.messageId,
|
||||
selectedResponse.messageId,
|
||||
followUpUser.messageId,
|
||||
followUpResponse.messageId,
|
||||
siblingResponse.messageId,
|
||||
pendingUser.messageId,
|
||||
initialResponse.messageId,
|
||||
]);
|
||||
expect(currentMessages.at(-1)).toEqual(
|
||||
expect.objectContaining({
|
||||
messageId: initialResponse.messageId,
|
||||
parentMessageId: pendingUser.messageId,
|
||||
content: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: expect.objectContaining({
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
}),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should not insert regenerate transport userMessage before preliminary OAuth steps', () => {
|
||||
const originalUser = createUserMessage({ messageId: 'original-user-message' });
|
||||
const originalResponse = createResponseMessage({
|
||||
messageId: 'original-response-message',
|
||||
parentMessageId: originalUser.messageId,
|
||||
});
|
||||
const initialResponse = createResponseMessage({
|
||||
messageId: 'original-response-message_',
|
||||
parentMessageId: originalUser.messageId,
|
||||
});
|
||||
const regenerateUserMessage = createUserMessage({
|
||||
messageId: 'synthetic-regenerate-user-message',
|
||||
parentMessageId: originalResponse.messageId,
|
||||
});
|
||||
|
||||
let currentMessages = [originalUser];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createToolCallRunStep({
|
||||
id: 'step-oauth-login',
|
||||
runId: USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-oauth',
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
const submission = createSubmission({
|
||||
userMessage: regenerateUserMessage,
|
||||
isRegenerate: true,
|
||||
messages: [originalUser, originalResponse],
|
||||
initialResponse,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
expect(currentMessages).toHaveLength(2);
|
||||
expect(currentMessages.map((message) => message.messageId)).toEqual([
|
||||
originalUser.messageId,
|
||||
initialResponse.messageId,
|
||||
]);
|
||||
expect(currentMessages).not.toContainEqual(
|
||||
expect.objectContaining({ messageId: regenerateUserMessage.messageId }),
|
||||
);
|
||||
expect(currentMessages[1]?.content).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: expect.objectContaining({
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should preserve regenerate history when OAuth steps arrive before optimistic state commits', () => {
|
||||
const originalUser = createUserMessage({ messageId: 'original-user-message' });
|
||||
const originalResponse = createResponseMessage({
|
||||
messageId: 'original-response-message',
|
||||
parentMessageId: originalUser.messageId,
|
||||
});
|
||||
const initialResponse = createResponseMessage({
|
||||
messageId: 'original-response-message_',
|
||||
parentMessageId: originalUser.messageId,
|
||||
});
|
||||
const regenerateUserMessage = createUserMessage({
|
||||
messageId: 'synthetic-regenerate-user-message',
|
||||
parentMessageId: originalResponse.messageId,
|
||||
});
|
||||
|
||||
mockGetMessages.mockReturnValue([]);
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createToolCallRunStep({
|
||||
id: 'step-oauth-login',
|
||||
runId: USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [
|
||||
{
|
||||
id: 'tool-call-oauth',
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
args: '',
|
||||
type: ToolCallTypes.TOOL_CALL,
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
const submission = createSubmission({
|
||||
userMessage: regenerateUserMessage,
|
||||
isRegenerate: true,
|
||||
messages: [originalUser],
|
||||
initialResponse,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
const lastCall = mockSetMessages.mock.calls[mockSetMessages.mock.calls.length - 1][0];
|
||||
expect(lastCall).toHaveLength(2);
|
||||
expect(lastCall.map((message: TMessage) => message.messageId)).toEqual([
|
||||
originalUser.messageId,
|
||||
initialResponse.messageId,
|
||||
]);
|
||||
expect(lastCall).not.toContainEqual(
|
||||
expect.objectContaining({ messageId: regenerateUserMessage.messageId }),
|
||||
);
|
||||
expect(lastCall[1]?.content).toContainEqual(
|
||||
expect.objectContaining({
|
||||
type: ContentTypes.TOOL_CALL,
|
||||
tool_call: expect.objectContaining({
|
||||
name: `oauth${Constants.mcp_delimiter}Google-Workspace`,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('does not seed non-tail regenerate run steps from the latest assistant message', () => {
|
||||
const originalUser = createUserMessage({ messageId: 'original-user-message' });
|
||||
const originalResponse = createResponseMessage({
|
||||
messageId: 'original-response-message',
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: 'Original response',
|
||||
},
|
||||
],
|
||||
});
|
||||
const followUpUser = createUserMessage({
|
||||
messageId: 'follow-up-user-message',
|
||||
parentMessageId: originalResponse.messageId,
|
||||
});
|
||||
const latestResponse = createResponseMessage({
|
||||
messageId: 'latest-response-message',
|
||||
parentMessageId: followUpUser.messageId,
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: 'Latest response content must not leak into regenerate',
|
||||
},
|
||||
],
|
||||
});
|
||||
const initialResponse = createResponseMessage({
|
||||
messageId: 'original-response-message_',
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [],
|
||||
});
|
||||
const regenerateUserMessage = createUserMessage({
|
||||
messageId: 'server-regenerate-user-message',
|
||||
parentMessageId: originalResponse.messageId,
|
||||
});
|
||||
|
||||
let currentMessages = [originalUser, originalResponse, followUpUser, latestResponse];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createRunStep({
|
||||
id: 'step-message-create',
|
||||
runId: 'server-regenerated-response-message',
|
||||
});
|
||||
const submission = createSubmission({
|
||||
userMessage: regenerateUserMessage,
|
||||
isRegenerate: true,
|
||||
messages: [originalUser],
|
||||
initialResponse,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
const regeneratedResponse = currentMessages.find(
|
||||
(message) => message.messageId === 'server-regenerated-response-message',
|
||||
);
|
||||
expect(regeneratedResponse).toEqual(
|
||||
expect.objectContaining({
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [],
|
||||
}),
|
||||
);
|
||||
expect(regeneratedResponse?.content).not.toEqual(latestResponse.content);
|
||||
expect(currentMessages.map((message) => message.messageId)).toEqual([
|
||||
originalUser.messageId,
|
||||
'server-regenerated-response-message',
|
||||
]);
|
||||
});
|
||||
|
||||
it('preserves sibling assistant responses when regenerating one response', () => {
|
||||
const originalUser = createUserMessage({ messageId: 'original-user-message' });
|
||||
const selectedResponse = createResponseMessage({
|
||||
messageId: 'selected-response-message',
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: 'Selected response',
|
||||
},
|
||||
],
|
||||
});
|
||||
const siblingResponse = createResponseMessage({
|
||||
messageId: 'sibling-response-message',
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [
|
||||
{
|
||||
type: ContentTypes.TEXT,
|
||||
text: 'Sibling response should stay visible',
|
||||
},
|
||||
],
|
||||
});
|
||||
const initialResponse = createResponseMessage({
|
||||
messageId: 'selected-response-message_',
|
||||
parentMessageId: originalUser.messageId,
|
||||
content: [],
|
||||
});
|
||||
const regenerateUserMessage = createUserMessage({
|
||||
messageId: 'server-regenerate-user-message',
|
||||
parentMessageId: selectedResponse.messageId,
|
||||
responseMessageId: initialResponse.messageId,
|
||||
});
|
||||
|
||||
let currentMessages = [originalUser, selectedResponse, siblingResponse];
|
||||
mockGetMessages.mockImplementation(() => currentMessages);
|
||||
mockSetMessages.mockImplementation((messages) => {
|
||||
currentMessages = messages;
|
||||
});
|
||||
|
||||
const { result } = renderHook(() => useStepHandler(createHookParams()));
|
||||
|
||||
const runStep = createRunStep({
|
||||
id: 'step-message-create',
|
||||
runId: 'server-regenerated-response-message',
|
||||
});
|
||||
const submission = createSubmission({
|
||||
userMessage: regenerateUserMessage,
|
||||
isRegenerate: true,
|
||||
messages: [originalUser, selectedResponse, siblingResponse],
|
||||
initialResponse,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, submission);
|
||||
});
|
||||
|
||||
expect(currentMessages.map((message) => message.messageId)).toEqual([
|
||||
originalUser.messageId,
|
||||
siblingResponse.messageId,
|
||||
'server-regenerated-response-message',
|
||||
]);
|
||||
expect(currentMessages).not.toContainEqual(
|
||||
expect.objectContaining({ messageId: selectedResponse.messageId }),
|
||||
);
|
||||
});
|
||||
|
||||
it('should propagate step metadata (agentId, groupId) for parallel rendering', () => {
|
||||
const responseMessage = createResponseMessage();
|
||||
mockGetMessages.mockReturnValue([responseMessage]);
|
||||
|
|
|
|||
|
|
@ -68,6 +68,81 @@ type TTitleEvent = {
|
|||
const hasRealTitle = (title?: string | null): title is string =>
|
||||
title != null && title !== '' && title !== 'New Chat';
|
||||
|
||||
export const buildCreatedInitialResponse = ({
|
||||
initialResponse,
|
||||
userMessage,
|
||||
isRegenerate = false,
|
||||
}: Pick<EventSubmission, 'initialResponse' | 'userMessage' | 'isRegenerate'>): TMessage => ({
|
||||
...initialResponse,
|
||||
parentMessageId:
|
||||
isRegenerate && initialResponse.parentMessageId
|
||||
? initialResponse.parentMessageId
|
||||
: userMessage.messageId,
|
||||
messageId:
|
||||
isRegenerate && initialResponse.messageId
|
||||
? initialResponse.messageId
|
||||
: `${userMessage.messageId}_`,
|
||||
conversationId: userMessage.conversationId ?? initialResponse.conversationId,
|
||||
});
|
||||
|
||||
export const isInitialNewConversationSubmission = ({
|
||||
userMessage,
|
||||
}: Pick<EventSubmission, 'userMessage'>): boolean =>
|
||||
userMessage?.parentMessageId === Constants.NO_PARENT;
|
||||
|
||||
export const mergeRegenerateFinalMessages = ({
|
||||
messages,
|
||||
responseMessage,
|
||||
initialResponseId,
|
||||
}: {
|
||||
messages: TMessage[];
|
||||
responseMessage: TMessage;
|
||||
initialResponseId?: string | null;
|
||||
}): TMessage[] => {
|
||||
const finalMessages: TMessage[] = [];
|
||||
let inserted = false;
|
||||
|
||||
for (const message of messages) {
|
||||
if (!message?.messageId || message.messageId === initialResponseId) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (message.messageId === responseMessage.messageId) {
|
||||
finalMessages.push(responseMessage);
|
||||
inserted = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
finalMessages.push(message);
|
||||
}
|
||||
|
||||
if (!inserted) {
|
||||
finalMessages.push(responseMessage);
|
||||
}
|
||||
|
||||
return finalMessages;
|
||||
};
|
||||
|
||||
export const getExistingConversationAbortMessages = ({
|
||||
messages,
|
||||
currentMessages,
|
||||
regenerateMessages,
|
||||
isRegenerate = false,
|
||||
}: Pick<EventSubmission, 'messages' | 'regenerateMessages' | 'isRegenerate'> & {
|
||||
currentMessages?: TMessage[];
|
||||
}): TMessage[] => {
|
||||
if (!isRegenerate) {
|
||||
return [...messages];
|
||||
}
|
||||
|
||||
if (regenerateMessages?.length) {
|
||||
return [...regenerateMessages];
|
||||
}
|
||||
|
||||
const sourceMessages = currentMessages?.length ? currentMessages : messages;
|
||||
return [...sourceMessages];
|
||||
};
|
||||
|
||||
export type EventHandlerParams = {
|
||||
isAddedRequest?: boolean;
|
||||
setCompleted: React.Dispatch<React.SetStateAction<Set<unknown>>>;
|
||||
|
|
@ -409,12 +484,11 @@ export default function useEventHandlers({
|
|||
* drops it, which is the right behavior: by finalize the real
|
||||
* `skill` tool_call is in `content` and takes over rendering.
|
||||
*/
|
||||
const initialResponse = {
|
||||
...submission.initialResponse,
|
||||
parentMessageId: userMessage.messageId,
|
||||
messageId: userMessage.messageId + '_',
|
||||
conversationId: userMessage.conversationId ?? submission.initialResponse.conversationId,
|
||||
};
|
||||
const initialResponse = buildCreatedInitialResponse({
|
||||
initialResponse: submission.initialResponse,
|
||||
userMessage,
|
||||
isRegenerate,
|
||||
});
|
||||
if (isRegenerate) {
|
||||
setMessages([...messages, initialResponse]);
|
||||
} else {
|
||||
|
|
@ -538,15 +612,40 @@ export default function useEventHandlers({
|
|||
const serverConversation = conversation as TConversation;
|
||||
|
||||
try {
|
||||
// Handle early abort - aborted during tool loading before any messages saved
|
||||
// Don't update conversation state, just reset UI and stay on new chat
|
||||
// Handle early abort - aborted before any response message was saved.
|
||||
if ((data as Record<string, unknown>).earlyAbort) {
|
||||
console.log(
|
||||
'[finalHandler] Early abort detected - no messages saved, staying on new chat',
|
||||
);
|
||||
console.log('[finalHandler] Early abort detected - no response message saved');
|
||||
setShowStopButton(false);
|
||||
setIsSubmitting(false);
|
||||
// Navigate to new chat if not already there
|
||||
|
||||
const currentConvoId = submissionConvo.conversationId;
|
||||
const isInitialNewConvo = isInitialNewConversationSubmission(submission);
|
||||
const isExistingConvo =
|
||||
currentConvoId && currentConvoId !== Constants.NEW_CONVO && !isInitialNewConvo;
|
||||
if (isExistingConvo) {
|
||||
const abortMessages = getExistingConversationAbortMessages({
|
||||
messages,
|
||||
isRegenerate,
|
||||
currentMessages: getMessages(),
|
||||
regenerateMessages: submission.regenerateMessages,
|
||||
});
|
||||
setMessages(abortMessages);
|
||||
queryClient.setQueryData<TMessage[]>(
|
||||
[QueryKeys.messages, currentConvoId],
|
||||
abortMessages,
|
||||
);
|
||||
setDraft({ id: currentConvoId, value: requestMessage?.text });
|
||||
return;
|
||||
}
|
||||
|
||||
if (currentConvoId && currentConvoId !== Constants.NEW_CONVO) {
|
||||
removeConvoFromAllQueries(queryClient, currentConvoId);
|
||||
queryClient.removeQueries({ queryKey: [QueryKeys.conversation, currentConvoId] });
|
||||
queryClient.removeQueries({ queryKey: [QueryKeys.messages, currentConvoId] });
|
||||
}
|
||||
setMessages([]);
|
||||
queryClient.setQueryData<TMessage[]>([QueryKeys.messages, Constants.NEW_CONVO], []);
|
||||
setDraft({ id: String(Constants.NEW_CONVO), value: requestMessage?.text });
|
||||
if (location.pathname !== `/c/${Constants.NEW_CONVO}`) {
|
||||
navigate(`/c/${Constants.NEW_CONVO}`, { replace: true });
|
||||
}
|
||||
|
|
@ -622,7 +721,11 @@ export default function useEventHandlers({
|
|||
if (runMessages) {
|
||||
finalMessages = [...runMessages];
|
||||
} else if (isRegenerate && responseMessage) {
|
||||
finalMessages = [...messages, responseMessage];
|
||||
finalMessages = mergeRegenerateFinalMessages({
|
||||
messages: submission.regenerateMessages ?? currentMessages ?? messages,
|
||||
responseMessage,
|
||||
initialResponseId: submission.initialResponse.messageId,
|
||||
});
|
||||
} else if (requestMessage != null && responseMessage != null) {
|
||||
finalMessages = [...messages, requestMessage, responseMessage];
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import {
|
|||
removeNullishValues,
|
||||
} from 'librechat-data-provider';
|
||||
import type {
|
||||
Agents,
|
||||
TMessage,
|
||||
TPayload,
|
||||
TSubmission,
|
||||
|
|
@ -131,9 +132,73 @@ const hasConcreteConversationId = (conversationId?: string | null) =>
|
|||
|
||||
const isInitialNewConversation = (submission: TSubmission) => {
|
||||
const conversationId = submission.conversation?.conversationId;
|
||||
return (
|
||||
submission.userMessage?.parentMessageId === Constants.NO_PARENT &&
|
||||
!hasConcreteConversationId(conversationId)
|
||||
return !hasConcreteConversationId(conversationId);
|
||||
};
|
||||
|
||||
const getToolCallName = (toolCall: unknown) =>
|
||||
toolCall != null && typeof toolCall === 'object' && 'name' in toolCall
|
||||
? (toolCall.name as unknown)
|
||||
: undefined;
|
||||
|
||||
const isOAuthToolCallName = (name: unknown) =>
|
||||
typeof name === 'string' && name.startsWith(`oauth${Constants.mcp_delimiter}`);
|
||||
|
||||
const hasOAuthToolCall = (toolCalls: unknown) =>
|
||||
Array.isArray(toolCalls) &&
|
||||
toolCalls.some((toolCall) => isOAuthToolCallName(getToolCallName(toolCall)));
|
||||
|
||||
const isOAuthStepEvent = (data: unknown) => {
|
||||
if (data == null || typeof data !== 'object' || !('event' in data)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const event = data.event;
|
||||
const payload = 'data' in data ? data.data : undefined;
|
||||
if (payload == null || typeof payload !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (event === StepEvents.ON_RUN_STEP) {
|
||||
const stepDetails = 'stepDetails' in payload ? payload.stepDetails : undefined;
|
||||
return (
|
||||
stepDetails != null &&
|
||||
typeof stepDetails === 'object' &&
|
||||
'tool_calls' in stepDetails &&
|
||||
hasOAuthToolCall(stepDetails.tool_calls)
|
||||
);
|
||||
}
|
||||
|
||||
if (event === StepEvents.ON_RUN_STEP_DELTA) {
|
||||
const delta = 'delta' in payload ? payload.delta : undefined;
|
||||
if (delta == null || typeof delta !== 'object') {
|
||||
return false;
|
||||
}
|
||||
return (
|
||||
('auth' in delta && delta.auth != null) ||
|
||||
('tool_calls' in delta && hasOAuthToolCall(delta.tool_calls))
|
||||
);
|
||||
}
|
||||
|
||||
if (event === StepEvents.ON_RUN_STEP_COMPLETED) {
|
||||
const result = 'result' in payload ? payload.result : undefined;
|
||||
if (result == null || typeof result !== 'object' || !('tool_call' in result)) {
|
||||
return false;
|
||||
}
|
||||
return isOAuthToolCallName(getToolCallName(result.tool_call));
|
||||
}
|
||||
|
||||
return false;
|
||||
};
|
||||
|
||||
const replaceNewConversationUrl = (conversationId: string) => {
|
||||
if (window.location.pathname !== `/c/${Constants.NEW_CONVO}`) {
|
||||
return;
|
||||
}
|
||||
|
||||
window.history.replaceState(
|
||||
window.history.state,
|
||||
'',
|
||||
`/c/${conversationId}${window.location.search}`,
|
||||
);
|
||||
};
|
||||
|
||||
|
|
@ -194,6 +259,95 @@ const hydrateSubmissionMessages = (
|
|||
: submission.initialResponse,
|
||||
});
|
||||
|
||||
const buildResumeEventSubmission = (
|
||||
currentSubmission: TSubmission,
|
||||
currentUserMessage: TMessage,
|
||||
resumeState?: Agents.ResumeState | null,
|
||||
): EventSubmission => {
|
||||
if (!resumeState) {
|
||||
return { ...currentSubmission, userMessage: currentUserMessage } as EventSubmission;
|
||||
}
|
||||
|
||||
const conversationId =
|
||||
resumeState.conversationId ??
|
||||
resumeState.userMessage?.conversationId ??
|
||||
currentUserMessage.conversationId ??
|
||||
currentSubmission.conversation?.conversationId;
|
||||
|
||||
const userMessage = {
|
||||
...currentUserMessage,
|
||||
...resumeState.userMessage,
|
||||
conversationId,
|
||||
isCreatedByUser: true,
|
||||
} as TMessage;
|
||||
|
||||
const responseMessageId =
|
||||
resumeState.responseMessageId ??
|
||||
currentSubmission.initialResponse?.messageId ??
|
||||
`${userMessage.messageId}_`;
|
||||
|
||||
const initialResponse = {
|
||||
...(currentSubmission.initialResponse as TMessage),
|
||||
messageId: responseMessageId,
|
||||
parentMessageId: userMessage.messageId,
|
||||
conversationId,
|
||||
content:
|
||||
resumeState.aggregatedContent ??
|
||||
(currentSubmission.initialResponse as TMessage | undefined)?.content,
|
||||
sender: resumeState.sender ?? currentSubmission.initialResponse?.sender,
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
|
||||
return {
|
||||
...currentSubmission,
|
||||
conversation: {
|
||||
...currentSubmission.conversation,
|
||||
conversationId,
|
||||
},
|
||||
userMessage,
|
||||
initialResponse,
|
||||
} as EventSubmission;
|
||||
};
|
||||
|
||||
const mergeResumeMessages = (
|
||||
messages: TMessage[],
|
||||
userMessage: TMessage,
|
||||
responseMessage: TMessage,
|
||||
): TMessage[] => {
|
||||
const nextMessages = [...messages];
|
||||
const userIndex = nextMessages.findIndex(
|
||||
(message) => message.messageId === userMessage.messageId,
|
||||
);
|
||||
const responseIndex = nextMessages.findIndex(
|
||||
(message) => message.messageId === responseMessage.messageId,
|
||||
);
|
||||
|
||||
if (userIndex >= 0) {
|
||||
nextMessages[userIndex] = { ...nextMessages[userIndex], ...userMessage };
|
||||
}
|
||||
|
||||
if (responseIndex >= 0) {
|
||||
nextMessages[responseIndex] = { ...nextMessages[responseIndex], ...responseMessage };
|
||||
}
|
||||
|
||||
if (userIndex >= 0 && responseIndex >= 0) {
|
||||
return nextMessages;
|
||||
}
|
||||
|
||||
if (userIndex >= 0) {
|
||||
const insertAt = userIndex + 1;
|
||||
nextMessages.splice(insertAt, 0, responseMessage);
|
||||
return nextMessages;
|
||||
}
|
||||
|
||||
if (responseIndex >= 0) {
|
||||
nextMessages.splice(responseIndex, 0, userMessage);
|
||||
return nextMessages;
|
||||
}
|
||||
|
||||
return [...nextMessages, userMessage, responseMessage];
|
||||
};
|
||||
|
||||
/**
|
||||
* Hook for resumable SSE streams.
|
||||
* Separates generation start (POST) from stream subscription (GET EventSource).
|
||||
|
|
@ -324,6 +478,17 @@ export default function useResumableSSE(
|
|||
(currentStreamId: string, currentSubmission: TSubmission, isResume = false) => {
|
||||
let { userMessage } = currentSubmission;
|
||||
let textIndex: number | null = null;
|
||||
const preCreatedStepEvents: Array<Parameters<typeof stepHandler>[0]> = [];
|
||||
const replayPreCreatedStepEvents = () => {
|
||||
if (preCreatedStepEvents.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const submission = { ...currentSubmission, userMessage } as EventSubmission;
|
||||
for (const event of preCreatedStepEvents.splice(0)) {
|
||||
stepHandler(event, submission);
|
||||
}
|
||||
};
|
||||
|
||||
const baseUrl = `${apiBaseUrl()}/api/agents/chat/stream/${encodeURIComponent(currentStreamId)}`;
|
||||
const url = isResume ? `${baseUrl}?resume=true` : baseUrl;
|
||||
|
|
@ -391,6 +556,7 @@ export default function useResumableSSE(
|
|||
overrideParentMessageId: userMessage.overrideParentMessageId,
|
||||
};
|
||||
createdHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
replayPreCreatedStepEvents();
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -408,6 +574,15 @@ export default function useResumableSSE(
|
|||
}
|
||||
|
||||
if (data.event != null) {
|
||||
if (!isResume && !createdStreamIdsRef.current.has(currentStreamId)) {
|
||||
if (isOAuthStepEvent(data)) {
|
||||
preCreatedStepEvents.push(data);
|
||||
stepHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
return;
|
||||
}
|
||||
preCreatedStepEvents.push(data);
|
||||
return;
|
||||
}
|
||||
stepHandler(data, { ...currentSubmission, userMessage } as EventSubmission);
|
||||
return;
|
||||
}
|
||||
|
|
@ -420,13 +595,18 @@ export default function useResumableSSE(
|
|||
|
||||
const runId = v4();
|
||||
setActiveRunId(runId);
|
||||
const resumeSubmission = buildResumeEventSubmission(
|
||||
currentSubmission,
|
||||
userMessage,
|
||||
data.resumeState,
|
||||
);
|
||||
currentSubmission = resumeSubmission;
|
||||
submissionRef.current = resumeSubmission;
|
||||
userMessage = resumeSubmission.userMessage;
|
||||
|
||||
if (data.resumeState?.runSteps) {
|
||||
for (const runStep of data.resumeState.runSteps) {
|
||||
stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, {
|
||||
...currentSubmission,
|
||||
userMessage,
|
||||
} as EventSubmission);
|
||||
stepHandler({ event: StepEvents.ON_RUN_STEP, data: runStep }, resumeSubmission);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -457,20 +637,20 @@ export default function useResumableSSE(
|
|||
});
|
||||
|
||||
if (responseIdx >= 0) {
|
||||
const updated = [...messages];
|
||||
const oldContent = updated[responseIdx]?.content;
|
||||
updated[responseIdx] = {
|
||||
...updated[responseIdx],
|
||||
const oldContent = messages[responseIdx]?.content;
|
||||
const responseMessage = {
|
||||
...messages[responseIdx],
|
||||
content: data.resumeState.aggregatedContent,
|
||||
};
|
||||
} as TMessage;
|
||||
const updated = mergeResumeMessages(messages, userMessage, responseMessage);
|
||||
console.log('[ResumableSSE] SYNC updating message', {
|
||||
messageId: updated[responseIdx]?.messageId,
|
||||
messageId: responseMessage.messageId,
|
||||
oldContentLength: Array.isArray(oldContent) ? oldContent.length : 0,
|
||||
newContentLength: data.resumeState.aggregatedContent?.length,
|
||||
});
|
||||
setMessages(updated);
|
||||
resetContentHandler();
|
||||
syncStepMessage(updated[responseIdx]);
|
||||
syncStepMessage(responseMessage);
|
||||
console.log('[ResumableSSE] SYNC complete, handlers synced');
|
||||
} else {
|
||||
const responseId = serverResponseId ?? `${userMsgId}_`;
|
||||
|
|
@ -482,7 +662,7 @@ export default function useResumableSSE(
|
|||
content: data.resumeState.aggregatedContent,
|
||||
isCreatedByUser: false,
|
||||
} as TMessage;
|
||||
setMessages([...messages, newMessage]);
|
||||
setMessages(mergeResumeMessages(messages, userMessage, newMessage));
|
||||
resetContentHandler();
|
||||
syncStepMessage(newMessage);
|
||||
}
|
||||
|
|
@ -492,16 +672,26 @@ export default function useResumableSSE(
|
|||
titleHandler(data.resumeState.titleEvent);
|
||||
}
|
||||
|
||||
if (data.resumeState?.replayEvents?.length > 0) {
|
||||
console.log(
|
||||
`[ResumableSSE] Replaying ${data.resumeState.replayEvents.length} resume events`,
|
||||
);
|
||||
for (const replayEvent of data.resumeState.replayEvents) {
|
||||
if (replayEvent.event != null) {
|
||||
stepHandler(replayEvent, resumeSubmission);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (data.pendingEvents?.length > 0) {
|
||||
console.log(`[ResumableSSE] Replaying ${data.pendingEvents.length} pending events`);
|
||||
const submission = { ...currentSubmission, userMessage } as EventSubmission;
|
||||
for (const pendingEvent of data.pendingEvents) {
|
||||
if (pendingEvent.event === 'title') {
|
||||
titleHandler(pendingEvent);
|
||||
} else if (pendingEvent.event != null) {
|
||||
stepHandler(pendingEvent, submission);
|
||||
stepHandler(pendingEvent, resumeSubmission);
|
||||
} else if (pendingEvent.type != null) {
|
||||
contentHandler({ data: pendingEvent, submission });
|
||||
contentHandler({ data: pendingEvent, submission: resumeSubmission });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -939,12 +1129,13 @@ export default function useResumableSSE(
|
|||
// Queue title generation if this is a new conversation (first message).
|
||||
// Skip temporary conversations — the server never generates titles for
|
||||
// them, so polling would 404 indefinitely.
|
||||
const isNewConvo = submission.userMessage?.parentMessageId === Constants.NO_PARENT;
|
||||
const isNewConvo = isInitialNewConversation(submission);
|
||||
if (isNewConvo && !submission.isTemporary) {
|
||||
queueTitleGeneration(newStreamId);
|
||||
}
|
||||
if (isInitialNewConversation(submission)) {
|
||||
optimisticStreamIdsRef.current.add(newStreamId);
|
||||
replaceNewConversationUrl(newStreamId);
|
||||
}
|
||||
const streamSubmission = addOptimisticConversation(newStreamId, submission);
|
||||
submissionRef.current = streamSubmission;
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import { useEffect, useRef } from 'react';
|
||||
import { useSetRecoilState, useRecoilValue } from 'recoil';
|
||||
import { useSetRecoilState, useRecoilValue, useRecoilCallback } from 'recoil';
|
||||
import { Constants, tMessageSchema, isAssistantsEndpoint } from 'librechat-data-provider';
|
||||
import type { TMessage, TConversation, TSubmission, Agents } from 'librechat-data-provider';
|
||||
import type { StreamStatusResponse } from '~/data-provider';
|
||||
import { getBranchSiblingIndexesForTarget } from '~/utils';
|
||||
import { useStreamStatus } from '~/data-provider';
|
||||
import store from '~/store';
|
||||
|
||||
|
|
@ -42,6 +43,41 @@ function resumeStateMatchesSubmission(
|
|||
return !!responseMessageId && resumeState.responseMessageId === responseMessageId;
|
||||
}
|
||||
|
||||
function getResumeBranchTargetMessageId(
|
||||
resumeState: Agents.ResumeState,
|
||||
messages: TMessage[],
|
||||
): string | null | undefined {
|
||||
const responseMessageId = resumeState.responseMessageId;
|
||||
if (!responseMessageId) {
|
||||
return resumeState.userMessage?.parentMessageId;
|
||||
}
|
||||
|
||||
const unpaddedResponseMessageId = responseMessageId.replace(/_+$/, '');
|
||||
let hasResponseMessage = false;
|
||||
let hasUnpaddedResponseMessage = false;
|
||||
|
||||
for (const message of messages) {
|
||||
if (message.messageId === responseMessageId) {
|
||||
hasResponseMessage = true;
|
||||
break;
|
||||
}
|
||||
|
||||
if (message.messageId === unpaddedResponseMessageId) {
|
||||
hasUnpaddedResponseMessage = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (hasResponseMessage) {
|
||||
return responseMessageId;
|
||||
}
|
||||
|
||||
if (hasUnpaddedResponseMessage) {
|
||||
return unpaddedResponseMessageId;
|
||||
}
|
||||
|
||||
return resumeState.userMessage?.parentMessageId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build a submission object from resume state for reconnected streams.
|
||||
* This provides the minimum data needed for useResumableSSE to subscribe.
|
||||
|
|
@ -148,6 +184,22 @@ export default function useResumeOnLoad(
|
|||
const resumableEnabled = !isAssistantsEndpoint(actualEndpoint);
|
||||
// Track conversations we've already processed (either resumed or skipped)
|
||||
const processedConvoRef = useRef<string | null>(null);
|
||||
const restoreResumeBranch = useRecoilCallback(
|
||||
({ set }) =>
|
||||
(resumeState: Agents.ResumeState, messages: TMessage[], activeConversationId: string) => {
|
||||
const targetMessageId = getResumeBranchTargetMessageId(resumeState, messages);
|
||||
const branchIndexes = getBranchSiblingIndexesForTarget(
|
||||
messages,
|
||||
targetMessageId,
|
||||
activeConversationId,
|
||||
);
|
||||
|
||||
for (const { parentMessageId, siblingIdx } of branchIndexes) {
|
||||
set(store.messagesSiblingIdxFamily(parentMessageId), siblingIdx);
|
||||
}
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
// Check for active stream when conversation changes
|
||||
const submissionConvoId = currentSubmission?.conversation?.conversationId;
|
||||
|
|
@ -266,6 +318,7 @@ export default function useResumeOnLoad(
|
|||
|
||||
// Build submission from resume state if available
|
||||
if (streamStatus.resumeState) {
|
||||
restoreResumeBranch(streamStatus.resumeState, messages, conversationId);
|
||||
const submission = buildSubmissionFromResumeState(
|
||||
streamStatus.resumeState,
|
||||
streamStatus.streamId,
|
||||
|
|
@ -308,6 +361,7 @@ export default function useResumeOnLoad(
|
|||
streamStatus,
|
||||
getMessages,
|
||||
setSubmission,
|
||||
restoreResumeBranch,
|
||||
]);
|
||||
|
||||
// Reset processedConvoRef when conversation changes to allow re-checking
|
||||
|
|
|
|||
|
|
@ -444,9 +444,91 @@ export default function useStepHandler({
|
|||
|
||||
const stepHandler = useCallback(
|
||||
(stepEvent: TStepEvent, submission: EventSubmission) => {
|
||||
const messages = getMessages() || [];
|
||||
const submissionMessages = submission.messages ?? [];
|
||||
const getEventMessages = (candidateMessages: TMessage[]) =>
|
||||
submission.isRegenerate ? candidateMessages : getCurrentMessages(candidateMessages);
|
||||
const messages = getEventMessages(submissionMessages);
|
||||
const { userMessage } = submission;
|
||||
let parentMessageId = userMessage.messageId;
|
||||
const getRegenerateResponseIds = (responseMessageId: string) => {
|
||||
const ids = new Set<string>();
|
||||
const addId = (id?: string | null) => {
|
||||
if (!id) {
|
||||
return;
|
||||
}
|
||||
ids.add(id);
|
||||
ids.add(id.replace(/_+$/, ''));
|
||||
};
|
||||
addId(responseMessageId);
|
||||
addId(submission.initialResponse?.messageId);
|
||||
addId(submission.userMessage?.responseMessageId);
|
||||
return ids;
|
||||
};
|
||||
const shouldRemoveRegenerateResponse = (message: TMessage, responseMessageId: string) =>
|
||||
submission.isRegenerate &&
|
||||
!message.isCreatedByUser &&
|
||||
getRegenerateResponseIds(responseMessageId).has(message.messageId);
|
||||
const ensureUserMessagePresent = (
|
||||
candidateMessages: TMessage[],
|
||||
responseMessageId: string,
|
||||
) => {
|
||||
if (
|
||||
submission.isRegenerate ||
|
||||
!userMessage?.messageId ||
|
||||
candidateMessages.some((message) => message.messageId === userMessage.messageId)
|
||||
) {
|
||||
return candidateMessages;
|
||||
}
|
||||
|
||||
const responseIndex = candidateMessages.findIndex(
|
||||
(message) => message.messageId === responseMessageId,
|
||||
);
|
||||
if (responseIndex < 0) {
|
||||
return [...candidateMessages, userMessage as TMessage];
|
||||
}
|
||||
|
||||
const nextMessages = [...candidateMessages];
|
||||
nextMessages.splice(responseIndex, 0, userMessage as TMessage);
|
||||
return nextMessages;
|
||||
};
|
||||
const getResponseBaseMessages = (
|
||||
candidateMessages: TMessage[],
|
||||
responseMessageId: string,
|
||||
ensureUserMessage = false,
|
||||
) => {
|
||||
const currentMessages = getEventMessages(candidateMessages);
|
||||
if (!submission.isRegenerate) {
|
||||
return ensureUserMessage
|
||||
? ensureUserMessagePresent(currentMessages, responseMessageId)
|
||||
: currentMessages;
|
||||
}
|
||||
return currentMessages.filter(
|
||||
(message) => !shouldRemoveRegenerateResponse(message, responseMessageId),
|
||||
);
|
||||
};
|
||||
const mergeResponseMessage = (
|
||||
candidateMessages: TMessage[],
|
||||
updatedResponse: TMessage,
|
||||
responseMessageId: string,
|
||||
options?: { ensureUserMessage?: boolean },
|
||||
) => {
|
||||
const currentMessages = getResponseBaseMessages(
|
||||
candidateMessages,
|
||||
responseMessageId,
|
||||
options?.ensureUserMessage === true,
|
||||
);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
);
|
||||
return hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
};
|
||||
let parentMessageId =
|
||||
submission.isRegenerate && submission.initialResponse?.parentMessageId
|
||||
? submission.initialResponse.parentMessageId
|
||||
: userMessage.messageId;
|
||||
|
||||
const currentTime = Date.now();
|
||||
if (currentTime - lastAnnouncementTimeRef.current > MESSAGE_UPDATE_INTERVAL) {
|
||||
|
|
@ -480,10 +562,13 @@ export default function useStepHandler({
|
|||
let response = messageMap.current.get(responseMessageId);
|
||||
|
||||
if (!response) {
|
||||
// Find the actual response message - check if last message is a response, otherwise use initialResponse
|
||||
// Find the actual response message. Regenerate submissions can target
|
||||
// an earlier branch while the visible history still ends at a later
|
||||
// assistant message, so never seed a regenerated response from the
|
||||
// conversation tail.
|
||||
const lastMessage = messages[messages.length - 1] as TMessage;
|
||||
const responseMessage =
|
||||
lastMessage && !lastMessage.isCreatedByUser
|
||||
!submission.isRegenerate && lastMessage && !lastMessage.isCreatedByUser
|
||||
? lastMessage
|
||||
: (submission?.initialResponse as TMessage);
|
||||
|
||||
|
|
@ -505,14 +590,20 @@ export default function useStepHandler({
|
|||
|
||||
// Get fresh messages to handle multi-tab scenarios where messages may have loaded
|
||||
// after this handler started (Tab 2 may have more complete history now)
|
||||
const freshMessages = getMessages() || [];
|
||||
const currentMessages = freshMessages.length > messages.length ? freshMessages : messages;
|
||||
const currentMessages = getResponseBaseMessages(messages, responseMessageId);
|
||||
|
||||
// Remove any existing response placeholder
|
||||
let updatedMessages = currentMessages.filter((m) => m.messageId !== responseMessageId);
|
||||
let updatedMessages = currentMessages.filter(
|
||||
(message) => !shouldRemoveRegenerateResponse(message, responseMessageId),
|
||||
);
|
||||
|
||||
// Ensure userMessage is present (multi-tab: Tab 2 may not have it yet)
|
||||
if (!updatedMessages.some((m) => m.messageId === userMessage.messageId)) {
|
||||
// Ensure userMessage is present (multi-tab: Tab 2 may not have it yet).
|
||||
// Regenerate reuses an existing user turn; its submission userMessage is only
|
||||
// a transport placeholder and must not become a new visible branch.
|
||||
if (
|
||||
!submission.isRegenerate &&
|
||||
!updatedMessages.some((m) => m.messageId === userMessage.messageId)
|
||||
) {
|
||||
updatedMessages = [...updatedMessages, userMessage as TMessage];
|
||||
}
|
||||
|
||||
|
|
@ -548,17 +639,11 @@ export default function useStepHandler({
|
|||
});
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
setMessages(
|
||||
mergeResponseMessage(messages, updatedResponse, responseMessageId, {
|
||||
ensureUserMessage: true,
|
||||
}),
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
|
||||
if (runStep.summary != null) {
|
||||
|
|
@ -580,8 +665,8 @@ export default function useStepHandler({
|
|||
);
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getMessages() || [];
|
||||
setMessages([...currentMessages.slice(0, -1), updatedResponse]);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
setMessages(mergeResponseMessage(currentMessages, updatedResponse, responseMessageId));
|
||||
}
|
||||
|
||||
const bufferedDeltas = pendingDeltaBuffer.current.get(runStep.id);
|
||||
|
|
@ -619,8 +704,8 @@ export default function useStepHandler({
|
|||
agentUpdateMeta,
|
||||
);
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getMessages() || [];
|
||||
setMessages([...currentMessages.slice(0, -1), updatedResponse]);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
setMessages(mergeResponseMessage(currentMessages, updatedResponse, responseMessageId));
|
||||
}
|
||||
} else if (stepEvent.event === StepEvents.ON_MESSAGE_DELTA) {
|
||||
const messageDelta = stepEvent.data;
|
||||
|
|
@ -662,8 +747,8 @@ export default function useStepHandler({
|
|||
getStepMetadata(runStep),
|
||||
);
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getMessages() || [];
|
||||
setMessages([...currentMessages.slice(0, -1), updatedResponse]);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
setMessages(mergeResponseMessage(currentMessages, updatedResponse, responseMessageId));
|
||||
}
|
||||
} else if (stepEvent.event === StepEvents.ON_REASONING_DELTA) {
|
||||
const reasoningDelta = stepEvent.data;
|
||||
|
|
@ -705,8 +790,8 @@ export default function useStepHandler({
|
|||
getStepMetadata(runStep),
|
||||
);
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getMessages() || [];
|
||||
setMessages([...currentMessages.slice(0, -1), updatedResponse]);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
setMessages(mergeResponseMessage(currentMessages, updatedResponse, responseMessageId));
|
||||
}
|
||||
} else if (stepEvent.event === StepEvents.ON_RUN_STEP_DELTA) {
|
||||
const runStepDelta = stepEvent.data;
|
||||
|
|
@ -761,17 +846,11 @@ export default function useStepHandler({
|
|||
});
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
setMessages(
|
||||
mergeResponseMessage(messages, updatedResponse, responseMessageId, {
|
||||
ensureUserMessage: true,
|
||||
}),
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
} else if (stepEvent.event === StepEvents.ON_RUN_STEP_COMPLETED) {
|
||||
const { result } = stepEvent.data;
|
||||
|
|
@ -810,17 +889,11 @@ export default function useStepHandler({
|
|||
);
|
||||
|
||||
messageMap.current.set(responseMessageId, updatedResponse);
|
||||
const currentMessages = getCurrentMessages(messages);
|
||||
const hasResponseMessage = currentMessages.some(
|
||||
(msg) => msg.messageId === responseMessageId,
|
||||
setMessages(
|
||||
mergeResponseMessage(messages, updatedResponse, responseMessageId, {
|
||||
ensureUserMessage: true,
|
||||
}),
|
||||
);
|
||||
const updatedMessages = hasResponseMessage
|
||||
? currentMessages.map((msg) =>
|
||||
msg.messageId === responseMessageId ? updatedResponse : msg,
|
||||
)
|
||||
: [...currentMessages, updatedResponse];
|
||||
|
||||
setMessages(updatedMessages);
|
||||
}
|
||||
} else if (stepEvent.event === StepEvents.ON_SUBAGENT_UPDATE) {
|
||||
applySubagentUpdate(stepEvent.data);
|
||||
|
|
@ -863,8 +936,8 @@ export default function useStepHandler({
|
|||
summarizeDeltaRaf.current = null;
|
||||
const latest = messageMap.current.get(responseMessageId);
|
||||
if (latest) {
|
||||
const msgs = getMessages() || [];
|
||||
setMessages([...msgs.slice(0, -1), latest]);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
setMessages(mergeResponseMessage(currentMessages, latest, responseMessageId));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
@ -882,8 +955,7 @@ export default function useStepHandler({
|
|||
return;
|
||||
}
|
||||
|
||||
const currentMessages = getMessages() || [];
|
||||
const targetIndex = currentMessages.findIndex((m) => m.messageId === completeMessageId);
|
||||
const currentMessages = submission.isRegenerate ? messages : getMessages() || [];
|
||||
|
||||
if (completeData.error) {
|
||||
const filtered = targetMessage.content.filter(
|
||||
|
|
@ -894,11 +966,7 @@ export default function useStepHandler({
|
|||
announcePolite({ message: 'summarize_failed', isStatus: true });
|
||||
const cleaned = { ...targetMessage, content: filtered };
|
||||
messageMap.current.set(completeMessageId, cleaned);
|
||||
if (targetIndex >= 0) {
|
||||
const updated = [...currentMessages];
|
||||
updated[targetIndex] = cleaned;
|
||||
setMessages(updated);
|
||||
}
|
||||
setMessages(mergeResponseMessage(currentMessages, cleaned, completeMessageId));
|
||||
}
|
||||
} else {
|
||||
let didFinalize = false;
|
||||
|
|
@ -912,13 +980,11 @@ export default function useStepHandler({
|
|||
}
|
||||
return part;
|
||||
});
|
||||
if (didFinalize && targetIndex >= 0) {
|
||||
if (didFinalize) {
|
||||
announcePolite({ message: 'summarize_completed', isStatus: true });
|
||||
const finalized = { ...targetMessage, content: updatedContent };
|
||||
messageMap.current.set(completeMessageId, finalized);
|
||||
const updated = [...currentMessages];
|
||||
updated[targetIndex] = finalized;
|
||||
setMessages(updated);
|
||||
setMessages(mergeResponseMessage(currentMessages, finalized, completeMessageId));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -21,6 +21,11 @@ export const STREAM_START_FAILED_METADATA_KEY = 'streamStartFailed';
|
|||
|
||||
type SiblingIndexLookup = (parentMessageId: string | null | undefined) => number;
|
||||
|
||||
export type BranchSiblingIndex = {
|
||||
parentMessageId: string | null | undefined;
|
||||
siblingIdx: number;
|
||||
};
|
||||
|
||||
export const selectActiveBranchTail = (
|
||||
messages: TMessage[] | null | undefined,
|
||||
rootSiblingKey: string | null | undefined,
|
||||
|
|
@ -83,6 +88,55 @@ export const getMessageBranchSiblingParentIds = (
|
|||
return Array.from(parentIds);
|
||||
};
|
||||
|
||||
export const getBranchSiblingIndexesForTarget = (
|
||||
messages: TMessage[] | null | undefined,
|
||||
targetMessageId: string | null | undefined,
|
||||
rootSiblingKey: string | null | undefined,
|
||||
): BranchSiblingIndex[] => {
|
||||
if (!targetMessageId) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const messagesTree = buildTree({ messages: messages ?? null });
|
||||
if (!messagesTree?.length) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const branchIndexes: BranchSiblingIndex[] = [];
|
||||
const findTargetPath = (
|
||||
siblings: TMessage[] | undefined,
|
||||
parentMessageId: string | null | undefined,
|
||||
): boolean => {
|
||||
if (!siblings?.length) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (let index = 0; index < siblings.length; index++) {
|
||||
const message = siblings[index];
|
||||
if (!message) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const isTarget = message.messageId === targetMessageId;
|
||||
const childHasTarget = findTargetPath(message.children, message.messageId);
|
||||
if (isTarget || childHasTarget) {
|
||||
if (siblings.length > 1) {
|
||||
branchIndexes.unshift({
|
||||
parentMessageId,
|
||||
siblingIdx: siblings.length - index - 1,
|
||||
});
|
||||
}
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
};
|
||||
|
||||
findTargetPath(messagesTree, rootSiblingKey);
|
||||
return branchIndexes;
|
||||
};
|
||||
|
||||
export const getLatestText = (message?: TMessage | null, includeIndex?: boolean): string => {
|
||||
if (!message) {
|
||||
return '';
|
||||
|
|
|
|||
|
|
@ -36,6 +36,44 @@ test.describe('core chat loop', () => {
|
|||
await expect(page.getByTestId('convo-item').first()).toBeVisible();
|
||||
});
|
||||
|
||||
test('can switch back to the previous branch after regenerating an earlier response', async ({
|
||||
page,
|
||||
}) => {
|
||||
test.setTimeout(90000);
|
||||
const firstMessage = 'branch root from e2e';
|
||||
const followUpMessage = 'follow-up on original branch from e2e';
|
||||
|
||||
await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
|
||||
await selectMockEndpoint(page, MOCK_ENDPOINTS[0]);
|
||||
|
||||
let response = await sendMessage(page, firstMessage);
|
||||
expect(response.ok()).toBeTruthy();
|
||||
await expect(mockReply(page)).toBeVisible();
|
||||
|
||||
response = await sendMessage(page, followUpMessage);
|
||||
expect(response.ok()).toBeTruthy();
|
||||
await expect(page.getByText(followUpMessage)).toBeVisible();
|
||||
|
||||
const firstAssistantMessage = page
|
||||
.getByTestId('messages-view')
|
||||
.locator('.message-render')
|
||||
.nth(1);
|
||||
await firstAssistantMessage.hover();
|
||||
const regenerateButton = firstAssistantMessage.locator('button[title="Regenerate"]').last();
|
||||
await expect(regenerateButton).toBeVisible();
|
||||
|
||||
const [regenerateResponse] = await Promise.all([
|
||||
page.waitForResponse(isAgentsStream, { timeout: 30000 }),
|
||||
regenerateButton.click(),
|
||||
]);
|
||||
expect(regenerateResponse.ok()).toBeTruthy();
|
||||
|
||||
await expect(page.getByText('2 / 2')).toBeVisible();
|
||||
await page.getByRole('button', { name: 'Previous sibling message' }).click();
|
||||
await expect(page.getByText('1 / 2')).toBeVisible();
|
||||
await expect(page.getByText(followUpMessage)).toBeVisible();
|
||||
});
|
||||
|
||||
test('keeps upload-to-provider CSV attached to the sent message and model input', async ({
|
||||
page,
|
||||
}) => {
|
||||
|
|
|
|||
|
|
@ -143,6 +143,105 @@ async function revisitConversationAndExpectMessages(
|
|||
await expectVisibleMessages(page, texts);
|
||||
}
|
||||
|
||||
async function mockActiveOAuthResumeStream({
|
||||
page,
|
||||
authUrl,
|
||||
conversationId,
|
||||
parentMessageId,
|
||||
pendingPrompt,
|
||||
pendingUserMessageId,
|
||||
}: {
|
||||
page: Page;
|
||||
authUrl: string;
|
||||
conversationId: string;
|
||||
parentMessageId: string;
|
||||
pendingPrompt: string;
|
||||
pendingUserMessageId: string;
|
||||
}) {
|
||||
const pendingResponseMessageId = `${pendingUserMessageId}_`;
|
||||
const toolCallId = `${pendingUserMessageId}:Google-Workspace`;
|
||||
const stepId = 'step_oauth_login_Google-Workspace';
|
||||
const resumeState = {
|
||||
runSteps: [],
|
||||
aggregatedContent: [],
|
||||
responseMessageId: pendingResponseMessageId,
|
||||
conversationId,
|
||||
userMessage: {
|
||||
messageId: pendingUserMessageId,
|
||||
parentMessageId,
|
||||
conversationId,
|
||||
text: pendingPrompt,
|
||||
},
|
||||
replayEvents: [
|
||||
{
|
||||
event: 'on_run_step',
|
||||
data: {
|
||||
runId: 'USE_PRELIM_RESPONSE_MESSAGE_ID',
|
||||
id: stepId,
|
||||
type: 'tool_calls',
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [
|
||||
{
|
||||
id: toolCallId,
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
type: 'tool_call_chunk',
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [
|
||||
{
|
||||
id: toolCallId,
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
type: 'tool_call_chunk',
|
||||
args: '',
|
||||
},
|
||||
],
|
||||
auth: authUrl,
|
||||
expires_at: Date.now() + 120000,
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
await page.route(`**/api/agents/chat/status/${conversationId}`, (route) =>
|
||||
route.fulfill({
|
||||
status: 200,
|
||||
contentType: 'application/json',
|
||||
body: JSON.stringify({
|
||||
active: true,
|
||||
streamId: conversationId,
|
||||
status: 'running',
|
||||
aggregatedContent: [],
|
||||
createdAt: Date.now(),
|
||||
resumeState,
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
await page.route(`**/api/agents/chat/stream/${conversationId}**`, (route) =>
|
||||
route.fulfill({
|
||||
status: 200,
|
||||
contentType: 'text/event-stream',
|
||||
body: `event: message\ndata: ${JSON.stringify({
|
||||
sync: true,
|
||||
resumeState,
|
||||
pendingEvents: [],
|
||||
})}\n\n`,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
async function openMockChat(page: Page) {
|
||||
await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
|
||||
await selectMockEndpoint(page, MOCK_ENDPOINTS[0]);
|
||||
|
|
@ -432,6 +531,53 @@ test.describe('message tree stream operations', () => {
|
|||
expect(messages.some((message) => messageText(message).includes(followReply))).toBe(false);
|
||||
});
|
||||
|
||||
test('resumes pending OAuth on the selected older branch after reload', async ({ page }) => {
|
||||
const label = uniqueLabel('oauth-branch');
|
||||
const rootPrompt = countedPrompt(`${label}-root`);
|
||||
const firstReply = countedReplyText(`${label}-root`, 1);
|
||||
const regeneratedReply = countedReplyText(`${label}-root`, 2);
|
||||
const followPrompt = replyPrompt(`${label}-follow`);
|
||||
const followReply = replyText(`${label}-follow`);
|
||||
const pendingPrompt = replyPrompt(`${label}-oauth`);
|
||||
|
||||
await openMockChat(page);
|
||||
await sendAndExpectReply(page, rootPrompt, firstReply);
|
||||
const conversationId = await conversationIdFromPage(page);
|
||||
await sendAndExpectReply(page, followPrompt, followReply);
|
||||
|
||||
await waitForGenerationStart(page, () =>
|
||||
clickMessageTitleButton(page, firstReply, 'Regenerate'),
|
||||
);
|
||||
await expect(messagesView(page).getByText(regeneratedReply)).toBeVisible({ timeout: 30000 });
|
||||
|
||||
await clickSibling(page, regeneratedReply, 'Previous');
|
||||
await expectVisibleMessages(page, [firstReply, followPrompt, followReply]);
|
||||
await expect(messagesView(page).getByText(regeneratedReply)).toBeHidden();
|
||||
|
||||
const messages = await waitForMessages(
|
||||
page,
|
||||
conversationId,
|
||||
(items) =>
|
||||
items.some((message) => messageText(message).includes(followReply)) &&
|
||||
items.some((message) => messageText(message).includes(regeneratedReply)),
|
||||
'two-branch conversation',
|
||||
);
|
||||
const branchOneTail = findMessage(messages, followReply, false);
|
||||
|
||||
await mockActiveOAuthResumeStream({
|
||||
page,
|
||||
conversationId,
|
||||
parentMessageId: branchOneTail.messageId,
|
||||
pendingPrompt,
|
||||
pendingUserMessageId: `${label}-pending-user`,
|
||||
authUrl: `https://auth.example.test/${label}`,
|
||||
});
|
||||
|
||||
await page.reload({ timeout: 10000 });
|
||||
await expectVisibleMessages(page, [firstReply, followPrompt, followReply, pendingPrompt]);
|
||||
await expect(messagesView(page).getByText(regeneratedReply)).toBeHidden();
|
||||
});
|
||||
|
||||
test('long threads retain regenerated and save-and-submit branches after revisit', async ({
|
||||
page,
|
||||
}) => {
|
||||
|
|
|
|||
|
|
@ -13,8 +13,8 @@ import {
|
|||
ReauthenticationRequiredError,
|
||||
resolveOboToken,
|
||||
} from '~/mcp/oauth';
|
||||
import { PENDING_STALE_MS, normalizeExpiresAt } from '~/flow/manager';
|
||||
import { sanitizeUrlForLogging, isClientRejectionMessage, isOAuthServer } from './utils';
|
||||
import { PENDING_STALE_MS, normalizeExpiresAt } from '~/flow/manager';
|
||||
import { withTimeout } from '~/utils/promise';
|
||||
import { MCPConnection } from './connection';
|
||||
import { processMCPEnv } from '~/utils';
|
||||
|
|
@ -46,7 +46,7 @@ export class MCPConnectionFactory {
|
|||
protected readonly flowManager?: FlowStateManager<MCPOAuthTokens | null>;
|
||||
protected readonly tokenMethods?: TokenMethods;
|
||||
protected readonly signal?: AbortSignal;
|
||||
protected readonly oauthStart?: (authURL: string) => Promise<void>;
|
||||
protected readonly oauthStart?: t.OAuthStartHandler;
|
||||
protected readonly oauthEnd?: () => Promise<void>;
|
||||
protected readonly returnOnOAuth?: boolean;
|
||||
protected readonly connectionTimeout?: number;
|
||||
|
|
@ -484,6 +484,15 @@ export class MCPConnectionFactory {
|
|||
};
|
||||
}
|
||||
|
||||
private getOAuthReplayExpiresAt(createdAt?: number): number | undefined {
|
||||
if (!createdAt) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const expiresAt = createdAt + PENDING_STALE_MS;
|
||||
return expiresAt > Date.now() ? expiresAt : undefined;
|
||||
}
|
||||
|
||||
/** Sets up OAuth event handlers for the connection */
|
||||
protected handleOAuthEvents(connection: MCPConnection): () => void {
|
||||
const oauthHandler = async (data: { serverUrl?: string }) => {
|
||||
|
|
@ -504,7 +513,24 @@ export class MCPConnectionFactory {
|
|||
logger.debug(
|
||||
`${this.logPrefix} Recent PENDING OAuth flow exists (${Math.round(pendingAge / 1000)}s old), skipping new initiation`,
|
||||
);
|
||||
connection.emit('oauthFailed', new Error('OAuth flow initiated - return early'));
|
||||
const flowMeta = existingFlow.metadata as MCPOAuthFlowMetadata | undefined;
|
||||
const storedAuthUrl = flowMeta?.authorizationUrl;
|
||||
if (storedAuthUrl && typeof this.oauthStart === 'function') {
|
||||
const expiresAt = this.getOAuthReplayExpiresAt(existingFlow.createdAt);
|
||||
if (!expiresAt) {
|
||||
logger.debug(`${this.logPrefix} PENDING OAuth flow expired before replay`);
|
||||
connection.emit(
|
||||
'oauthFailed',
|
||||
new Error('Pending OAuth flow expired before replay'),
|
||||
);
|
||||
return;
|
||||
}
|
||||
logger.info(
|
||||
`${this.logPrefix} Re-issuing stored authorization URL while reusing PENDING flow`,
|
||||
);
|
||||
await this.oauthStart(storedAuthUrl, { expiresAt });
|
||||
}
|
||||
connection.emit('oauthFailed', new Error('Pending OAuth flow reused - return early'));
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -771,10 +797,14 @@ export class MCPConnectionFactory {
|
|||
|
||||
const storedAuthUrl = flowMeta?.authorizationUrl;
|
||||
if (storedAuthUrl && typeof this.oauthStart === 'function') {
|
||||
const expiresAt = this.getOAuthReplayExpiresAt(existingFlow.createdAt);
|
||||
if (!expiresAt) {
|
||||
throw new Error('Pending OAuth flow expired before replay');
|
||||
}
|
||||
logger.info(
|
||||
`${this.logPrefix} Re-issuing stored authorization URL to caller while joining PENDING flow`,
|
||||
);
|
||||
await this.oauthStart(storedAuthUrl);
|
||||
await this.oauthStart(storedAuthUrl, { expiresAt });
|
||||
}
|
||||
|
||||
reusedStoredClient = flowMeta?.reusedStoredClient === true;
|
||||
|
|
|
|||
|
|
@ -10,18 +10,18 @@ import type { FlowStateManager } from '~/flow/manager';
|
|||
import type { MCPOAuthTokens } from './oauth';
|
||||
import type { RequestBody } from '~/types';
|
||||
import type * as t from './types';
|
||||
import { isUserSourced, requiresOAuthMachinery, requiresUserScopedConnection } from './utils';
|
||||
import { MCPServersInitializer } from './registry/MCPServersInitializer';
|
||||
import { OboTokenResolutionError, resolveOboToken } from '~/mcp/oauth';
|
||||
import { MCPServerInspector } from './registry/MCPServerInspector';
|
||||
import { MCPServersRegistry } from './registry/MCPServersRegistry';
|
||||
import { UserConnectionManager } from './UserConnectionManager';
|
||||
import { ConnectionsRepository } from './ConnectionsRepository';
|
||||
import { MCPConnectionFactory } from './MCPConnectionFactory';
|
||||
import { OboTokenResolutionError, resolveOboToken } from '~/mcp/oauth';
|
||||
import { preProcessGraphTokens } from '~/utils/graph';
|
||||
import { formatToolContent } from './parsers';
|
||||
import { MCPConnection } from './connection';
|
||||
import { processMCPEnv } from '~/utils/env';
|
||||
import { isUserSourced, requiresOAuthMachinery, requiresUserScopedConnection } from './utils';
|
||||
|
||||
function createOboToolCallErrorMessage(
|
||||
logPrefix: string,
|
||||
|
|
@ -327,7 +327,7 @@ Please follow these instructions when using tools from the respective MCP server
|
|||
tokenMethods?: TokenMethods;
|
||||
customUserVars?: Record<string, string>;
|
||||
flowManager: FlowStateManager<MCPOAuthTokens | null>;
|
||||
oauthStart?: (authURL: string) => Promise<void>;
|
||||
oauthStart?: t.OAuthStartHandler;
|
||||
oauthEnd?: () => Promise<void>;
|
||||
graphTokenResolver?: GraphTokenResolver;
|
||||
oboTokenResolver?: OboTokenResolver;
|
||||
|
|
|
|||
|
|
@ -1,13 +1,34 @@
|
|||
import { logger } from '@librechat/data-schemas';
|
||||
import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js';
|
||||
import type { MCPOAuthFlowMetadata } from '~/mcp/oauth';
|
||||
import type { FlowState } from '~/flow/types';
|
||||
import type * as t from './types';
|
||||
import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry';
|
||||
import { ConnectionsRepository } from '~/mcp/ConnectionsRepository';
|
||||
import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory';
|
||||
import { isUserSourced, requiresOAuthMachinery } from './utils';
|
||||
import { PENDING_STALE_MS } from '~/flow/manager';
|
||||
import { MCPOAuthHandler } from '~/mcp/oauth';
|
||||
import { MCPConnection } from './connection';
|
||||
import { mcpConfig } from './mcpConfig';
|
||||
|
||||
type PendingOAuthStart = {
|
||||
authURL: string;
|
||||
options?: t.OAuthStartOptions;
|
||||
};
|
||||
|
||||
type PendingOAuthState = {
|
||||
oauthStarts: Set<t.OAuthStartHandler>;
|
||||
emittedAuthUrls: WeakMap<t.OAuthStartHandler, string>;
|
||||
primaryOAuthStart?: t.OAuthStartHandler;
|
||||
lastOAuthStart?: PendingOAuthStart;
|
||||
};
|
||||
|
||||
type PendingConnection = {
|
||||
promise: Promise<MCPConnection>;
|
||||
oauth: PendingOAuthState;
|
||||
};
|
||||
|
||||
/**
|
||||
* Abstract base class for managing user-specific MCP connections with lifecycle management.
|
||||
* Only meant to be extended by MCPManager.
|
||||
|
|
@ -23,7 +44,7 @@ export abstract class UserConnectionManager {
|
|||
/** Last activity timestamp for users (not per server) */
|
||||
protected userLastActivity: Map<string, number> = new Map();
|
||||
/** In-flight connection promises keyed by `userId:serverName` — coalesces concurrent attempts */
|
||||
protected pendingConnections: Map<string, Promise<MCPConnection>> = new Map();
|
||||
protected pendingConnections: Map<string, PendingConnection> = new Map();
|
||||
|
||||
/** Updates the last activity timestamp for a user */
|
||||
protected updateUserLastActivity(userId: string): void {
|
||||
|
|
@ -48,25 +69,197 @@ export abstract class UserConnectionManager {
|
|||
const pending = this.pendingConnections.get(lockKey);
|
||||
if (pending) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Joining in-flight connection attempt`);
|
||||
return pending;
|
||||
await this.addPendingOAuthStart(pending.oauth, opts, userId);
|
||||
return pending.promise;
|
||||
}
|
||||
}
|
||||
|
||||
const connectionPromise = this.createUserConnectionInternal(opts, userId);
|
||||
const pendingOAuth = this.createPendingOAuthState(opts.oauthStart);
|
||||
const connectionPromise = this.createUserConnectionInternal(
|
||||
{
|
||||
...opts,
|
||||
oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth),
|
||||
},
|
||||
userId,
|
||||
);
|
||||
|
||||
if (!forceNew) {
|
||||
this.pendingConnections.set(lockKey, connectionPromise);
|
||||
this.pendingConnections.set(lockKey, { promise: connectionPromise, oauth: pendingOAuth });
|
||||
}
|
||||
|
||||
try {
|
||||
return await connectionPromise;
|
||||
} finally {
|
||||
if (!forceNew && this.pendingConnections.get(lockKey) === connectionPromise) {
|
||||
if (!forceNew && this.pendingConnections.get(lockKey)?.promise === connectionPromise) {
|
||||
this.pendingConnections.delete(lockKey);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private createPendingOAuthState(oauthStart?: t.OAuthStartHandler): PendingOAuthState {
|
||||
return {
|
||||
oauthStarts: oauthStart ? new Set([oauthStart]) : new Set(),
|
||||
emittedAuthUrls: new WeakMap<t.OAuthStartHandler, string>(),
|
||||
primaryOAuthStart: oauthStart,
|
||||
};
|
||||
}
|
||||
|
||||
private createPendingOAuthStart(
|
||||
serverName: string,
|
||||
userId: string,
|
||||
pendingOAuth: PendingOAuthState,
|
||||
): t.OAuthStartHandler {
|
||||
return async (authURL, options) => {
|
||||
pendingOAuth.lastOAuthStart = { authURL, options };
|
||||
|
||||
let primaryError: unknown;
|
||||
const oauthStarts = Array.from(pendingOAuth.oauthStarts);
|
||||
for (const oauthStart of oauthStarts) {
|
||||
try {
|
||||
await this.emitPendingOAuthStart(pendingOAuth, oauthStart, authURL, options);
|
||||
} catch (error) {
|
||||
if (oauthStart === pendingOAuth.primaryOAuthStart) {
|
||||
primaryError = error;
|
||||
} else {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to notify joined OAuth listener`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (primaryError) {
|
||||
throw primaryError;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private async addPendingOAuthStart(
|
||||
pendingOAuth: PendingOAuthState,
|
||||
opts: t.UserMCPConnectionOptions,
|
||||
userId: string,
|
||||
): Promise<void> {
|
||||
const { oauthStart, serverName } = opts;
|
||||
if (typeof oauthStart !== 'function') {
|
||||
return;
|
||||
}
|
||||
|
||||
pendingOAuth.oauthStarts.add(oauthStart);
|
||||
const lastOAuthStart = pendingOAuth.lastOAuthStart;
|
||||
if (lastOAuthStart) {
|
||||
try {
|
||||
const pendingOAuthStart =
|
||||
lastOAuthStart.options?.expiresAt == null
|
||||
? await this.getFlowPendingOAuthStart(opts, userId)
|
||||
: undefined;
|
||||
const replayOAuthStart =
|
||||
pendingOAuthStart?.authURL === lastOAuthStart.authURL
|
||||
? pendingOAuthStart
|
||||
: lastOAuthStart;
|
||||
await this.emitPendingOAuthStart(
|
||||
pendingOAuth,
|
||||
oauthStart,
|
||||
replayOAuthStart.authURL,
|
||||
replayOAuthStart.options,
|
||||
);
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
await this.reissuePendingOAuthStart(opts, userId, pendingOAuth);
|
||||
}
|
||||
|
||||
private async emitPendingOAuthStart(
|
||||
pendingOAuth: PendingOAuthState,
|
||||
oauthStart: t.OAuthStartHandler,
|
||||
authURL: string,
|
||||
options?: t.OAuthStartOptions,
|
||||
): Promise<void> {
|
||||
if (pendingOAuth.emittedAuthUrls.get(oauthStart) === authURL) {
|
||||
return;
|
||||
}
|
||||
pendingOAuth.emittedAuthUrls.set(oauthStart, authURL);
|
||||
await oauthStart(authURL, options);
|
||||
}
|
||||
|
||||
private getPendingOAuthStart(flow: FlowState | null | undefined): PendingOAuthStart | undefined {
|
||||
if (flow?.status !== 'PENDING') {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const expiresAt = flow.createdAt + PENDING_STALE_MS;
|
||||
if (expiresAt <= Date.now()) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const metadata = flow.metadata as MCPOAuthFlowMetadata | undefined;
|
||||
const authorizationUrl = metadata?.authorizationUrl;
|
||||
if (!authorizationUrl) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return { authURL: authorizationUrl, options: { expiresAt } };
|
||||
}
|
||||
|
||||
private async getFlowPendingOAuthStart(
|
||||
{ flowManager, serverName }: Pick<t.UserMCPConnectionOptions, 'flowManager' | 'serverName'>,
|
||||
userId: string,
|
||||
): Promise<PendingOAuthStart | undefined> {
|
||||
if (!flowManager) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const flowId = MCPOAuthHandler.generateFlowId(userId, serverName);
|
||||
const existingFlow = await flowManager.getFlowState(flowId, 'mcp_oauth');
|
||||
return this.getPendingOAuthStart(existingFlow);
|
||||
}
|
||||
|
||||
private async reissuePendingOAuthStart(
|
||||
{ flowManager, oauthStart, serverName }: t.UserMCPConnectionOptions,
|
||||
userId: string,
|
||||
pendingOAuth?: PendingOAuthState,
|
||||
): Promise<void> {
|
||||
if (!flowManager || typeof oauthStart !== 'function') {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const pendingOAuthStart = await this.getFlowPendingOAuthStart(
|
||||
{ flowManager, serverName },
|
||||
userId,
|
||||
);
|
||||
if (!pendingOAuthStart) {
|
||||
return;
|
||||
}
|
||||
|
||||
logger.info(
|
||||
`[MCP][User: ${userId}][${serverName}] Re-issuing stored authorization URL while joining in-flight connection`,
|
||||
);
|
||||
if (pendingOAuth) {
|
||||
pendingOAuth.lastOAuthStart = pendingOAuthStart;
|
||||
await this.emitPendingOAuthStart(
|
||||
pendingOAuth,
|
||||
oauthStart,
|
||||
pendingOAuthStart.authURL,
|
||||
pendingOAuthStart.options,
|
||||
);
|
||||
} else {
|
||||
await oauthStart(pendingOAuthStart.authURL, pendingOAuthStart.options);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private async createUserConnectionInternal(
|
||||
{
|
||||
serverName,
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@ import type { FlowStateManager } from '~/flow/manager';
|
|||
import type { MCPOAuthTokens } from '~/mcp/oauth';
|
||||
import type * as t from '~/mcp/types';
|
||||
import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory';
|
||||
import { MCPConnection } from '~/mcp/connection';
|
||||
import { MCPOAuthHandler, MCPTokenStorage } from '~/mcp/oauth';
|
||||
import { PENDING_STALE_MS } from '~/flow/manager';
|
||||
import { MCPConnection } from '~/mcp/connection';
|
||||
import { processMCPEnv } from '~/utils';
|
||||
|
||||
jest.mock('~/mcp/connection');
|
||||
|
|
@ -398,7 +399,7 @@ describe('MCPConnectionFactory', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('should skip new OAuth flow initiation when a PENDING flow already exists (returnOnOAuth)', async () => {
|
||||
it('should re-issue the stored OAuth URL when a PENDING flow already exists (returnOnOAuth)', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
|
|
@ -416,7 +417,10 @@ describe('MCPConnectionFactory', () => {
|
|||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
type: 'mcp_oauth',
|
||||
metadata: { codeVerifier: 'existing-verifier' },
|
||||
metadata: {
|
||||
codeVerifier: 'existing-verifier',
|
||||
authorizationUrl: 'https://auth.example.com/existing',
|
||||
},
|
||||
createdAt: Date.now(),
|
||||
});
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
|
@ -438,10 +442,73 @@ describe('MCPConnectionFactory', () => {
|
|||
await oauthRequiredHandler!({ serverUrl: 'https://api.example.com' });
|
||||
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled();
|
||||
expect(oauthOptions.oauthStart).toHaveBeenCalledWith(
|
||||
'https://auth.example.com/existing',
|
||||
expect.objectContaining({ expiresAt: expect.any(Number) }),
|
||||
);
|
||||
expect(mockFlowManager.deleteFlow).not.toHaveBeenCalled();
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith(
|
||||
'oauthFailed',
|
||||
expect.objectContaining({ message: 'OAuth flow initiated - return early' }),
|
||||
expect.objectContaining({ message: 'Pending OAuth flow reused - return early' }),
|
||||
);
|
||||
});
|
||||
|
||||
it('should emit oauthFailed when a PENDING flow expires before replay (returnOnOAuth)', async () => {
|
||||
const createdAt = 1000;
|
||||
const firstNow = createdAt + PENDING_STALE_MS - 1;
|
||||
const expiredNow = createdAt + PENDING_STALE_MS;
|
||||
jest
|
||||
.spyOn(Date, 'now')
|
||||
.mockReturnValueOnce(firstNow)
|
||||
.mockReturnValueOnce(expiredNow)
|
||||
.mockReturnValue(expiredNow);
|
||||
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
user: mockUser,
|
||||
};
|
||||
|
||||
const oauthOptions: t.OAuthConnectionOptions = {
|
||||
user: mockUser,
|
||||
useOAuth: true,
|
||||
returnOnOAuth: true,
|
||||
oauthStart: jest.fn(),
|
||||
flowManager: mockFlowManager,
|
||||
};
|
||||
|
||||
mockFlowManager.getFlowState.mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
type: 'mcp_oauth',
|
||||
metadata: {
|
||||
codeVerifier: 'existing-verifier',
|
||||
authorizationUrl: 'https://auth.example.com/existing',
|
||||
},
|
||||
createdAt,
|
||||
});
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
||||
let oauthRequiredHandler: (data: Record<string, unknown>) => Promise<void>;
|
||||
mockConnectionInstance.on.mockImplementation((event, handler) => {
|
||||
if (event === 'oauthRequired') {
|
||||
oauthRequiredHandler = handler as (data: Record<string, unknown>) => Promise<void>;
|
||||
}
|
||||
return mockConnectionInstance;
|
||||
});
|
||||
|
||||
try {
|
||||
await MCPConnectionFactory.create(basicOptions, oauthOptions);
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
await oauthRequiredHandler!({ serverUrl: 'https://api.example.com' });
|
||||
|
||||
expect(oauthOptions.oauthStart).not.toHaveBeenCalled();
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled();
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith(
|
||||
'oauthFailed',
|
||||
expect.objectContaining({ message: 'Pending OAuth flow expired before replay' }),
|
||||
);
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@ const mockLogger = logger as jest.Mocked<typeof logger>;
|
|||
describe('MCP OAuth Race Condition Fixes', () => {
|
||||
afterEach(() => {
|
||||
jest.clearAllMocks();
|
||||
jest.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe('Fix 1: Connection mutex coalesces concurrent attempts', () => {
|
||||
|
|
@ -77,7 +78,7 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
dbId: undefined,
|
||||
};
|
||||
|
||||
jest
|
||||
const registrySpy = jest
|
||||
.spyOn(
|
||||
// eslint-disable-next-line @typescript-eslint/no-require-imports
|
||||
require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry,
|
||||
|
|
@ -97,27 +98,213 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
return mockConnection as never;
|
||||
});
|
||||
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, { ttl: 30000, ci: true });
|
||||
const user = { id: 'user-1' };
|
||||
const opts = {
|
||||
serverName: 'test-server',
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
try {
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, {
|
||||
ttl: 30000,
|
||||
ci: true,
|
||||
});
|
||||
const user = { id: 'user-1' };
|
||||
const opts = {
|
||||
serverName: 'test-server',
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
};
|
||||
|
||||
const [conn1, conn2, conn3] = await Promise.all([
|
||||
manager.getUserConnection(opts),
|
||||
manager.getUserConnection(opts),
|
||||
manager.getUserConnection(opts),
|
||||
]);
|
||||
|
||||
expect(conn1).toBe(conn2);
|
||||
expect(conn2).toBe(conn3);
|
||||
expect(createSpy).toHaveBeenCalledTimes(1);
|
||||
expect(manager.createCallCount).toBe(1);
|
||||
} finally {
|
||||
createSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('should re-issue the pending OAuth URL when joining an in-flight connection', async () => {
|
||||
const { UserConnectionManager } = await import('~/mcp/UserConnectionManager');
|
||||
|
||||
class TestManager extends UserConnectionManager {}
|
||||
|
||||
const manager = new TestManager();
|
||||
|
||||
const mockConnection = {
|
||||
isConnected: jest.fn().mockResolvedValue(true),
|
||||
disconnect: jest.fn().mockResolvedValue(undefined),
|
||||
isStale: jest.fn().mockReturnValue(false),
|
||||
};
|
||||
|
||||
const [conn1, conn2, conn3] = await Promise.all([
|
||||
manager.getUserConnection(opts),
|
||||
manager.getUserConnection(opts),
|
||||
manager.getUserConnection(opts),
|
||||
]);
|
||||
const mockAppConnections = { has: jest.fn().mockResolvedValue(false) };
|
||||
manager.appConnections = mockAppConnections as never;
|
||||
|
||||
expect(conn1).toBe(conn2);
|
||||
expect(conn2).toBe(conn3);
|
||||
expect(createSpy).toHaveBeenCalledTimes(1);
|
||||
expect(manager.createCallCount).toBe(1);
|
||||
const mockConfig = {
|
||||
type: 'streamable-http',
|
||||
url: 'http://localhost:9999/',
|
||||
updatedAt: undefined,
|
||||
dbId: undefined,
|
||||
};
|
||||
|
||||
createSpy.mockRestore();
|
||||
const registrySpy = jest
|
||||
.spyOn(
|
||||
// eslint-disable-next-line @typescript-eslint/no-require-imports
|
||||
require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry,
|
||||
'getInstance',
|
||||
)
|
||||
.mockReturnValue({
|
||||
getServerConfig: jest.fn().mockResolvedValue(mockConfig),
|
||||
shouldEnableSSRFProtection: jest.fn().mockReturnValue(false),
|
||||
getAllowedDomains: jest.fn().mockReturnValue(null),
|
||||
getAllowedAddresses: jest.fn().mockReturnValue(null),
|
||||
});
|
||||
|
||||
const { MCPConnectionFactory } = await import('~/mcp/MCPConnectionFactory');
|
||||
const createSpy = jest.spyOn(MCPConnectionFactory, 'create').mockImplementation(async () => {
|
||||
await new Promise((r) => setTimeout(r, 100));
|
||||
return mockConnection as never;
|
||||
});
|
||||
|
||||
try {
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, {
|
||||
ttl: 30000,
|
||||
ci: true,
|
||||
});
|
||||
const user = { id: 'user-oauth' };
|
||||
const serverName = 'test-server';
|
||||
const authorizationUrl = 'https://auth.example.com/existing';
|
||||
await flowManager.initFlow(`${user.id}:${serverName}`, 'mcp_oauth', { authorizationUrl });
|
||||
|
||||
const firstConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
});
|
||||
const oauthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const joinedConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
oauthStart,
|
||||
});
|
||||
|
||||
const [conn1, conn2] = await Promise.all([firstConnection, joinedConnection]);
|
||||
|
||||
expect(conn1).toBe(conn2);
|
||||
expect(oauthStart).toHaveBeenCalledWith(
|
||||
authorizationUrl,
|
||||
expect.objectContaining({ expiresAt: expect.any(Number) }),
|
||||
);
|
||||
expect(createSpy).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
createSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('should preserve pending OAuth expiry when joining after the prompt was emitted', async () => {
|
||||
const { UserConnectionManager } = await import('~/mcp/UserConnectionManager');
|
||||
|
||||
class TestManager extends UserConnectionManager {}
|
||||
|
||||
const manager = new TestManager();
|
||||
|
||||
const mockConnection = {
|
||||
isConnected: jest.fn().mockResolvedValue(true),
|
||||
disconnect: jest.fn().mockResolvedValue(undefined),
|
||||
isStale: jest.fn().mockReturnValue(false),
|
||||
};
|
||||
|
||||
const mockAppConnections = { has: jest.fn().mockResolvedValue(false) };
|
||||
manager.appConnections = mockAppConnections as never;
|
||||
|
||||
const mockConfig = {
|
||||
type: 'streamable-http',
|
||||
url: 'http://localhost:9999/',
|
||||
requiresOAuth: true,
|
||||
updatedAt: undefined,
|
||||
dbId: undefined,
|
||||
};
|
||||
|
||||
const registrySpy = jest
|
||||
.spyOn(
|
||||
// eslint-disable-next-line @typescript-eslint/no-require-imports
|
||||
require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry,
|
||||
'getInstance',
|
||||
)
|
||||
.mockReturnValue({
|
||||
getServerConfig: jest.fn().mockResolvedValue(mockConfig),
|
||||
shouldEnableSSRFProtection: jest.fn().mockReturnValue(false),
|
||||
getAllowedDomains: jest.fn().mockReturnValue(null),
|
||||
getAllowedAddresses: jest.fn().mockReturnValue(null),
|
||||
});
|
||||
|
||||
let releaseConnection: () => void = () => undefined;
|
||||
const connectionReleased = new Promise<void>((resolve) => {
|
||||
releaseConnection = resolve;
|
||||
});
|
||||
|
||||
const { MCPConnectionFactory } = await import('~/mcp/MCPConnectionFactory');
|
||||
const authorizationUrl = 'https://auth.example.com/current';
|
||||
const createSpy = jest
|
||||
.spyOn(MCPConnectionFactory, 'create')
|
||||
.mockImplementation(async (_basicOptions, oauthOptions) => {
|
||||
if (oauthOptions && 'oauthStart' in oauthOptions) {
|
||||
await oauthOptions.oauthStart?.(authorizationUrl);
|
||||
}
|
||||
await connectionReleased;
|
||||
return mockConnection as never;
|
||||
});
|
||||
|
||||
try {
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, {
|
||||
ttl: 30000,
|
||||
ci: true,
|
||||
});
|
||||
const user = { id: 'user-oauth-emitted' };
|
||||
const serverName = 'test-server';
|
||||
await flowManager.initFlow(`${user.id}:${serverName}`, 'mcp_oauth', { authorizationUrl });
|
||||
|
||||
const firstOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const firstConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
oauthStart: firstOAuthStart,
|
||||
});
|
||||
for (let i = 0; i < 20 && firstOAuthStart.mock.calls.length === 0; i++) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
}
|
||||
expect(firstOAuthStart).toHaveBeenCalledWith(authorizationUrl, undefined);
|
||||
|
||||
const joinedOAuthStart = jest.fn().mockResolvedValue(undefined);
|
||||
const joinedConnection = manager.getUserConnection({
|
||||
serverName,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
oauthStart: joinedOAuthStart,
|
||||
});
|
||||
|
||||
releaseConnection();
|
||||
const [conn1, conn2] = await Promise.all([firstConnection, joinedConnection]);
|
||||
|
||||
expect(conn1).toBe(conn2);
|
||||
expect(joinedOAuthStart).toHaveBeenCalledWith(
|
||||
authorizationUrl,
|
||||
expect.objectContaining({ expiresAt: expect.any(Number) }),
|
||||
);
|
||||
expect(createSpy).toHaveBeenCalledTimes(1);
|
||||
} finally {
|
||||
releaseConnection();
|
||||
createSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
|
||||
it('should not coalesce when forceNew is true', async () => {
|
||||
|
|
@ -144,7 +331,7 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
dbId: undefined,
|
||||
};
|
||||
|
||||
jest
|
||||
const registrySpy = jest
|
||||
.spyOn(
|
||||
// eslint-disable-next-line @typescript-eslint/no-require-imports
|
||||
require('~/mcp/registry/MCPServersRegistry').MCPServersRegistry,
|
||||
|
|
@ -158,33 +345,41 @@ describe('MCP OAuth Race Condition Fixes', () => {
|
|||
});
|
||||
|
||||
const { MCPConnectionFactory } = await import('~/mcp/MCPConnectionFactory');
|
||||
jest.spyOn(MCPConnectionFactory, 'create').mockImplementation(async () => {
|
||||
const createSpy = jest.spyOn(MCPConnectionFactory, 'create').mockImplementation(async () => {
|
||||
callCount++;
|
||||
await new Promise((r) => setTimeout(r, 50));
|
||||
return makeConnection() as never;
|
||||
});
|
||||
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, { ttl: 30000, ci: true });
|
||||
const user = { id: 'user-2' };
|
||||
try {
|
||||
const store = new MockKeyv();
|
||||
const flowManager = new FlowStateManager(store as unknown as Keyv, {
|
||||
ttl: 30000,
|
||||
ci: true,
|
||||
});
|
||||
const user = { id: 'user-2' };
|
||||
|
||||
const [conn1, conn2] = await Promise.all([
|
||||
manager.getUserConnection({
|
||||
serverName: 'test-server',
|
||||
forceNew: true,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
}),
|
||||
manager.getUserConnection({
|
||||
serverName: 'test-server',
|
||||
forceNew: true,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
}),
|
||||
]);
|
||||
const [conn1, conn2] = await Promise.all([
|
||||
manager.getUserConnection({
|
||||
serverName: 'test-server',
|
||||
forceNew: true,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
}),
|
||||
manager.getUserConnection({
|
||||
serverName: 'test-server',
|
||||
forceNew: true,
|
||||
user: user as never,
|
||||
flowManager: flowManager as never,
|
||||
}),
|
||||
]);
|
||||
|
||||
expect(callCount).toBe(2);
|
||||
expect(conn1).not.toBe(conn2);
|
||||
expect(callCount).toBe(2);
|
||||
expect(conn1).not.toBe(conn2);
|
||||
} finally {
|
||||
createSpy.mockRestore();
|
||||
registrySpy.mockRestore();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
158
packages/api/src/mcp/oauth/events.ts
Normal file
158
packages/api/src/mcp/oauth/events.ts
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
import { Constants, Time } from 'librechat-data-provider';
|
||||
import { GraphEvents, StepTypes } from '@librechat/agents';
|
||||
import type * as t from '~/types';
|
||||
import { buildOAuthToolCallName } from '~/mcp/utils';
|
||||
|
||||
export type OAuthPromptOptions = {
|
||||
expiresAt?: number;
|
||||
};
|
||||
|
||||
export type OAuthToolCall = {
|
||||
id?: string;
|
||||
name: string;
|
||||
type: 'tool_call' | 'tool_call_chunk';
|
||||
args?: string;
|
||||
output?: string;
|
||||
};
|
||||
|
||||
export function getOAuthPromptExpiresAt(options?: OAuthPromptOptions, now = Date.now()): number {
|
||||
return typeof options?.expiresAt === 'number' && Number.isFinite(options.expiresAt)
|
||||
? options.expiresAt
|
||||
: now + Time.TWO_MINUTES;
|
||||
}
|
||||
|
||||
export function getMCPServerNamesFromTools(tools?: unknown[] | null): Set<string> {
|
||||
const serverNames = new Set<string>();
|
||||
|
||||
for (const tool of tools ?? []) {
|
||||
if (typeof tool !== 'string') {
|
||||
continue;
|
||||
}
|
||||
|
||||
const delimiterIndex = tool.indexOf(Constants.mcp_delimiter);
|
||||
if (delimiterIndex === -1) {
|
||||
continue;
|
||||
}
|
||||
|
||||
serverNames.add(tool.slice(delimiterIndex + Constants.mcp_delimiter.length));
|
||||
}
|
||||
|
||||
return serverNames;
|
||||
}
|
||||
|
||||
export function buildMCPAuthStepId(serverName: string): string {
|
||||
return `step_oauth_login_${serverName}`;
|
||||
}
|
||||
|
||||
export function buildMCPAuthToolCall({
|
||||
id,
|
||||
args,
|
||||
output,
|
||||
serverName,
|
||||
type = 'tool_call_chunk',
|
||||
}: {
|
||||
id?: string;
|
||||
args?: string;
|
||||
output?: string;
|
||||
serverName: string;
|
||||
type?: OAuthToolCall['type'];
|
||||
}): OAuthToolCall {
|
||||
return {
|
||||
...(id != null ? { id } : {}),
|
||||
name: buildOAuthToolCallName(serverName),
|
||||
...(args != null ? { args } : {}),
|
||||
...(output != null ? { output } : {}),
|
||||
type,
|
||||
};
|
||||
}
|
||||
|
||||
export function buildMCPAuthRunStepEvent({
|
||||
runId = Constants.USE_PRELIM_RESPONSE_MESSAGE_ID,
|
||||
stepId,
|
||||
toolCall,
|
||||
index = 0,
|
||||
}: {
|
||||
runId?: string | null;
|
||||
stepId: string;
|
||||
toolCall: OAuthToolCall;
|
||||
index?: number;
|
||||
}): t.ServerSentEvent {
|
||||
return {
|
||||
event: GraphEvents.ON_RUN_STEP,
|
||||
data: {
|
||||
runId,
|
||||
id: stepId,
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
index,
|
||||
stepDetails: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [toolCall],
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export function buildMCPAuthRunStepDeltaEvent({
|
||||
authURL,
|
||||
stepId,
|
||||
toolCall,
|
||||
options,
|
||||
}: {
|
||||
authURL: string;
|
||||
stepId: string;
|
||||
toolCall: OAuthToolCall;
|
||||
options?: OAuthPromptOptions;
|
||||
}): t.ServerSentEvent {
|
||||
return {
|
||||
event: GraphEvents.ON_RUN_STEP_DELTA,
|
||||
data: {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [{ ...toolCall, args: '' }],
|
||||
auth: authURL,
|
||||
expires_at: getOAuthPromptExpiresAt(options),
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export function buildMCPAuthRunStepEndDeltaEvent({
|
||||
stepId,
|
||||
toolCall,
|
||||
}: {
|
||||
stepId: string;
|
||||
toolCall: OAuthToolCall;
|
||||
}): t.ServerSentEvent {
|
||||
return {
|
||||
event: GraphEvents.ON_RUN_STEP_DELTA,
|
||||
data: {
|
||||
id: stepId,
|
||||
delta: {
|
||||
type: StepTypes.TOOL_CALLS,
|
||||
tool_calls: [{ ...toolCall }],
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export function buildMCPAuthRunStepCompletedEvent({
|
||||
stepId,
|
||||
toolCall,
|
||||
index = 0,
|
||||
}: {
|
||||
stepId: string;
|
||||
toolCall: OAuthToolCall;
|
||||
index?: number;
|
||||
}): t.ServerSentEvent {
|
||||
return {
|
||||
event: GraphEvents.ON_RUN_STEP_COMPLETED,
|
||||
data: {
|
||||
result: {
|
||||
id: stepId,
|
||||
index,
|
||||
tool_call: toolCall,
|
||||
},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -4,3 +4,5 @@ export * from './tokens';
|
|||
export * from './detectOAuth';
|
||||
export * from './methods';
|
||||
export * from './obo';
|
||||
export * from './pending';
|
||||
export * from './events';
|
||||
|
|
|
|||
70
packages/api/src/mcp/oauth/pending.ts
Normal file
70
packages/api/src/mcp/oauth/pending.ts
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
import { logger } from '@librechat/data-schemas';
|
||||
|
||||
import type { MCPOAuthFlowMetadata, MCPOAuthTokens } from './types';
|
||||
import type { FlowStateManager } from '~/flow/manager';
|
||||
import type { FlowState } from '~/flow/types';
|
||||
import type * as t from '~/mcp/types';
|
||||
|
||||
import { PENDING_STALE_MS } from '~/flow/manager';
|
||||
import { MCPOAuthHandler } from './handler';
|
||||
|
||||
const MCP_OAUTH_FLOW_TYPE = 'mcp_oauth';
|
||||
|
||||
type PendingOAuthFlowState = Pick<
|
||||
FlowState<MCPOAuthTokens | null>,
|
||||
'createdAt' | 'metadata' | 'status'
|
||||
>;
|
||||
|
||||
type PendingOAuthFlowManager = Pick<FlowStateManager<MCPOAuthTokens | null>, 'getFlowState'>;
|
||||
|
||||
export type PendingOAuthStart = {
|
||||
authURL: string;
|
||||
options?: t.OAuthStartOptions;
|
||||
};
|
||||
|
||||
export type ReplayablePendingMCPOAuthStartOptions = {
|
||||
flowManager?: PendingOAuthFlowManager | null;
|
||||
userId: string;
|
||||
serverName: string;
|
||||
};
|
||||
|
||||
export function getReplayablePendingMCPOAuthStartFromFlow(
|
||||
flow: PendingOAuthFlowState | null | undefined,
|
||||
now = Date.now(),
|
||||
): PendingOAuthStart | undefined {
|
||||
if (flow?.status !== 'PENDING') {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const expiresAt = flow.createdAt + PENDING_STALE_MS;
|
||||
if (expiresAt <= now) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const metadata = flow.metadata as Partial<MCPOAuthFlowMetadata>;
|
||||
const authorizationUrl = metadata.authorizationUrl;
|
||||
if (!authorizationUrl) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return { authURL: authorizationUrl, options: { expiresAt } };
|
||||
}
|
||||
|
||||
export async function getReplayablePendingMCPOAuthStart({
|
||||
flowManager,
|
||||
userId,
|
||||
serverName,
|
||||
}: ReplayablePendingMCPOAuthStartOptions): Promise<PendingOAuthStart | undefined> {
|
||||
if (!flowManager || typeof flowManager.getFlowState !== 'function') {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
try {
|
||||
const flowId = MCPOAuthHandler.generateFlowId(userId, serverName);
|
||||
const flowState = await flowManager.getFlowState(flowId, MCP_OAUTH_FLOW_TYPE);
|
||||
return getReplayablePendingMCPOAuthStartFromFlow(flowState);
|
||||
} catch (error) {
|
||||
logger.warn(`[MCP OAuth] Failed to inspect pending flow for ${serverName}:`, error);
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
|
@ -19,10 +19,10 @@ import type {
|
|||
import type { SearchResultData, UIResource, TPlugin } from 'librechat-data-provider';
|
||||
import type { TokenMethods, IUser } from '@librechat/data-schemas';
|
||||
import type { LCTool } from '@librechat/agents';
|
||||
import type { OboTokenResolver, OboTrustChecker } from '~/mcp/oauth/obo';
|
||||
import type { FlowStateManager } from '~/flow/manager';
|
||||
import type { RequestBody } from '~/types/http';
|
||||
import type * as o from '~/mcp/oauth/types';
|
||||
import type { OboTokenResolver, OboTrustChecker } from '~/mcp/oauth/obo';
|
||||
|
||||
export type StdioOptions = z.infer<typeof StdioOptionsSchema>;
|
||||
export type WebSocketOptions = z.infer<typeof WebSocketOptionsSchema>;
|
||||
|
|
@ -201,12 +201,18 @@ export interface UserConnectionContext {
|
|||
connectionTimeout?: number;
|
||||
}
|
||||
|
||||
export interface OAuthStartOptions {
|
||||
expiresAt?: number;
|
||||
}
|
||||
|
||||
export type OAuthStartHandler = (authURL: string, options?: OAuthStartOptions) => Promise<void>;
|
||||
|
||||
export interface OAuthConnectionOptions extends UserConnectionContext {
|
||||
useOAuth: true;
|
||||
flowManager: FlowStateManager<o.MCPOAuthTokens | null>;
|
||||
tokenMethods?: TokenMethods;
|
||||
signal?: AbortSignal;
|
||||
oauthStart?: (authURL: string) => Promise<void>;
|
||||
oauthStart?: OAuthStartHandler;
|
||||
oauthEnd?: () => Promise<void>;
|
||||
returnOnOAuth?: boolean;
|
||||
oboTokenResolver?: OboTokenResolver;
|
||||
|
|
@ -221,7 +227,7 @@ export interface UserMCPConnectionOptions extends UserConnectionContext {
|
|||
flowManager?: FlowStateManager<o.MCPOAuthTokens | null>;
|
||||
tokenMethods?: TokenMethods;
|
||||
signal?: AbortSignal;
|
||||
oauthStart?: (authURL: string) => Promise<void>;
|
||||
oauthStart?: OAuthStartHandler;
|
||||
oauthEnd?: () => Promise<void>;
|
||||
returnOnOAuth?: boolean;
|
||||
oboTokenResolver?: OboTokenResolver;
|
||||
|
|
@ -234,7 +240,7 @@ export interface ToolDiscoveryOptions {
|
|||
flowManager?: FlowStateManager<o.MCPOAuthTokens | null>;
|
||||
tokenMethods?: TokenMethods;
|
||||
signal?: AbortSignal;
|
||||
oauthStart?: (authURL: string) => Promise<void>;
|
||||
oauthStart?: OAuthStartHandler;
|
||||
customUserVars?: Record<string, string>;
|
||||
requestBody?: RequestBody;
|
||||
connectionTimeout?: number;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import { Constants, parseTextParts } from 'librechat-data-provider';
|
||||
import { logger, getTenantId, SYSTEM_TENANT_ID } from '@librechat/data-schemas';
|
||||
import type { StandardGraph } from '@librechat/agents';
|
||||
import { parseTextParts } from 'librechat-data-provider';
|
||||
import type { Agents, TMessageContentParts } from 'librechat-data-provider';
|
||||
import type { StandardGraph } from '@librechat/agents';
|
||||
import type {
|
||||
SerializableJobData,
|
||||
IEventTransport,
|
||||
|
|
@ -9,6 +9,7 @@ import type {
|
|||
AbortResult,
|
||||
IJobStore,
|
||||
} from './interfaces/IJobStore';
|
||||
import type { GenerationJobStore } from '~/app/metrics';
|
||||
import type * as t from '~/types';
|
||||
import {
|
||||
recordGenerationJob,
|
||||
|
|
@ -16,12 +17,85 @@ import {
|
|||
recordGenerationStreamSubscription,
|
||||
setGenerationJobsInFlight,
|
||||
} from '~/app/metrics';
|
||||
import type { GenerationJobStore } from '~/app/metrics';
|
||||
import { InMemoryEventTransport } from './implementations/InMemoryEventTransport';
|
||||
import { InMemoryJobStore } from './implementations/InMemoryJobStore';
|
||||
import { filterPersistableAbortContent } from './abortContent';
|
||||
|
||||
/** Error surfaced to any client still attached when a stale/hung job is reaped. */
|
||||
const REAPED_JOB_ERROR = 'Generation timed out';
|
||||
const OAUTH_TOOL_CALL_PREFIX = `oauth${Constants.mcp_delimiter}`;
|
||||
|
||||
function getToolCallName(toolCall: unknown): unknown {
|
||||
return toolCall != null && typeof toolCall === 'object' && 'name' in toolCall
|
||||
? toolCall.name
|
||||
: undefined;
|
||||
}
|
||||
|
||||
function hasOAuthToolCall(toolCalls: unknown): boolean {
|
||||
return (
|
||||
Array.isArray(toolCalls) &&
|
||||
toolCalls.some((toolCall) => {
|
||||
const name = getToolCallName(toolCall);
|
||||
return typeof name === 'string' && name.startsWith(OAUTH_TOOL_CALL_PREFIX);
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
function getReplayStepId(event: t.ServerSentEvent): unknown {
|
||||
if (!('event' in event) || !event.data || typeof event.data !== 'object') {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (event.event === 'on_run_step' || event.event === 'on_run_step_delta') {
|
||||
return 'id' in event.data ? event.data.id : undefined;
|
||||
}
|
||||
|
||||
if (event.event === 'on_run_step_completed') {
|
||||
const result = 'result' in event.data ? event.data.result : undefined;
|
||||
return result != null && typeof result === 'object' && 'id' in result ? result.id : undefined;
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function isOAuthReplayEvent(event: t.ServerSentEvent): boolean {
|
||||
if (!('event' in event) || !event.data || typeof event.data !== 'object') {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (event.event === 'on_run_step') {
|
||||
const stepDetails = 'stepDetails' in event.data ? event.data.stepDetails : undefined;
|
||||
return (
|
||||
stepDetails != null &&
|
||||
typeof stepDetails === 'object' &&
|
||||
'tool_calls' in stepDetails &&
|
||||
hasOAuthToolCall(stepDetails.tool_calls)
|
||||
);
|
||||
}
|
||||
|
||||
if (event.event === 'on_run_step_delta') {
|
||||
const delta = 'delta' in event.data ? event.data.delta : undefined;
|
||||
if (delta == null || typeof delta !== 'object') {
|
||||
return false;
|
||||
}
|
||||
if (!('tool_calls' in delta) || !hasOAuthToolCall(delta.tool_calls)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
if (event.event === 'on_run_step_completed') {
|
||||
const result = 'result' in event.data ? event.data.result : undefined;
|
||||
if (result == null || typeof result !== 'object' || !('tool_call' in result)) {
|
||||
return false;
|
||||
}
|
||||
const name = getToolCallName(result.tool_call);
|
||||
return typeof name === 'string' && name.startsWith(OAUTH_TOOL_CALL_PREFIX);
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configuration options for GenerationJobManager
|
||||
|
|
@ -100,6 +174,9 @@ class GenerationJobManagerClass {
|
|||
/** Jobs actively generating in this process. */
|
||||
private runningJobs = new Set<string>();
|
||||
|
||||
/** Serializes replay-event read/modify/write updates per stream. */
|
||||
private replayEventWriteQueues = new Map<string, Promise<void>>();
|
||||
|
||||
private cleanupInterval: NodeJS.Timeout | null = null;
|
||||
|
||||
/** Whether we're using Redis stores */
|
||||
|
|
@ -567,6 +644,7 @@ class GenerationJobManagerClass {
|
|||
// Clear content state and run step buffer (Redis only)
|
||||
this.jobStore.clearContentState(streamId);
|
||||
this.runStepBuffers?.delete(streamId);
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
|
||||
// For error jobs, DON'T delete immediately - keep around so late-connecting
|
||||
// clients can receive the error. This handles the race condition where error
|
||||
|
|
@ -652,16 +730,20 @@ class GenerationJobManagerClass {
|
|||
/** Content before clearing state */
|
||||
const result = await this.jobStore.getContentParts(streamId);
|
||||
const content = result?.content ?? [];
|
||||
const abortContent = filterPersistableAbortContent(content);
|
||||
const shouldPersistAbortContent = abortContent.length > 0;
|
||||
|
||||
/** Collected usage for all models */
|
||||
const collectedUsage = this.jobStore.getCollectedUsage(streamId);
|
||||
|
||||
/** Text from content parts for fallback token counting */
|
||||
const text = parseTextParts(content as TMessageContentParts[]);
|
||||
const text = shouldPersistAbortContent
|
||||
? parseTextParts(abortContent as TMessageContentParts[])
|
||||
: '';
|
||||
|
||||
/** Detect "early abort" - aborted before any generation happened (e.g., during tool loading)
|
||||
In this case, no messages were saved to DB, so frontend shouldn't navigate to conversation */
|
||||
const isEarlyAbort = content.length === 0 && !jobData.responseMessageId;
|
||||
const isEarlyAbort = !shouldPersistAbortContent && jobData.createdEventEmitted !== true;
|
||||
|
||||
/** Final event for abort */
|
||||
const userMessageId = jobData.userMessage?.messageId;
|
||||
|
|
@ -686,7 +768,7 @@ class GenerationJobManagerClass {
|
|||
messageId: jobData.responseMessageId ?? `${userMessageId ?? 'aborted'}_`,
|
||||
parentMessageId: userMessageId,
|
||||
conversationId: jobData.conversationId,
|
||||
content,
|
||||
content: abortContent,
|
||||
sender: jobData.sender ?? 'AI',
|
||||
unfinished: true,
|
||||
error: false,
|
||||
|
|
@ -704,6 +786,7 @@ class GenerationJobManagerClass {
|
|||
await this.eventTransport.emitDone(streamId, abortFinalEvent);
|
||||
this.jobStore.clearContentState(streamId);
|
||||
this.runStepBuffers?.delete(streamId);
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
|
||||
// Immediate cleanup if configured (default: true)
|
||||
if (this._cleanupOnComplete) {
|
||||
|
|
@ -726,7 +809,7 @@ class GenerationJobManagerClass {
|
|||
return {
|
||||
success: true,
|
||||
jobData,
|
||||
content,
|
||||
content: abortContent,
|
||||
finalEvent: abortFinalEvent,
|
||||
text,
|
||||
collectedUsage,
|
||||
|
|
@ -962,6 +1045,7 @@ class GenerationJobManagerClass {
|
|||
|
||||
await this.trackUserMessage(streamId, event);
|
||||
await this.trackTitleEvent(streamId, event);
|
||||
await this.trackReplayEvent(streamId, event);
|
||||
|
||||
// For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability)
|
||||
if (this._isRedis) {
|
||||
|
|
@ -1059,6 +1143,71 @@ class GenerationJobManagerClass {
|
|||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Persist replay-only stream events that are needed to reconstruct active
|
||||
* UI state on resume but are not represented by aggregated message content.
|
||||
*/
|
||||
private async trackReplayEvent(streamId: string, event: t.ServerSentEvent): Promise<void> {
|
||||
if (!isOAuthReplayEvent(event)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const previousWrite = this.replayEventWriteQueues.get(streamId) ?? Promise.resolve();
|
||||
const nextWrite = previousWrite
|
||||
.catch(() => {
|
||||
// Keep the queue moving even if a prior replay metadata write failed.
|
||||
})
|
||||
.then(() => this.persistReplayEvent(streamId, event));
|
||||
|
||||
this.replayEventWriteQueues.set(streamId, nextWrite);
|
||||
|
||||
try {
|
||||
await nextWrite;
|
||||
} finally {
|
||||
if (this.replayEventWriteQueues.get(streamId) === nextWrite) {
|
||||
this.replayEventWriteQueues.delete(streamId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async persistReplayEvent(streamId: string, event: t.ServerSentEvent): Promise<void> {
|
||||
const jobData = await this.jobStore.getJob(streamId);
|
||||
if (!jobData) {
|
||||
return;
|
||||
}
|
||||
|
||||
let replayEvents: t.ServerSentEvent[] = [];
|
||||
if (jobData.replayEvents) {
|
||||
try {
|
||||
replayEvents = JSON.parse(jobData.replayEvents) as t.ServerSentEvent[];
|
||||
} catch {
|
||||
replayEvents = [];
|
||||
}
|
||||
}
|
||||
|
||||
const stepId = getReplayStepId(event);
|
||||
const eventName = 'event' in event ? event.event : undefined;
|
||||
const existingIndex =
|
||||
stepId == null
|
||||
? -1
|
||||
: replayEvents.findIndex((candidate) => {
|
||||
if (!('event' in candidate) || candidate.event !== eventName) {
|
||||
return false;
|
||||
}
|
||||
return getReplayStepId(candidate) === stepId;
|
||||
});
|
||||
|
||||
if (existingIndex >= 0) {
|
||||
replayEvents[existingIndex] = event;
|
||||
} else {
|
||||
replayEvents.push(event);
|
||||
}
|
||||
|
||||
await this.jobStore.updateJob(streamId, {
|
||||
replayEvents: JSON.stringify(replayEvents),
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Persist user message metadata from the created event.
|
||||
* Awaited in emitChunk so the HSET commits before the PUBLISH,
|
||||
|
|
@ -1072,6 +1221,7 @@ class GenerationJobManagerClass {
|
|||
|
||||
const { message } = event;
|
||||
const updates: Partial<SerializableJobData> = {
|
||||
createdEventEmitted: true,
|
||||
userMessage: {
|
||||
messageId: message.messageId,
|
||||
parentMessageId: message.parentMessageId,
|
||||
|
|
@ -1176,6 +1326,14 @@ class GenerationJobManagerClass {
|
|||
// Ignore malformed persisted title events.
|
||||
}
|
||||
}
|
||||
let replayEvents: t.ResumeState['replayEvents'];
|
||||
if (jobData.replayEvents) {
|
||||
try {
|
||||
replayEvents = JSON.parse(jobData.replayEvents) as t.ResumeState['replayEvents'];
|
||||
} catch {
|
||||
// Ignore malformed persisted replay events.
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug(`[GenerationJobManager] getResumeState:`, {
|
||||
streamId,
|
||||
|
|
@ -1191,6 +1349,7 @@ class GenerationJobManagerClass {
|
|||
conversationId: jobData.conversationId,
|
||||
sender: jobData.sender,
|
||||
titleEvent,
|
||||
replayEvents,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -1407,6 +1566,7 @@ class GenerationJobManagerClass {
|
|||
this.runningJobs.clear();
|
||||
this.syncRunningJobMetrics();
|
||||
this.runStepBuffers?.clear();
|
||||
this.replayEventWriteQueues.clear();
|
||||
|
||||
logger.debug('[GenerationJobManager] Destroyed');
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,218 @@
|
|||
import type { ServerSentEvent } from '~/types';
|
||||
import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport';
|
||||
import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore';
|
||||
import { GenerationJobManagerClass } from '~/stream/GenerationJobManager';
|
||||
|
||||
jest.spyOn(console, 'log').mockImplementation();
|
||||
|
||||
function createInMemoryManager(): GenerationJobManagerClass {
|
||||
const manager = new GenerationJobManagerClass();
|
||||
manager.configure({
|
||||
jobStore: new InMemoryJobStore({ ttlAfterComplete: 60000 }),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
});
|
||||
manager.initialize();
|
||||
return manager;
|
||||
}
|
||||
|
||||
class SnapshotReplayJobStore extends InMemoryJobStore {
|
||||
async getJob(streamId: string) {
|
||||
const job = await super.getJob(streamId);
|
||||
return job ? { ...job } : null;
|
||||
}
|
||||
|
||||
async updateJob(streamId: string, updates: Parameters<InMemoryJobStore['updateJob']>[1]) {
|
||||
if (updates.replayEvents) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 5));
|
||||
}
|
||||
await super.updateJob(streamId, updates);
|
||||
}
|
||||
}
|
||||
|
||||
function createSnapshotReplayManager(): GenerationJobManagerClass {
|
||||
const manager = new GenerationJobManagerClass();
|
||||
manager.configure({
|
||||
jobStore: new SnapshotReplayJobStore({ ttlAfterComplete: 60000 }),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
});
|
||||
manager.initialize();
|
||||
return manager;
|
||||
}
|
||||
|
||||
describe('GenerationJobManager resume replay events', () => {
|
||||
let manager: GenerationJobManagerClass | undefined;
|
||||
|
||||
afterEach(async () => {
|
||||
await manager?.destroy();
|
||||
manager = undefined;
|
||||
});
|
||||
|
||||
test('includes OAuth run step and delta replay events in resume state', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `oauth-delta-resume-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
const runStepEvent = {
|
||||
event: 'on_run_step',
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
runId: 'USE_PRELIM_RESPONSE_MESSAGE_ID',
|
||||
index: 0,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-oauth', name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
},
|
||||
},
|
||||
} satisfies ServerSentEvent;
|
||||
const authEvent = {
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
} satisfies ServerSentEvent;
|
||||
|
||||
await manager.emitChunk(streamId, runStepEvent);
|
||||
await manager.emitChunk(streamId, authEvent);
|
||||
await manager.emitChunk(streamId, {
|
||||
event: 'on_run_step',
|
||||
data: {
|
||||
id: 'step-regular',
|
||||
runId: 'response-1',
|
||||
index: 1,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: 'call-regular', name: 'regular-tool', args: '{}' }],
|
||||
},
|
||||
},
|
||||
});
|
||||
await manager.emitChunk(streamId, {
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: 'step-no-auth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ name: 'regular-tool', args: '{}' }],
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
|
||||
expect(resumeState?.replayEvents).toEqual([runStepEvent, authEvent]);
|
||||
});
|
||||
|
||||
test('replaces OAuth replay event for the same step id', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `oauth-delta-replace-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
await manager.emitChunk(streamId, {
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
tool_calls: [{ name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/first',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const replacementEvent = {
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: 'step-oauth',
|
||||
delta: {
|
||||
tool_calls: [{ name: 'oauth_mcp_Google-Workspace', args: '' }],
|
||||
auth: 'https://auth.example.com/latest',
|
||||
expires_at: 1780792000,
|
||||
},
|
||||
},
|
||||
} satisfies ServerSentEvent;
|
||||
|
||||
await manager.emitChunk(streamId, replacementEvent);
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
|
||||
expect(resumeState?.replayEvents).toEqual([replacementEvent]);
|
||||
});
|
||||
|
||||
test('does not replay non-MCP action auth deltas', async () => {
|
||||
manager = createInMemoryManager();
|
||||
const streamId = `action-auth-delta-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
await manager.emitChunk(streamId, {
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: 'step-action-auth',
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ name: 'google_calendar_action_api_example_com', args: '' }],
|
||||
auth: 'https://auth.example.com/action',
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
|
||||
expect(resumeState?.replayEvents).toBeUndefined();
|
||||
});
|
||||
|
||||
test('serializes replay event updates for concurrent MCP OAuth prompts', async () => {
|
||||
manager = createSnapshotReplayManager();
|
||||
const streamId = `oauth-delta-concurrent-${Date.now()}`;
|
||||
await manager.createJob(streamId, 'user-1', streamId);
|
||||
|
||||
const createRunStepEvent = (serverName: string, index: number) =>
|
||||
({
|
||||
event: 'on_run_step',
|
||||
data: {
|
||||
id: `step-oauth-${serverName}`,
|
||||
runId: 'USE_PRELIM_RESPONSE_MESSAGE_ID',
|
||||
index,
|
||||
stepDetails: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ id: `call-${serverName}`, name: `oauth_mcp_${serverName}`, args: '' }],
|
||||
},
|
||||
},
|
||||
}) satisfies ServerSentEvent;
|
||||
|
||||
const createAuthEvent = (serverName: string) =>
|
||||
({
|
||||
event: 'on_run_step_delta',
|
||||
data: {
|
||||
id: `step-oauth-${serverName}`,
|
||||
delta: {
|
||||
type: 'tool_calls',
|
||||
tool_calls: [{ name: `oauth_mcp_${serverName}`, args: '' }],
|
||||
auth: `https://auth.example.com/${serverName}`,
|
||||
expires_at: 1780791946,
|
||||
},
|
||||
},
|
||||
}) satisfies ServerSentEvent;
|
||||
|
||||
const events = [
|
||||
createRunStepEvent('Google-Workspace', 0),
|
||||
createAuthEvent('Google-Workspace'),
|
||||
createRunStepEvent('clickhouse-docs', 1),
|
||||
createAuthEvent('clickhouse-docs'),
|
||||
];
|
||||
|
||||
await Promise.all(events.map((event) => manager!.emitChunk(streamId, event)));
|
||||
|
||||
const resumeState = await manager.getResumeState(streamId);
|
||||
|
||||
expect(resumeState?.replayEvents).toHaveLength(events.length);
|
||||
expect(resumeState?.replayEvents).toEqual(expect.arrayContaining(events));
|
||||
});
|
||||
});
|
||||
|
|
@ -282,6 +282,182 @@ describe('AbortJob - Text and CollectedUsage', () => {
|
|||
await GenerationJobManager.destroy();
|
||||
});
|
||||
|
||||
it('should treat OAuth-only content as an early abort', async () => {
|
||||
const { GenerationJobManager } = await import('../GenerationJobManager');
|
||||
const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore');
|
||||
const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport');
|
||||
|
||||
GenerationJobManager.configure({
|
||||
jobStore: new InMemoryJobStore(),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
cleanupOnComplete: false,
|
||||
});
|
||||
|
||||
GenerationJobManager.initialize();
|
||||
|
||||
const streamId = `oauth-abort-${Date.now()}`;
|
||||
await GenerationJobManager.createJob(streamId, 'user-1');
|
||||
|
||||
GenerationJobManager.setContentParts(streamId, [
|
||||
{
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
type: 'tool_call',
|
||||
id: 'oauth-call-1',
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
args: '',
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
},
|
||||
},
|
||||
] as never);
|
||||
|
||||
const abortResult = await GenerationJobManager.abortJob(streamId);
|
||||
|
||||
expect(abortResult.success).toBe(true);
|
||||
expect(abortResult.content).toEqual([]);
|
||||
expect(abortResult.text).toBe('');
|
||||
expect(abortResult.finalEvent).toEqual(
|
||||
expect.objectContaining({
|
||||
earlyAbort: true,
|
||||
responseMessage: null,
|
||||
}),
|
||||
);
|
||||
|
||||
await GenerationJobManager.destroy();
|
||||
});
|
||||
|
||||
it('should keep post-created OAuth-only aborts tied to the created turn', async () => {
|
||||
const { GenerationJobManager } = await import('../GenerationJobManager');
|
||||
const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore');
|
||||
const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport');
|
||||
|
||||
GenerationJobManager.configure({
|
||||
jobStore: new InMemoryJobStore(),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
cleanupOnComplete: false,
|
||||
});
|
||||
|
||||
GenerationJobManager.initialize();
|
||||
|
||||
const streamId = `created-oauth-abort-${Date.now()}`;
|
||||
await GenerationJobManager.createJob(streamId, 'user-1');
|
||||
await GenerationJobManager.updateMetadata(streamId, {
|
||||
responseMessageId: 'response-message-1',
|
||||
userMessage: {
|
||||
messageId: 'user-message-1',
|
||||
parentMessageId: 'parent-message-1',
|
||||
conversationId: streamId,
|
||||
text: 'Use Google Workspace',
|
||||
},
|
||||
});
|
||||
await GenerationJobManager.emitChunk(streamId, {
|
||||
created: true,
|
||||
streamId,
|
||||
message: {
|
||||
messageId: 'user-message-1',
|
||||
parentMessageId: 'parent-message-1',
|
||||
conversationId: streamId,
|
||||
text: 'Use Google Workspace',
|
||||
sender: 'User',
|
||||
isCreatedByUser: true,
|
||||
},
|
||||
});
|
||||
|
||||
GenerationJobManager.setContentParts(streamId, [
|
||||
{
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
type: 'tool_call',
|
||||
id: 'oauth-call-1',
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
args: '',
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
},
|
||||
},
|
||||
] as never);
|
||||
|
||||
const abortResult = await GenerationJobManager.abortJob(streamId);
|
||||
|
||||
expect(abortResult.success).toBe(true);
|
||||
expect(abortResult.content).toEqual([]);
|
||||
expect(abortResult.text).toBe('');
|
||||
expect(abortResult.finalEvent).toEqual(
|
||||
expect.objectContaining({
|
||||
earlyAbort: false,
|
||||
conversation: { conversationId: streamId },
|
||||
requestMessage: expect.objectContaining({
|
||||
messageId: 'user-message-1',
|
||||
}),
|
||||
responseMessage: expect.objectContaining({
|
||||
messageId: 'response-message-1',
|
||||
content: [],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
await GenerationJobManager.destroy();
|
||||
});
|
||||
|
||||
it('should strip OAuth prompts but persist real abort content', async () => {
|
||||
const { GenerationJobManager } = await import('../GenerationJobManager');
|
||||
const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore');
|
||||
const { InMemoryEventTransport } = await import('../implementations/InMemoryEventTransport');
|
||||
|
||||
GenerationJobManager.configure({
|
||||
jobStore: new InMemoryJobStore(),
|
||||
eventTransport: new InMemoryEventTransport(),
|
||||
isRedis: false,
|
||||
cleanupOnComplete: false,
|
||||
});
|
||||
|
||||
GenerationJobManager.initialize();
|
||||
|
||||
const streamId = `mixed-oauth-abort-${Date.now()}`;
|
||||
await GenerationJobManager.createJob(streamId, 'user-1');
|
||||
await GenerationJobManager.updateMetadata(streamId, {
|
||||
responseMessageId: 'response-message-1',
|
||||
userMessage: {
|
||||
messageId: 'user-message-1',
|
||||
parentMessageId: 'parent-message-1',
|
||||
conversationId: streamId,
|
||||
text: 'Use Google Workspace',
|
||||
},
|
||||
});
|
||||
|
||||
const textPart = { type: 'text', text: 'Partial response...' };
|
||||
GenerationJobManager.setContentParts(streamId, [
|
||||
{
|
||||
type: 'tool_call',
|
||||
tool_call: {
|
||||
type: 'tool_call',
|
||||
id: 'oauth-call-1',
|
||||
name: 'oauth_mcp_Google-Workspace',
|
||||
args: '',
|
||||
auth: 'https://auth.example.com/oauth',
|
||||
},
|
||||
},
|
||||
textPart,
|
||||
] as never);
|
||||
|
||||
const abortResult = await GenerationJobManager.abortJob(streamId);
|
||||
|
||||
expect(abortResult.success).toBe(true);
|
||||
expect(abortResult.content).toEqual([textPart]);
|
||||
expect(abortResult.text).toBe('Partial response...');
|
||||
expect(abortResult.finalEvent).toEqual(
|
||||
expect.objectContaining({
|
||||
earlyAbort: false,
|
||||
responseMessage: expect.objectContaining({
|
||||
content: [textPart],
|
||||
}),
|
||||
}),
|
||||
);
|
||||
|
||||
await GenerationJobManager.destroy();
|
||||
});
|
||||
|
||||
it('should return both text and collectedUsage on abort', async () => {
|
||||
const { GenerationJobManager } = await import('../GenerationJobManager');
|
||||
const { InMemoryJobStore } = await import('../implementations/InMemoryJobStore');
|
||||
|
|
|
|||
80
packages/api/src/stream/abortContent.ts
Normal file
80
packages/api/src/stream/abortContent.ts
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
import { Constants, ContentTypes } from 'librechat-data-provider';
|
||||
import type { Agents, TMessageContentParts } from 'librechat-data-provider';
|
||||
|
||||
const OAUTH_TOOL_CALL_PREFIX = `oauth${Constants.mcp_delimiter}`;
|
||||
|
||||
type PersistableAbortContentPart = Agents.MessageContentComplex | TMessageContentParts;
|
||||
type AbortContentPart = PersistableAbortContentPart | null | undefined;
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return value != null && typeof value === 'object';
|
||||
}
|
||||
|
||||
function getTextValue(value: unknown): string {
|
||||
if (typeof value === 'string') {
|
||||
return value;
|
||||
}
|
||||
if (isRecord(value) && typeof value.value === 'string') {
|
||||
return value.value;
|
||||
}
|
||||
return '';
|
||||
}
|
||||
|
||||
function getToolCall(part: Record<string, unknown>): Record<string, unknown> | null {
|
||||
if (isRecord(part.tool_call)) {
|
||||
return part.tool_call;
|
||||
}
|
||||
|
||||
const keyedToolCall = part[ContentTypes.TOOL_CALL];
|
||||
return isRecord(keyedToolCall) ? keyedToolCall : null;
|
||||
}
|
||||
|
||||
function isOAuthPromptContentPart(part: AbortContentPart): boolean {
|
||||
if (!isRecord(part) || part.type !== ContentTypes.TOOL_CALL) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const partRecord = part as Record<string, unknown>;
|
||||
const toolCall = getToolCall(partRecord);
|
||||
if (!toolCall) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const auth = typeof toolCall.auth === 'string' ? toolCall.auth : partRecord.auth;
|
||||
if (typeof auth === 'string' && auth.trim().length > 0) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const name = toolCall.name;
|
||||
return typeof name === 'string' && name.startsWith(OAUTH_TOOL_CALL_PREFIX);
|
||||
}
|
||||
|
||||
function isPersistableAbortContentPart(
|
||||
part: AbortContentPart,
|
||||
): part is PersistableAbortContentPart {
|
||||
if (!isRecord(part)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (isOAuthPromptContentPart(part)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (part.type === ContentTypes.TEXT) {
|
||||
return getTextValue(part.text).trim().length > 0;
|
||||
}
|
||||
|
||||
if (part.type === ContentTypes.THINK) {
|
||||
return getTextValue(part.think).trim().length > 0;
|
||||
}
|
||||
|
||||
return Object.keys(part).length > 0;
|
||||
}
|
||||
|
||||
export function hasPersistableAbortContent(content: unknown): boolean {
|
||||
return filterPersistableAbortContent(content).length > 0;
|
||||
}
|
||||
|
||||
export function filterPersistableAbortContent(content: unknown): PersistableAbortContentPart[] {
|
||||
return Array.isArray(content) ? content.filter(isPersistableAbortContentPart) : [];
|
||||
}
|
||||
|
|
@ -913,6 +913,7 @@ export class RedisJobStore implements IJobStore {
|
|||
error: data.error || undefined,
|
||||
userMessage: data.userMessage ? JSON.parse(data.userMessage) : undefined,
|
||||
responseMessageId: data.responseMessageId || undefined,
|
||||
createdEventEmitted: data.createdEventEmitted === '1',
|
||||
sender: data.sender || undefined,
|
||||
syncSent: data.syncSent === '1',
|
||||
finalEvent: data.finalEvent || undefined,
|
||||
|
|
@ -920,6 +921,8 @@ export class RedisJobStore implements IJobStore {
|
|||
iconURL: data.iconURL || undefined,
|
||||
model: data.model || undefined,
|
||||
promptTokens: data.promptTokens ? parseInt(data.promptTokens, 10) : undefined,
|
||||
titleEvent: data.titleEvent || undefined,
|
||||
replayEvents: data.replayEvents || undefined,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ export type {
|
|||
|
||||
export { createStreamServices } from './createStreamServices';
|
||||
export type { StreamServicesConfig, StreamServices } from './createStreamServices';
|
||||
export { filterPersistableAbortContent, hasPersistableAbortContent } from './abortContent';
|
||||
|
||||
// Implementations (for advanced use cases)
|
||||
export { InMemoryJobStore } from './implementations/InMemoryJobStore';
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import type { Agents } from 'librechat-data-provider';
|
||||
import type { StandardGraph } from '@librechat/agents';
|
||||
import type { Agents } from 'librechat-data-provider';
|
||||
|
||||
/**
|
||||
* Job status enum
|
||||
|
|
@ -30,6 +30,9 @@ export interface SerializableJobData {
|
|||
/** Response message ID for reconnection */
|
||||
responseMessageId?: string;
|
||||
|
||||
/** Whether the user-message created event has been emitted */
|
||||
createdEventEmitted?: boolean;
|
||||
|
||||
/** Sender name for UI display */
|
||||
sender?: string;
|
||||
|
||||
|
|
@ -42,6 +45,9 @@ export interface SerializableJobData {
|
|||
/** Serialized title event for replay during active-stream resume */
|
||||
titleEvent?: string;
|
||||
|
||||
/** Serialized replay-only stream events for active-stream resume */
|
||||
replayEvents?: string;
|
||||
|
||||
/** Endpoint metadata for abort handling - avoids storing functions */
|
||||
endpoint?: string;
|
||||
iconURL?: string;
|
||||
|
|
@ -149,6 +155,11 @@ export interface ResumeState {
|
|||
title?: string;
|
||||
};
|
||||
};
|
||||
replayEvents?: Array<{
|
||||
event: string;
|
||||
data?: unknown;
|
||||
[key: string]: unknown;
|
||||
}>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -147,6 +147,8 @@ export type TSubmission = {
|
|||
isContinued?: boolean;
|
||||
isTemporary: boolean;
|
||||
messages: TMessage[];
|
||||
/** Client-only full message context used to restore branch siblings after scoped regenerate. */
|
||||
regenerateMessages?: TMessage[];
|
||||
isRegenerate?: boolean;
|
||||
initialResponse?: TMessage;
|
||||
conversation: Partial<TConversation>;
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
/* eslint-disable @typescript-eslint/no-namespace */
|
||||
import { StepTypes, ContentTypes, ToolCallTypes } from './runs';
|
||||
import type { FunctionToolCall, SummaryContentPart } from './assistants';
|
||||
import type { TAttachment, TPlugin } from 'src/schemas';
|
||||
import { StepTypes, ContentTypes, ToolCallTypes } from './runs';
|
||||
|
||||
export namespace Agents {
|
||||
export type MessageType = 'human' | 'ai' | 'generic' | 'system' | 'function' | 'tool' | 'remove';
|
||||
|
|
@ -224,6 +224,11 @@ export namespace Agents {
|
|||
title?: string;
|
||||
};
|
||||
};
|
||||
replayEvents?: Array<{
|
||||
event: string;
|
||||
data?: unknown;
|
||||
[key: string]: unknown;
|
||||
}>;
|
||||
}
|
||||
/**
|
||||
* Represents a run step delta i.e. any changed fields on a run step during
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue