From d8427ffc5e0570a46e5b3ec2387df0b235faad0f Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 26 Jul 2026 21:58:25 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=82=20test:=20Cover=20Tool=20Approval?= =?UTF-8?q?=20Workflows=20End=20to=20End=20(#14427)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: cover tool approval workflows end to end * fix: preserve tool approval state across resume * fix: preserve agent context in mock stream responses * fix: preserve nested approvals in collapsed groups --- api/app/clients/tools/util/handleTools.js | 1 + .../clients/tools/util/handleTools.test.js | 3 + .../agents/__tests__/callbacks.spec.js | 58 ++ .../agents/__tests__/resume.spec.js | 27 + api/server/controllers/agents/callbacks.js | 106 ++- api/server/controllers/agents/client.js | 37 +- api/server/controllers/agents/client.test.js | 23 + api/server/controllers/agents/request.js | 60 +- api/server/controllers/agents/resume.js | 13 +- api/server/services/ActionService.js | 10 +- api/server/services/ActionService.spec.js | 115 ++- .../services/Endpoints/agents/initialize.js | 19 +- .../Endpoints/agents/initialize.spec.js | 48 + api/server/services/MCP.js | 44 +- api/server/services/MCP.spec.js | 110 ++- api/server/services/ToolService.js | 48 +- api/server/services/Tools/search.js | 16 +- api/server/services/Tools/search.spec.js | 63 ++ .../services/__tests__/ToolService.spec.js | 98 ++ api/typedefs.js | 1 + .../Chat/Messages/Content/ApprovalContext.tsx | 16 +- .../components/Chat/Messages/Content/Part.tsx | 1 + .../Messages/Content/Parts/SubagentCall.tsx | 1 + .../Chat/Messages/Content/ToolApproval.tsx | 6 +- .../Chat/Messages/Content/ToolCall.tsx | 10 +- .../Chat/Messages/Content/ToolCallGroup.tsx | 57 +- .../__tests__/ApprovalContext.test.tsx | 76 ++ .../Content/__tests__/ToolCallGroup.test.tsx | 193 ++++ client/src/utils/approval.spec.ts | 31 + client/src/utils/approval.ts | 4 + e2e/config/librechat.e2e.yaml | 12 + e2e/setup/fake-mcp-server.js | 33 + e2e/setup/fake-model.js | 124 ++- e2e/setup/tool-approval-hook.js | 29 + e2e/specs/mock/tool-approvals.spec.ts | 835 ++++++++++++++++++ packages/api/src/agents/hitl/resume.spec.ts | 138 +++ packages/api/src/agents/hitl/resume.ts | 77 ++ packages/api/src/agents/memory.spec.ts | 72 +- packages/api/src/agents/memory.ts | 25 +- .../api/src/stream/GenerationJobManager.ts | 89 +- .../GenerationJobManager.resumeReplay.spec.ts | 240 ++++- .../RedisJobStore.stream_integration.spec.ts | 136 +++ 42 files changed, 2960 insertions(+), 145 deletions(-) create mode 100644 api/server/services/Tools/search.spec.js create mode 100644 client/src/components/Chat/Messages/Content/__tests__/ApprovalContext.test.tsx create mode 100644 e2e/setup/tool-approval-hook.js create mode 100644 e2e/specs/mock/tool-approvals.spec.ts diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index 56eb345b8f..38a8d8b26b 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -503,6 +503,7 @@ const loadTools = async ({ requestScopedConnections, res: options.res, streamId: options.req?._resumableStreamId || null, + jobCreatedAt: options.jobCreatedAt, model: agent?.model ?? model, serverName: config.serverName, provider: agent?.provider ?? endpoint, diff --git a/api/app/clients/tools/util/handleTools.test.js b/api/app/clients/tools/util/handleTools.test.js index a812443faa..50be5798ec 100644 --- a/api/app/clients/tools/util/handleTools.test.js +++ b/api/app/clients/tools/util/handleTools.test.js @@ -319,6 +319,7 @@ describe('Tool Handlers', () => { const serverName = 'body-scoped'; const toolKey = `search${Constants.mcp_delimiter}${serverName}`; const requestBody = { conversationId: 'conv-123', messageId: 'msg-123' }; + const jobCreatedAt = 1234; const serverConfig = { type: 'streamable-http', url: 'https://api.example.com/messages/{{LIBRECHAT_BODY_MESSAGEID}}/mcp', @@ -336,6 +337,7 @@ describe('Tool Handlers', () => { user: { id: fakeUser._id.toString(), role: 'USER' }, body: requestBody, }, + jobCreatedAt, }, }); @@ -348,6 +350,7 @@ describe('Tool Handlers', () => { expect(mockCreateMCPTool).toHaveBeenCalledWith( expect.objectContaining({ requestBody, + jobCreatedAt, toolKey, config: serverConfig, }), diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index 433e1dc2e8..a183a212f5 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -7,6 +7,9 @@ jest.mock('nanoid', () => ({ jest.mock('@librechat/api', () => ({ sendEvent: jest.fn(), + GenerationJobManager: { + emitChunk: jest.fn(), + }, HOST_FILE_AUTHORING_ARTIFACT_KEY: '__librechat_file_authoring', getToolInputValidationDetails: jest.fn((result, validationError) => validationError != null @@ -74,6 +77,61 @@ jest.mock('~/server/services/Files/process', () => ({ saveBase64Image: jest.fn(), })); +describe('resumable event generation fencing', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('forwards the originating job epoch with run-step events', async () => { + const { GenerationJobManager } = require('@librechat/api'); + const { GraphEvents } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + const data = { + id: 'step-1', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-1', name: 'approval_probe', args: '{}' }], + }, + }; + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent: jest.fn(), + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'conversation-1', + jobCreatedAt: 1234, + }); + + await handlers[GraphEvents.ON_RUN_STEP].handle(GraphEvents.ON_RUN_STEP, data); + + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'conversation-1', + { event: GraphEvents.ON_RUN_STEP, data }, + { expectedCreatedAt: 1234 }, + ); + }); + + it('forwards the originating job epoch with deferred attachments', () => { + const { GenerationJobManager } = require('@librechat/api'); + const { createAttachmentEmitter } = require('../callbacks'); + const attachment = { file_id: 'file-1', status: 'ready' }; + const emitAttachment = createAttachmentEmitter({ + res: { write: jest.fn() }, + streamId: 'conversation-1', + jobCreatedAt: 1234, + }); + + emitAttachment(attachment); + + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'conversation-1', + { event: 'attachment', data: attachment }, + { expectedCreatedAt: 1234 }, + ); + }); +}); + describe('createToolEndCallback', () => { let req, res, artifactPromises, createToolEndCallback; let logger; diff --git a/api/server/controllers/agents/__tests__/resume.spec.js b/api/server/controllers/agents/__tests__/resume.spec.js index 5ad47f8912..a52181feb0 100644 --- a/api/server/controllers/agents/__tests__/resume.spec.js +++ b/api/server/controllers/agents/__tests__/resume.spec.js @@ -646,6 +646,33 @@ describe('ResumeAgentController (POST /agents/chat/resume)', () => { ); }); + it('passes persisted run steps into the rebuilt run for tool-result correlation', async () => { + const runSteps = [ + { + id: 'step-approval', + index: 1, + type: 'tool_calls', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'tc1', name: 'approval_probe', args: '{}' }], + }, + usage: null, + }, + ]; + mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); + mockGenerationJobManager.getResumeState.mockResolvedValue({ + aggregatedContent: [], + runSteps, + }); + + await post(approveBody()); + await settled; + await flush(); + + const client = await mockInitializeClient.mock.results[0].value.then((r) => r.client); + expect(client.resumeCompletion).toHaveBeenCalledWith(expect.objectContaining({ runSteps })); + }); + it('restores the paused user message files before reconstruction (execute-code files)', async () => { mockGenerationJobManager.getJob.mockResolvedValue(makeToolApprovalJob()); // The resume body carries no files; the controller must source them from the diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 1c04253481..f451954c5b 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -222,11 +222,12 @@ function checkIfLastAgent(last_agent_id, langgraph_node) { * @param {ServerResponse} res - The server response object * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode * @param {Object} eventData - The event data to send + * @param {number} [expectedCreatedAt] - The generation epoch that produced the event * @returns {Promise} */ -async function emitEvent(res, streamId, eventData) { +async function emitEvent(res, streamId, eventData, expectedCreatedAt) { if (streamId) { - await GenerationJobManager.emitChunk(streamId, eventData); + await GenerationJobManager.emitChunk(streamId, eventData, { expectedCreatedAt }); } else { sendEvent(res, eventData); } @@ -239,13 +240,12 @@ async function emitEvent(res, streamId, eventData) { * running state. Only signals while a fired prewarm remains unresolved * ({@link shouldSignalSandboxStart}); stateless deployments never fire one * and completed boots clear the marker, so both stay on the generic label. - * @param {ServerResponse} res - The server response object - * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode + * @param {(eventData: Object) => Promise} emitForJob - Generation-fenced event emitter * @param {StreamEventData} data - The `on_run_step` event data * @param {GraphRunnableConfig['configurable']} [metadata] The runnable metadata * @returns {Promise} */ -async function maybeEmitSandboxStarting(res, streamId, data, metadata) { +async function maybeEmitSandboxStarting(emitForJob, data, metadata) { const conversationId = metadata?.thread_id; if (!conversationId || !(await shouldSignalSandboxStart(conversationId))) { return; @@ -256,7 +256,7 @@ async function maybeEmitSandboxStarting(res, streamId, data, metadata) { if (!toolCall?.id || name == null || !isCodeSessionToolName(name)) { continue; } - await emitEvent(res, streamId, { + await emitForJob({ event: StepEvents.ON_SANDBOX_STARTING, data: { tool_call_id: toolCall.id, runId: metadata?.run_id }, }); @@ -323,6 +323,7 @@ function feedSubagentAggregator(aggregator, event) { * @param {ToolEndCallback} options.toolEndCallback - Callback to use when tool ends. * @param {Array} options.collectedUsage - The list of collected usage metadata. * @param {string | null} [options.streamId] - The stream ID for resumable mode, or null for standard mode. + * @param {number} [options.jobCreatedAt] - The generation epoch that owns emitted events. * @param {ToolExecuteOptions} [options.toolExecuteOptions] - Options for event-driven tool execution. * @param {UsageCostDeps} [options.usageCost] - Pricing context for authoritative per-event cost. * @param {{ latest: TContextUsageEvent | null, count: number }} [options.contextUsageSink] - Mutable @@ -343,6 +344,7 @@ function getDefaultHandlers({ collectedUsage, collectedThoughtSignatures = null, streamId = null, + jobCreatedAt, toolExecuteOptions = null, summarizationOptions = null, subagentAggregatorsByToolCallId = null, @@ -355,6 +357,7 @@ function getDefaultHandlers({ `[getDefaultHandlers] Missing required options: res: ${!res}, aggregateContent: ${!aggregateContent}`, ); } + const emitForJob = (eventData) => emitEvent(res, streamId, eventData, jobCreatedAt); /** * Emit a token-usage event, attaching the authoritative per-event USD cost * when cost display is enabled. The backend is the single source of truth @@ -385,7 +388,7 @@ function getDefaultHandlers({ if (usageEmitSink) { usageEmitSink.push(payload); } - return emitEvent(res, streamId, { event: UsageEvents.ON_TOKEN_USAGE, data: payload }); + return emitForJob({ event: UsageEvents.ON_TOKEN_USAGE, data: payload }); }; const handlers = { [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler( @@ -404,17 +407,17 @@ function getDefaultHandlers({ handle: async (event, data, metadata) => { aggregateContent({ event, data }); if (data?.stepDetails.type === StepTypes.TOOL_CALLS) { - await emitEvent(res, streamId, { event, data }); - await maybeEmitSandboxStarting(res, streamId, data, metadata); + await emitForJob({ event, data }); + await maybeEmitSandboxStarting(emitForJob, data, metadata); } else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (!metadata?.hide_sequential_outputs) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else { const agentName = metadata?.name ?? 'Agent'; const isToolCall = data?.stepDetails.type === StepTypes.TOOL_CALLS; const action = isToolCall ? 'performing a task...' : 'thinking...'; - await emitEvent(res, streamId, { + await emitForJob({ event: 'on_agent_update', data: { runId: metadata?.run_id, @@ -434,11 +437,11 @@ function getDefaultHandlers({ handle: async (event, data, metadata) => { aggregateContent({ event, data }); if (data?.delta.type === StepTypes.TOOL_CALLS) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (!metadata?.hide_sequential_outputs) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } }, }, @@ -477,11 +480,11 @@ function getDefaultHandlers({ } } if (data?.result != null) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (!metadata?.hide_sequential_outputs) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } }, }, @@ -495,9 +498,9 @@ function getDefaultHandlers({ handle: async (event, data, metadata) => { aggregateContent({ event, data }); if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (!metadata?.hide_sequential_outputs) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } }, }, @@ -511,9 +514,9 @@ function getDefaultHandlers({ handle: async (event, data, metadata) => { aggregateContent({ event, data }); if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } else if (!metadata?.hide_sequential_outputs) { - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } }, }, @@ -568,14 +571,14 @@ function getDefaultHandlers({ ); } } - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); }, }; if (summarizationOptions?.enabled !== false) { handlers[GraphEvents.ON_SUMMARIZE_START] = { handle: async (_event, data) => { - await emitEvent(res, streamId, { + await emitForJob({ event: GraphEvents.ON_SUMMARIZE_START, data, }); @@ -584,7 +587,7 @@ function getDefaultHandlers({ handlers[GraphEvents.ON_SUMMARIZE_DELTA] = { handle: async (_event, data) => { aggregateContent({ event: GraphEvents.ON_SUMMARIZE_DELTA, data }); - await emitEvent(res, streamId, { + await emitForJob({ event: GraphEvents.ON_SUMMARIZE_DELTA, data, }); @@ -593,7 +596,7 @@ function getDefaultHandlers({ handlers[GraphEvents.ON_SUMMARIZE_COMPLETE] = { handle: async (_event, data) => { aggregateContent({ event: GraphEvents.ON_SUMMARIZE_COMPLETE, data }); - await emitEvent(res, streamId, { + await emitForJob({ event: GraphEvents.ON_SUMMARIZE_COMPLETE, data, }); @@ -634,7 +637,7 @@ function getDefaultHandlers({ contextUsageSink.count = (contextUsageSink.count ?? 0) + 1; contextUsageSink.latestUsageIndex = usageEmitSink?.length ?? 0; } - await emitEvent(res, streamId, { event, data }); + await emitForJob({ event, data }); } }, }; @@ -649,10 +652,15 @@ function getDefaultHandlers({ * @param {ServerResponse} res - The server response object * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode * @param {Object} attachment - The attachment data + * @param {number} [expectedCreatedAt] - The generation epoch that produced the attachment */ -function writeAttachment(res, streamId, attachment) { +function writeAttachment(res, streamId, attachment, expectedCreatedAt) { if (streamId) { - GenerationJobManager.emitChunk(streamId, { event: 'attachment', data: attachment }); + GenerationJobManager.emitChunk( + streamId, + { event: 'attachment', data: attachment }, + { expectedCreatedAt }, + ); } else { res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`); } @@ -699,12 +707,13 @@ function isStreamWritable(res, streamId) { * @param {ServerResponse} res * @param {string | null} streamId * @param {Object} attachment - Updated attachment payload (must carry `file_id`). + * @param {number} [expectedCreatedAt] - The generation epoch that produced the attachment */ -function writeAttachmentUpdate(res, streamId, attachment) { +function writeAttachmentUpdate(res, streamId, attachment, expectedCreatedAt) { if (!isStreamWritable(res, streamId)) { return; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, expectedCreatedAt); } /** @@ -714,9 +723,10 @@ function writeAttachmentUpdate(res, streamId, attachment) { * @param {ServerResponse} params.res * @param {Promise[]} params.artifactPromises * @param {string | null} [params.streamId] - The stream ID for resumable mode, or null for standard mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted attachments. * @returns {ToolEndCallback} The tool end callback. */ -function createToolEndCallback({ req, res, artifactPromises, streamId = null }) { +function createToolEndCallback({ req, res, artifactPromises, streamId = null, jobCreatedAt }) { /** * @type {ToolEndCallback} */ @@ -747,7 +757,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) if (!streamId && !res.headersSent) { return attachment; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); return attachment; })().catch((error) => { logger.error('Error processing file citations:', error); @@ -769,7 +779,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) if (!streamId && !res.headersSent) { return attachment; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); return attachment; })().catch((error) => { logger.error('Error processing artifact content:', error); @@ -791,7 +801,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) if (!streamId && !res.headersSent) { return attachment; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); return attachment; })().catch((error) => { logger.error('Error processing artifact content:', error); @@ -813,7 +823,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) if (!streamId && !res.headersSent) { return attachment; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); return attachment; })().catch((error) => { logger.error('Error processing memory artifact content:', error); @@ -858,7 +868,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) return null; } - writeAttachment(res, streamId, fileMetadata); + writeAttachment(res, streamId, fileMetadata, jobCreatedAt); return fileMetadata; })().catch((error) => { logger.error('Error processing artifact content:', error); @@ -935,7 +945,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) * IIFE catch but logged as noise). Same gate the Responses * path uses below. */ if (isStreamWritable(res, streamId)) { - writeAttachment(res, streamId, fileMetadata); + writeAttachment(res, streamId, fileMetadata, jobCreatedAt); } /* Deferred preview rendering: extraction continues running * even after the HTTP response closes. If the stream is still @@ -958,11 +968,16 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) fileId: fileMetadata.file_id, previewRevision: result?.previewRevision, onResolved: (updated) => { - writeAttachmentUpdate(res, streamId, { - ...updated, - messageId: metadata.run_id, - toolCallId, - }); + writeAttachmentUpdate( + res, + streamId, + { + ...updated, + messageId: metadata.run_id, + toolCallId, + }, + jobCreatedAt, + ); }, }); return fileMetadata; @@ -983,14 +998,15 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) * @param {Object} params * @param {ServerResponse} params.res * @param {string | null} [params.streamId] + * @param {number} [params.jobCreatedAt] * @returns {(attachment: Object) => void} */ -function createAttachmentEmitter({ res, streamId = null }) { +function createAttachmentEmitter({ res, streamId = null, jobCreatedAt }) { return (attachment) => { if (!attachment || !isStreamWritable(res, streamId)) { return; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); }; } diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 7db770bf37..f1cb8c783a 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -44,6 +44,7 @@ const { deleteAgentCheckpoint, agentRequestsAskUserQuestion, attachAskUserQuestionArgs, + hydrateResumeRunSteps, createContentIndexOffsetHandlers, createSteerIndexOffsetHandlers, createSteerDrainHook, @@ -140,6 +141,7 @@ class AgentClient extends BaseClient { const { agentConfigs, contentParts, + stepMap, collectedUsage, collectedThoughtSignatures, artifactPromises, @@ -169,6 +171,10 @@ class AgentClient extends BaseClient { this.toolInputValidationErrors = toolInputValidationErrors; /** @type {MessageContentComplex[]} */ this.contentParts = contentParts; + /** Original run-step identity used by the content aggregator to attach + * completion events to their rendered content indices. + * @type {Map | undefined} */ + this.stepMap = stepMap; /** @type {Array} */ this.collectedUsage = collectedUsage; /** Vertex Gemini 3 thought signatures captured during the run, keyed by @@ -302,7 +308,7 @@ class AgentClient extends BaseClient { conversationId: this.conversationId, }, }, - { durable: true }, + { durable: true, expectedCreatedAt: this.jobCreatedAt }, ); } @@ -983,6 +989,7 @@ class AgentClient extends BaseClient { config, messageId, streamId, + jobCreatedAt: this.jobCreatedAt, conversationId, memoryMethods: { setMemory: db.setMemory, @@ -1333,10 +1340,14 @@ class AgentClient extends BaseClient { const emit = (async () => { try { if (streamId) { - await GenerationJobManager.emitChunk(streamId, { - event: UsageEvents.ON_TOKEN_USAGE, - data, - }); + await GenerationJobManager.emitChunk( + streamId, + { + event: UsageEvents.ON_TOKEN_USAGE, + data, + }, + { expectedCreatedAt: this.jobCreatedAt }, + ); } else { sendEvent(res, { event: UsageEvents.ON_TOKEN_USAGE, data }); } @@ -1533,10 +1544,14 @@ class AgentClient extends BaseClient { logger.error(`[AgentClient] Failed to release request slot on pause ${streamId}`, err); } } - await GenerationJobManager.emitChunk(streamId, { - event: ApprovalEvents.ON_PENDING_ACTION, - data: toClientPendingAction(pendingAction), - }); + await GenerationJobManager.emitChunk( + streamId, + { + event: ApprovalEvents.ON_PENDING_ACTION, + data: toClientPendingAction(pendingAction), + }, + { expectedCreatedAt: this.jobCreatedAt }, + ); // Steers queued before this pause stay IN the store for the whole approval // window: `resumeState.pendingSteers` re-seeds the client's chips on // reload, and the resumed run drains them at its first tool boundary. @@ -2067,12 +2082,14 @@ class AgentClient extends BaseClient { * @param {object} params * @param {Agents.ToolApprovalDecisionMap | { answer: string }} params.resumeValue * @param {Array} [params.seedContent] - content aggregated before the pause + * @param {Array} [params.runSteps] - run steps emitted before the pause * @param {AbortController} [params.abortController] * @param {Pick} [params.commandOptions] */ async resumeCompletion({ resumeValue, seedContent = [], + runSteps = [], abortController = null, commandOptions, userMCPAuthMap, @@ -2203,6 +2220,8 @@ class AgentClient extends BaseClient { throw new Error('Failed to create run for resume'); } + hydrateResumeRunSteps(runSteps, this.stepMap, run.Graph, seedContent); + this.run = run; if (this._resolveRun) { this._resolveRun(run); diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index f07ec03e9a..954c45b747 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -2922,6 +2922,29 @@ describe('AgentClient - titleConvo', () => { ); }); + it('should bind memory processing to the current generation epoch', async () => { + mockReq._resumableStreamId = 'convo-123'; + mockCheckAccess.mockResolvedValue(true); + mockInitializeAgent.mockResolvedValue({ + ...mockAgent, + provider: EModelEndpoint.openAI, + }); + mockCreateMemoryProcessor.mockResolvedValue([undefined, jest.fn()]); + + client = new AgentClient({ ...mockOptions, jobCreatedAt: 1234 }); + client.conversationId = 'convo-123'; + client.responseMessageId = 'response-123'; + + await client.useMemory(); + + expect(mockCreateMemoryProcessor).toHaveBeenCalledWith( + expect.objectContaining({ + streamId: 'convo-123', + jobCreatedAt: 1234, + }), + ); + }); + it('should load different agent when memory config agent.id differs from current agent id', async () => { const differentAgentId = 'different-agent-456'; const differentAgent = { diff --git a/api/server/controllers/agents/request.js b/api/server/controllers/agents/request.js index a77cdedf66..172a895f4a 100644 --- a/api/server/controllers/agents/request.js +++ b/api/server/controllers/agents/request.js @@ -592,13 +592,17 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit if (titleAbortController.signal.aborted) { return; } - await GenerationJobManager.emitChunk(streamId, { - event: 'title', - data: { - conversationId: titleConversationId, - title, + await GenerationJobManager.emitChunk( + streamId, + { + event: 'title', + data: { + conversationId: titleConversationId, + title, + }, }, - }); + { expectedCreatedAt: jobCreatedAt }, + ); })().catch((err) => { logger.error('[ResumableAgentController] Error emitting title event', err); }); @@ -643,27 +647,31 @@ const ResumableAgentController = async (req, res, next, initializeClient, addTit logger.error('[ResumableAgentController] Failed to persist start metadata', err); }); - GenerationJobManager.emitChunk(streamId, { - created: true, - // Skill selections aren't on `userMessage` yet at onStart (BaseClient adds - // them later), so attach them from the request — this is the message - // `trackUserMessage` persists as the authoritative job.metadata.userMessage, - // and it's what the live client renders the user bubble from. - message: { - ...userMessage, - // Carry files so trackUserMessage (the authoritative writer) persists them on - // job.metadata.userMessage for a HITL resume (see the updateMetadata above). - ...(Array.isArray(req.body?.files) && - req.body.files.length > 0 && { files: req.body.files }), - ...(Array.isArray(req.body?.manualSkills) && - req.body.manualSkills.length > 0 && { manualSkills: req.body.manualSkills }), - ...(Array.isArray(req.body?.alwaysAppliedSkills) && - req.body.alwaysAppliedSkills.length > 0 && { - alwaysAppliedSkills: req.body.alwaysAppliedSkills, - }), - }, + GenerationJobManager.emitChunk( streamId, - }).catch((err) => { + { + created: true, + // Skill selections aren't on `userMessage` yet at onStart (BaseClient adds + // them later), so attach them from the request — this is the message + // `trackUserMessage` persists as the authoritative job.metadata.userMessage, + // and it's what the live client renders the user bubble from. + message: { + ...userMessage, + // Carry files so trackUserMessage (the authoritative writer) persists them on + // job.metadata.userMessage for a HITL resume (see the updateMetadata above). + ...(Array.isArray(req.body?.files) && + req.body.files.length > 0 && { files: req.body.files }), + ...(Array.isArray(req.body?.manualSkills) && + req.body.manualSkills.length > 0 && { manualSkills: req.body.manualSkills }), + ...(Array.isArray(req.body?.alwaysAppliedSkills) && + req.body.alwaysAppliedSkills.length > 0 && { + alwaysAppliedSkills: req.body.alwaysAppliedSkills, + }), + }, + streamId, + }, + { expectedCreatedAt: jobCreatedAt }, + ).catch((err) => { logger.error('[ResumableAgentController] Failed to queue created event', err); }); }; diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 124a43ad55..a33a47be8d 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -327,10 +327,14 @@ async function finalizeResumedTurn({ client, onTitleGenerated: ({ conversationId: titleConvoId, title }) => { conversation.title = title; - return GenerationJobManager.emitChunk(streamId, { - event: 'title', - data: { conversationId: titleConvoId, title }, - }); + return GenerationJobManager.emitChunk( + streamId, + { + event: 'title', + data: { conversationId: titleConvoId, title }, + }, + { expectedCreatedAt: job.createdAt }, + ); }, }); } catch (err) { @@ -695,6 +699,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) await client.resumeCompletion({ resumeValue: mapped.resumeValue, seedContent, + runSteps: resumeState?.runSteps ?? [], abortController: job.abortController, // Carry the user's MCP auth so approved MCP tools run with their credentials. userMCPAuthMap: result.userMCPAuthMap, diff --git a/api/server/services/ActionService.js b/api/server/services/ActionService.js index db8a536d8e..c39533fa7e 100644 --- a/api/server/services/ActionService.js +++ b/api/server/services/ActionService.js @@ -175,6 +175,7 @@ async function loadActionSets(searchParams) { * @param {import('zod').ZodTypeAny | undefined} [params.zodSchema] - The Zod schema for tool input validation/definition * @param {{ oauth_client_id?: string; oauth_client_secret?: string; }} params.encrypted - The encrypted values for the action. * @param {string | null} [params.streamId] - The stream ID for resumable streams. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. * @param {boolean} [params.useSSRFProtection] - When true, uses SSRF-safe HTTP agents that validate resolved IPs at connect time. * @param {string[] | null} [params.allowedAddresses] - Optional admin exemption list of host:port pairs that bypass the SSRF private-IP block. * @returns { Promise unknown}> } An object with `_call` method to execute the tool input. @@ -189,6 +190,7 @@ async function createActionTool({ description, encrypted, streamId = null, + jobCreatedAt, useSSRFProtection = false, allowedAddresses, }) { @@ -250,7 +252,9 @@ async function createActionTool({ async () => { const eventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data }; if (streamId) { - await GenerationJobManager.emitChunk(streamId, eventData); + await GenerationJobManager.emitChunk(streamId, eventData, { + expectedCreatedAt: jobCreatedAt, + }); } else { sendEvent(res, eventData); } @@ -281,7 +285,9 @@ async function createActionTool({ data.delta.expires_at = undefined; const successEventData = { event: GraphEvents.ON_RUN_STEP_DELTA, data }; if (streamId) { - await GenerationJobManager.emitChunk(streamId, successEventData); + await GenerationJobManager.emitChunk(streamId, successEventData, { + expectedCreatedAt: jobCreatedAt, + }); } else { sendEvent(res, successEventData); } diff --git a/api/server/services/ActionService.spec.js b/api/server/services/ActionService.spec.js index 52419975f7..73c4cf6a5a 100644 --- a/api/server/services/ActionService.spec.js +++ b/api/server/services/ActionService.spec.js @@ -1,14 +1,51 @@ const { Constants, actionDelimiter, actionDomainSeparator } = require('librechat-data-provider'); -const { domainParser, legacyDomainEncode, validateAndUpdateTool } = require('./ActionService'); + +const mockEmitChunk = jest.fn(); +const mockFindToken = jest.fn(); +const mockActionFlowManager = { + createFlowWithHandler: jest.fn(), + createFlow: jest.fn(), +}; jest.mock('keyv'); +jest.mock('jsonwebtoken', () => ({ + sign: jest.fn(() => 'signed-state'), +})); + +jest.mock('@librechat/agents', () => ({ + ...jest.requireActual('@librechat/agents'), + sleep: jest.fn().mockResolvedValue(undefined), +})); + +jest.mock('@librechat/api', () => ({ + ...jest.requireActual('@librechat/api'), + validateActionOAuthMetadata: jest.fn().mockResolvedValue(undefined), + GenerationJobManager: { + emitChunk: (...args) => mockEmitChunk(...args), + }, +})); + jest.mock('~/models', () => ({ getActions: jest.fn(), + findToken: (...args) => mockFindToken(...args), + updateToken: jest.fn(), + createToken: jest.fn(), deleteActions: jest.fn(), + deleteAssistant: jest.fn(), +})); + +jest.mock('~/config', () => ({ + getActionFlowStateManager: jest.fn(() => mockActionFlowManager), })); const { getActions } = require('~/models'); +const { + createActionTool, + domainParser, + legacyDomainEncode, + validateAndUpdateTool, +} = require('./ActionService'); let mockDomainCache = {}; jest.mock('~/cache/getLogStores', () => { @@ -24,6 +61,10 @@ jest.mock('~/cache/getLogStores', () => { beforeEach(() => { mockDomainCache = {}; getActions.mockReset(); + mockEmitChunk.mockReset(); + mockFindToken.mockReset(); + mockActionFlowManager.createFlowWithHandler.mockReset(); + mockActionFlowManager.createFlow.mockReset(); }); const SEP = actionDomainSeparator; @@ -202,6 +243,78 @@ describe('legacyDomainEncode', () => { }); }); +describe('createActionTool OAuth events', () => { + it('fences resumable login and completion deltas to the owning job epoch', async () => { + const streamId = 'action-oauth-stream'; + const jobCreatedAt = 1234; + const preparedExecutor = { + setAuth: jest.fn().mockResolvedValue(undefined), + execute: jest.fn().mockResolvedValue({ data: { ok: true } }), + }; + const requestBuilder = { + createExecutor: jest.fn(() => ({ + setParams: jest.fn(() => preparedExecutor), + })), + }; + mockFindToken.mockResolvedValue(null); + mockActionFlowManager.createFlowWithHandler.mockImplementation( + async (_flowId, _type, handler) => handler(), + ); + mockActionFlowManager.createFlow.mockResolvedValue({ + access_token: 'access-token', + refresh_token: 'refresh-token', + expires_in: 3600, + }); + + const actionTool = await createActionTool({ + userId: 'action-user', + res: {}, + action: { + action_id: 'action-1', + metadata: { + domain: 'https://api.example.com', + oauth_client_id: 'client-id', + auth: { + type: 'oauth', + authorization_url: 'https://auth.example.com/authorize', + client_url: 'https://auth.example.com/token', + scope: 'read', + }, + }, + }, + requestBuilder, + encrypted: { + oauth_client_id: 'encrypted-client-id', + oauth_client_secret: 'encrypted-client-secret', + }, + streamId, + jobCreatedAt, + }); + + await actionTool._call( + {}, + { + metadata: { + thread_id: 'thread-1', + run_id: 'run-1', + }, + toolCall: { + id: 'tool-call-1', + stepId: 'step-1', + name: 'action-tool', + type: 'tool_call', + }, + }, + ); + + expect(mockEmitChunk).toHaveBeenCalledTimes(2); + for (const [emittedStreamId, , options] of mockEmitChunk.mock.calls) { + expect(emittedStreamId).toBe(streamId); + expect(options).toEqual({ expectedCreatedAt: jobCreatedAt }); + } + }); +}); + describe('validateAndUpdateTool', () => { const mockReq = { user: { id: 'user123' } }; diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 9cc8cba7d1..cdf86f85ba 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -59,8 +59,9 @@ const db = require('~/models'); * @param {string | null} [streamId] - The stream ID for resumable mode * @param {boolean} [definitionsOnly=false] - When true, returns only serializable * tool definitions without creating full tool instances (for event-driven mode) + * @param {number} [jobCreatedAt] - The generation epoch that owns emitted tool events */ -function createToolLoader(signal, streamId = null, definitionsOnly = false) { +function createToolLoader(signal, streamId = null, definitionsOnly = false, jobCreatedAt) { /** * @param {object} params * @param {ServerRequest} params.req @@ -96,6 +97,7 @@ function createToolLoader(signal, streamId = null, definitionsOnly = false) { agent, signal, streamId, + jobCreatedAt, tool_resources, definitionsOnly, }); @@ -143,7 +145,13 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt /** @type {Map} */ const toolInputValidationErrors = new Map(); const { contentParts, aggregateContent, stepMap } = createContentAggregator(); - const toolEndCallback = createToolEndCallback({ req, res, artifactPromises, streamId }); + const toolEndCallback = createToolEndCallback({ + req, + res, + artifactPromises, + streamId, + jobCreatedAt, + }); /** Query accessible skill IDs once per run (shared across all agents). * Skills activate under strict opt-in semantics — see @@ -269,6 +277,7 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt userMCPAuthMap: ctx.userMCPAuthMap, tool_resources: ctx.tool_resources, actionsEnabled: ctx.actionsEnabled, + jobCreatedAt, }); logger.debug(`[ON_TOOL_EXECUTE] loaded ${result.loadedTools?.length ?? 0} tools`); @@ -288,7 +297,7 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt req, updateToolCallResult: db.updateToolCallResult, }), - emitAttachment: createAttachmentEmitter({ res, streamId }), + emitAttachment: createAttachmentEmitter({ res, streamId, jobCreatedAt }), ...getSkillToolDeps(), }; @@ -337,6 +346,7 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt collectedUsage, collectedThoughtSignatures, streamId, + jobCreatedAt, subagentAggregatorsByToolCallId, usageCost, contextUsageSink, @@ -364,7 +374,7 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt const allowedProviders = new Set(appConfig?.endpoints?.[EModelEndpoint.agents]?.allowedProviders); /** Event-driven mode: only load tool definitions, not full instances */ - const loadTools = createToolLoader(signal, streamId, true); + const loadTools = createToolLoader(signal, streamId, true, jobCreatedAt); /** @type {Array} */ const requestFiles = req.body.files ?? []; /** @type {string} */ @@ -1009,6 +1019,7 @@ const initializeClient = async ({ req, res, signal, endpointOption, jobCreatedAt res, sender, contentParts, + stepMap, agentConfigs, eventHandlers, collectedUsage, diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index dd3a4a86ca..871a7ecc3d 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -42,11 +42,13 @@ jest.mock('@librechat/api', () => ({ * `ON_TOOL_EXECUTE` pipeline with a real subagent id and observe whether * the tool context (agent, tool_resources, skill ACLs) was preserved. */ let capturedToolExecuteOptions; +let capturedDefaultHandlerOptions; jest.mock('~/server/controllers/agents/callbacks', () => ({ createToolEndCallback: jest.fn(() => jest.fn()), createAttachmentEmitter: jest.fn(() => jest.fn()), createBackgroundCodeResultHandler: jest.fn(() => jest.fn()), getDefaultHandlers: jest.fn((opts) => { + capturedDefaultHandlerOptions = opts; capturedToolExecuteOptions = opts?.toolExecuteOptions; return {}; }), @@ -109,6 +111,7 @@ describe('initializeClient — processAgent ACL gate', () => { await mongoose.connection.dropDatabase(); jest.clearAllMocks(); agentClientArgs = undefined; + capturedDefaultHandlerOptions = undefined; testUser = await User.create({ email: 'test@example.com', @@ -151,6 +154,51 @@ describe('initializeClient — processAgent ACL gate', () => { maxContextTokens: 4096, }); + it('threads the owning job epoch into resumable event handlers', async () => { + const { + createAttachmentEmitter, + createToolEndCallback, + } = require('~/server/controllers/agents/callbacks'); + mockInitializeAgent.mockResolvedValue(makePrimaryConfig([])); + const req = makeReq(); + req._resumableStreamId = 'conv_1'; + + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + jobCreatedAt: 1234, + }); + + expect(capturedDefaultHandlerOptions).toEqual( + expect.objectContaining({ + streamId: 'conv_1', + jobCreatedAt: 1234, + }), + ); + expect(createToolEndCallback).toHaveBeenCalledWith( + expect.objectContaining({ + streamId: 'conv_1', + jobCreatedAt: 1234, + }), + ); + expect(createAttachmentEmitter).toHaveBeenCalledWith({ + res: {}, + streamId: 'conv_1', + jobCreatedAt: 1234, + }); + + mockLoadToolsForExecution.mockResolvedValue({ loadedTools: [], configurable: {} }); + await capturedToolExecuteOptions.loadTools([], PRIMARY_ID); + expect(mockLoadToolsForExecution).toHaveBeenCalledWith( + expect.objectContaining({ + streamId: 'conv_1', + jobCreatedAt: 1234, + }), + ); + }); + it('should skip handoff agent and filter its edge when user lacks VIEW access', async () => { await createAgent({ id: TARGET_ID, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index f6fc35f1ed..9ac936f57c 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -244,8 +244,9 @@ function isEmptyObjectSchema(jsonSchema) { * @param {string} params.stepId - The ID of the step in the flow. * @param {ToolCallChunk} params.toolCall - The tool call object containing tool information. * @param {string | null} [params.streamId] - The stream ID for resumable mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. */ -function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null }) { +function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null, jobCreatedAt }) { /** * @param {string} authURL - The URL to redirect the user for OAuth authentication. * @param {{ expiresAt?: number }} [options] @@ -254,7 +255,9 @@ function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null }) { return async function (authURL, options) { const eventData = buildMCPAuthRunStepDeltaEvent({ authURL, stepId, toolCall, options }); if (streamId) { - await GenerationJobManager.emitChunk(streamId, eventData); + await GenerationJobManager.emitChunk(streamId, eventData, { + expectedCreatedAt: jobCreatedAt, + }); } else { sendEvent(res, eventData); } @@ -269,13 +272,24 @@ function createRunStepDeltaEmitter({ res, stepId, toolCall, streamId = null }) { * @param {ToolCallChunk} params.toolCall - The tool call object containing tool information. * @param {number} [params.index] * @param {string | null} [params.streamId] - The stream ID for resumable mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. * @returns {() => Promise} */ -function createRunStepEmitter({ res, runId, stepId, toolCall, index, streamId = null }) { +function createRunStepEmitter({ + res, + runId, + stepId, + toolCall, + index, + streamId = null, + jobCreatedAt, +}) { return async function () { const eventData = buildMCPAuthRunStepEvent({ runId, stepId, toolCall, index }); if (streamId) { - await GenerationJobManager.emitChunk(streamId, eventData); + await GenerationJobManager.emitChunk(streamId, eventData, { + expectedCreatedAt: jobCreatedAt, + }); } else { sendEvent(res, eventData); } @@ -333,12 +347,15 @@ function createOAuthStart({ flowId, flowManager, callback }) { * @param {string} params.stepId - The ID of the step in the flow. * @param {ToolCallChunk} params.toolCall - The tool call object containing tool information. * @param {string | null} [params.streamId] - The stream ID for resumable mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. */ -function createOAuthEnd({ res, stepId, toolCall, streamId = null }) { +function createOAuthEnd({ res, stepId, toolCall, streamId = null, jobCreatedAt }) { return async function () { const eventData = buildMCPAuthRunStepEndDeltaEvent({ stepId, toolCall }); if (streamId) { - await GenerationJobManager.emitChunk(streamId, eventData); + await GenerationJobManager.emitChunk(streamId, eventData, { + expectedCreatedAt: jobCreatedAt, + }); } else { sendEvent(res, eventData); } @@ -386,6 +403,7 @@ function createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }) { * @param {string} params.model * @param {number} [params.index] * @param {string | null} [params.streamId] - The stream ID for resumable mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. * @param {Record>} [params.userMCPAuthMap] * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] * @param {import('@librechat/api').ParsedServerConfig} [params.serverConfig] - Used to bypass reconnect throttling for request-scoped servers. @@ -403,6 +421,7 @@ async function reconnectServer({ requestBody, requestScopedConnections, streamId = null, + jobCreatedAt, }) { logger.debug( `[MCP][reconnectServer] serverName: ${serverName}, user: ${user?.id}, hasUserMCPAuthMap: ${!!userMCPAuthMap}`, @@ -456,12 +475,14 @@ async function reconnectServer({ stepId, toolCall, streamId, + jobCreatedAt, }); const runStepDeltaEmitter = createRunStepDeltaEmitter({ res, stepId, toolCall, streamId, + jobCreatedAt, }); const callback = createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }); const oauthStart = createOAuthStart({ @@ -508,6 +529,7 @@ async function reconnectServer({ * @param {number} [params.index] * @param {AbortSignal} [params.signal] * @param {string | null} [params.streamId] - The stream ID for resumable mode. + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. * @param {import('@librechat/api').ParsedServerConfig} [params.config] * @param {import('@librechat/api').RequestBody} [params.requestBody] * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] @@ -528,6 +550,7 @@ async function createMCPTools({ requestBody, requestScopedConnections, streamId = null, + jobCreatedAt, }) { const serverConfig = config ?? (await getMCPServersRegistry().getServerConfig(serverName, user?.id, configServers)); @@ -567,6 +590,7 @@ async function createMCPTools({ requestBody, requestScopedConnections, streamId, + jobCreatedAt, }); if (result === null) { logger.debug(`[MCP][${serverName}] Reconnect throttled, skipping tool creation.`); @@ -587,6 +611,7 @@ async function createMCPTools({ userMCPAuthMap, configServers, streamId, + jobCreatedAt, availableTools: result.availableTools, toolKey: `${tool.name}${Constants.mcp_delimiter}${serverName}`, requestBody, @@ -619,6 +644,7 @@ async function createMCPTools({ * @param {Record>} [params.userMCPAuthMap] * @param {import('@librechat/api').ParsedServerConfig} [params.config] * @param {(availableTools: LCAvailableTools) => void} [params.onAvailableTools] + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted events. * @returns { Promise unknown}> } An object with `_call` method to execute the tool input. */ async function createMCPTool({ @@ -637,6 +663,7 @@ async function createMCPTool({ configServers, onAvailableTools, streamId = null, + jobCreatedAt, }) { const [toolName, serverName] = toolKey.split(Constants.mcp_delimiter); @@ -694,6 +721,7 @@ async function createMCPTool({ requestBody, requestScopedConnections, streamId, + jobCreatedAt, }); if (result?.availableTools) { onAvailableTools?.(result.availableTools); @@ -725,6 +753,7 @@ async function createMCPTool({ serverConfig, toolDefinition, streamId, + jobCreatedAt, }); } @@ -740,6 +769,7 @@ function createToolInstance({ toolDefinition, provider: capturedProvider, streamId = null, + jobCreatedAt, }) { /** @type {LCTool} */ const { description, parameters } = toolDefinition; @@ -795,6 +825,7 @@ function createToolInstance({ stepId, toolCall, streamId, + jobCreatedAt, }); const oauthStart = createOAuthStart({ flowId, @@ -806,6 +837,7 @@ function createToolInstance({ stepId, toolCall, streamId, + jobCreatedAt, }); if (derivedSignal) { diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index 30fbc6442b..1e529189c6 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -39,7 +39,7 @@ jest.mock('@librechat/api', () => { }); const { logger } = require('@librechat/data-schemas'); -const { MCPOAuthHandler } = require('@librechat/api'); +const { MCPOAuthHandler, GenerationJobManager } = require('@librechat/api'); const { CacheKeys, Constants, Permissions, PermissionTypes } = require('librechat-data-provider'); const D = Constants.mcp_delimiter; const { @@ -837,6 +837,44 @@ describe('User parameter passing tests', () => { expect(mockReinitMCPServer.mock.calls[0][0].user).toBe(mockUser); }); + it('fences resumable tool-loading OAuth events to the owning job epoch', async () => { + const mockUser = { id: 'epoch-loading-user', name: 'Epoch Loading User' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const streamId = 'epoch-loading-stream'; + const jobCreatedAt = 1234; + const flowManager = { + getFlowState: jest.fn().mockResolvedValue(null), + createFlowWithHandler: jest.fn(async (_flowId, _type, handler) => handler()), + failFlow: jest.fn(), + }; + mockGetFlowStateManager.mockReturnValue(flowManager); + mockReinitMCPServer.mockImplementation(async ({ oauthStart }) => { + await oauthStart('https://auth.example.com/loading'); + return { tools: [], availableTools: {} }; + }); + + await createMCPTools({ + res: mockRes, + user: mockUser, + serverName: 'epoch-loading-server', + provider: 'openai', + userMCPAuthMap: {}, + config: { type: 'stdio' }, + streamId, + jobCreatedAt, + }); + + expect(GenerationJobManager.emitChunk).toHaveBeenCalledTimes(2); + expect(GenerationJobManager.emitChunk.mock.calls.map(([, event]) => event.event)).toEqual([ + 'on_run_step', + 'on_run_step_delta', + ]); + for (const [emittedStreamId, , options] of GenerationJobManager.emitChunk.mock.calls) { + expect(emittedStreamId).toBe(streamId); + expect(options).toEqual({ expectedCreatedAt: jobCreatedAt }); + } + }); + it('should fail tenant-scoped OAuth flows when tool loading is aborted', async () => { const mockUser = { id: 'tenant-user', name: 'Tenant User' }; const mockRes = { write: jest.fn(), flush: jest.fn() }; @@ -1054,6 +1092,76 @@ describe('User parameter passing tests', () => { expect(mockGetMCPManager).not.toHaveBeenCalled(); }); + it('fences resumable tool-call OAuth events to the owning job epoch', async () => { + const mockUser = { id: 'epoch-tool-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const streamId = 'epoch-tool-stream'; + const jobCreatedAt = 5678; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + const flowManager = { + getFlowState: jest.fn().mockResolvedValue(null), + createFlowWithHandler: jest.fn(async (_flowId, _type, handler) => handler()), + failFlow: jest.fn(), + }; + mockGetFlowStateManager.mockReturnValue(flowManager); + mockGetMCPManager.mockReturnValue({ + callTool: jest.fn(async ({ oauthStart, oauthEnd }) => { + await oauthStart('https://auth.example.com/tool-call'); + await oauthEnd(); + return ['ok', null]; + }), + }); + + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + toolKey: `test-tool${D}epoch-tool-server`, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`test-tool${D}epoch-tool-server`]: { + function: { + description: 'Epoch-fenced tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + streamId, + jobCreatedAt, + }); + + await mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { + provider: 'openai', + thread_id: 'thread-epoch', + run_id: 'run-epoch', + }, + toolCall: { + id: 'tool-call-epoch', + stepId: 'step-epoch', + name: 'test-tool', + type: 'tool_call', + }, + }, + ); + + expect(GenerationJobManager.emitChunk).toHaveBeenCalledTimes(2); + for (const [emittedStreamId, , options] of GenerationJobManager.emitChunk.mock.calls) { + expect(emittedStreamId).toBe(streamId); + expect(options).toEqual({ expectedCreatedAt: jobCreatedAt }); + } + }); + it('should reuse request-scoped MCP permission checks across tool executions', async () => { const mockUser = { id: 'mcp-allowed-user', role: 'USER' }; const mockReq = { user: mockUser }; diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index be67b60fa7..ec28d15b55 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -532,6 +532,7 @@ const isBuiltInTool = (toolName) => * @param {ServerResponse} [params.res] - The response object for SSE events * @param {Object} params.agent - The agent configuration * @param {string|null} [params.streamId] - Stream ID for resumable mode + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events * @returns {Promise<{ * toolDefinitions?: import('@librechat/api').LCTool[]; * toolRegistry?: Map; @@ -540,7 +541,14 @@ const isBuiltInTool = (toolName) => * hasDeferredTools?: boolean; * }>} */ -async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, tool_resources }) { +async function loadToolDefinitionsWrapper({ + req, + res, + agent, + streamId = null, + jobCreatedAt, + tool_resources, +}) { if (!agent.tools || agent.tools.length === 0) { return { toolDefinitions: [] }; } @@ -653,8 +661,12 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to }); if (streamId) { - await GenerationJobManager.emitChunk(streamId, runStepEvent); - await GenerationJobManager.emitChunk(streamId, runStepDeltaEvent); + await GenerationJobManager.emitChunk(streamId, runStepEvent, { + expectedCreatedAt: jobCreatedAt, + }); + await GenerationJobManager.emitChunk(streamId, runStepDeltaEvent, { + expectedCreatedAt: jobCreatedAt, + }); } else if (res && !res.writableEnded) { sendEvent(res, runStepEvent); sendEvent(res, runStepDeltaEvent); @@ -683,7 +695,9 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to }); if (streamId) { - await GenerationJobManager.emitChunk(streamId, runStepCompletedEvent); + await GenerationJobManager.emitChunk(streamId, runStepCompletedEvent, { + expectedCreatedAt: jobCreatedAt, + }); } else if (res && !res.writableEnded) { sendEvent(res, runStepCompletedEvent); } else { @@ -1074,6 +1088,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to * @param {Object} [params.tool_resources] - Tool resources * @param {string} [params.openAIApiKey] - OpenAI API key * @param {string|null} [params.streamId] - Stream ID for resumable mode + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events * @param {boolean} [params.definitionsOnly=true] - When true, returns only serializable * tool definitions without creating full tool instances. Use for event-driven mode * where tools are loaded on-demand during execution. @@ -1086,10 +1101,18 @@ async function loadAgentTools({ tool_resources, openAIApiKey, streamId = null, + jobCreatedAt, definitionsOnly = true, }) { if (definitionsOnly) { - return loadToolDefinitionsWrapper({ req, res, agent, streamId, tool_resources }); + return loadToolDefinitionsWrapper({ + req, + res, + agent, + streamId, + jobCreatedAt, + tool_resources, + }); } if (!agent.tools || agent.tools.length === 0) { @@ -1155,7 +1178,7 @@ async function loadAgentTools({ /** @type {ReturnType} */ let webSearchCallbacks; if (includesWebSearch) { - webSearchCallbacks = createOnSearchResults(res, streamId); + webSearchCallbacks = createOnSearchResults(res, streamId, jobCreatedAt); } /** @type {Record>} */ @@ -1178,6 +1201,7 @@ async function loadAgentTools({ options: { req, res, + jobCreatedAt, openAIApiKey, tool_resources, processFileURL, @@ -1387,6 +1411,7 @@ async function loadAgentTools({ name: toolName, description: functionSignature.description, streamId, + jobCreatedAt, useSSRFProtection: !Array.isArray(_allowedDomains) || _allowedDomains.length === 0, allowedAddresses: _allowedAddresses, }); @@ -1440,6 +1465,7 @@ async function loadAgentTools({ * @param {Record>} [params.userMCPAuthMap] - User MCP auth map * @param {Object} [params.tool_resources] - Tool resources * @param {string|null} [params.streamId] - Stream ID for web search callbacks + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events * @param {boolean} [params.actionsEnabled] - Whether the actions capability is enabled * @returns {Promise<{ loadedTools: Array, configurable: Object }>} */ @@ -1456,6 +1482,7 @@ async function loadToolsForExecution({ userMCPAuthMap, tool_resources, streamId = null, + jobCreatedAt, actionsEnabled, }) { const appConfig = req.config; @@ -1611,7 +1638,9 @@ async function loadToolsForExecution({ if (regularToolNames.length > 0) { const includesWebSearch = regularToolNames.includes(Tools.web_search); - const webSearchCallbacks = includesWebSearch ? createOnSearchResults(res, streamId) : undefined; + const webSearchCallbacks = includesWebSearch + ? createOnSearchResults(res, streamId, jobCreatedAt) + : undefined; const { loadedTools } = await loadTools({ agent, @@ -1623,6 +1652,7 @@ async function loadToolsForExecution({ options: { req, res, + jobCreatedAt, tool_resources, processFileURL, uploadImageBuffer, @@ -1648,6 +1678,7 @@ async function loadToolsForExecution({ agent, appConfig, streamId, + jobCreatedAt, actionToolNames, }); allLoadedTools.push(...actionTools); @@ -1686,6 +1717,7 @@ async function loadToolsForExecution({ * @param {Object} params.agent - The agent object * @param {Object} params.appConfig - App configuration * @param {string|null} params.streamId - Stream ID + * @param {number} [params.jobCreatedAt] - The generation epoch that owns emitted tool events * @param {string[]} params.actionToolNames - Action tool names to load * @returns {Promise} Loaded action tools */ @@ -1695,6 +1727,7 @@ async function loadActionToolsForExecution({ agent, appConfig, streamId, + jobCreatedAt, actionToolNames, }) { const loadedActionTools = []; @@ -1787,6 +1820,7 @@ async function loadActionToolsForExecution({ res, action, streamId, + jobCreatedAt, zodSchema, encrypted, requestBuilder, diff --git a/api/server/services/Tools/search.js b/api/server/services/Tools/search.js index c4cdfc752f..c6e031e9ee 100644 --- a/api/server/services/Tools/search.js +++ b/api/server/services/Tools/search.js @@ -8,10 +8,15 @@ const { GenerationJobManager } = require('@librechat/api'); * @param {import('http').ServerResponse} res - The server response object * @param {string | null} streamId - The stream ID for resumable mode, or null for standard mode * @param {Object} attachment - The attachment data + * @param {number} [jobCreatedAt] - The generation epoch that owns the attachment */ -function writeAttachment(res, streamId, attachment) { +function writeAttachment(res, streamId, attachment, jobCreatedAt) { if (streamId) { - GenerationJobManager.emitChunk(streamId, { event: 'attachment', data: attachment }); + GenerationJobManager.emitChunk( + streamId, + { event: 'attachment', data: attachment }, + { expectedCreatedAt: jobCreatedAt }, + ); } else { res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`); } @@ -21,9 +26,10 @@ function writeAttachment(res, streamId, attachment) { * Creates a function to handle search results and stream them as attachments * @param {import('http').ServerResponse} res - The HTTP server response object * @param {string | null} [streamId] - The stream ID for resumable mode, or null for standard mode + * @param {number} [jobCreatedAt] - The generation epoch that owns emitted attachments * @returns {{ onSearchResults: function(SearchResult, GraphRunnableConfig): void; onGetHighlights: function(string): void}} - Function that takes search results and returns or streams an attachment */ -function createOnSearchResults(res, streamId = null) { +function createOnSearchResults(res, streamId = null, jobCreatedAt) { const context = { sourceMap: new Map(), searchResultData: undefined, @@ -86,7 +92,7 @@ function createOnSearchResults(res, streamId = null) { if (!res.headersSent) { return attachment; } - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); } /** @@ -108,7 +114,7 @@ function createOnSearchResults(res, streamId = null) { } const attachment = buildAttachment(context); - writeAttachment(res, streamId, attachment); + writeAttachment(res, streamId, attachment, jobCreatedAt); } return { diff --git a/api/server/services/Tools/search.spec.js b/api/server/services/Tools/search.spec.js new file mode 100644 index 0000000000..8925d84062 --- /dev/null +++ b/api/server/services/Tools/search.spec.js @@ -0,0 +1,63 @@ +const mockEmitChunk = jest.fn(); + +jest.mock('nanoid', () => ({ + nanoid: jest.fn(() => 'search-attachment'), +})); + +jest.mock('@librechat/api', () => ({ + GenerationJobManager: { + emitChunk: (...args) => mockEmitChunk(...args), + }, +})); + +const { createOnSearchResults } = require('./search'); + +describe('createOnSearchResults', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('fences resumable search attachments to the owning generation', () => { + const callbacks = createOnSearchResults({ headersSent: true }, 'conversation-1', 1234); + const runnableConfig = { + metadata: { + user_id: 'user-1', + thread_id: 'conversation-1', + run_id: 'response-1', + }, + toolCall: { + id: 'tool-call-1', + name: 'web_search', + turn: 0, + }, + }; + + callbacks.onSearchResults( + { + success: true, + data: { + organic: [{ link: 'https://example.com' }], + topStories: [], + }, + }, + runnableConfig, + ); + callbacks.onGetHighlights('https://example.com'); + + expect(mockEmitChunk).toHaveBeenCalledTimes(2); + for (const call of mockEmitChunk.mock.calls) { + expect(call).toEqual([ + 'conversation-1', + { + event: 'attachment', + data: expect.objectContaining({ + messageId: 'response-1', + toolCallId: 'tool-call-1', + conversationId: 'conversation-1', + }), + }, + { expectedCreatedAt: 1234 }, + ]); + } + }); +}); diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index db8c2c8cb1..6e902de5cd 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -101,6 +101,7 @@ const { processRequiredActions, resolveAgentCapabilities, } = require('../ToolService'); +const { createOnSearchResults } = require('~/server/services/Tools/search'); const { reinitMCPServer } = require('~/server/services/Tools/mcp'); const { PENDING_STALE_MS } = require('@librechat/api'); @@ -425,6 +426,66 @@ describe('ToolService - Action Capability Gating', () => { ]); }); + it('fences resumable MCP OAuth definition events to the owning job epoch', async () => { + const req = createMockReq([AgentCapabilities.tools]); + const res = { writableEnded: false }; + const serverName = 'Epoch-Server'; + const streamId = 'stream-epoch'; + const jobCreatedAt = 1234; + const mcpTool = `search${Constants.mcp_delimiter}${serverName}`; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig([AgentCapabilities.tools])); + mockResolveConfigServers.mockResolvedValue({ + [serverName]: { + type: 'streamable-http', + url: `https://mcp.example.com/${serverName}`, + requiresOAuth: true, + }, + }); + mockLoadToolDefinitions + .mockImplementationOnce(async (_args, deps) => { + await deps.getOrFetchMCPServerTools(req.user.id, serverName); + return { + toolDefinitions: [], + toolRegistry: new Map(), + hasDeferredTools: false, + }; + }) + .mockResolvedValue({ + toolDefinitions: [mcpTool], + toolRegistry: new Map(), + hasDeferredTools: false, + }); + reinitMCPServer.mockImplementation(async ({ returnOnOAuth, oauthStart, oauthEnd }) => { + await oauthStart(`https://auth.example.com/${serverName}`); + if (returnOnOAuth === false) { + await oauthEnd(); + return { availableTools: { [mcpTool]: {} } }; + } + return { availableTools: null }; + }); + + await loadAgentTools({ + req, + res, + agent: { id: 'agent_123', tools: [mcpTool] }, + definitionsOnly: true, + streamId, + jobCreatedAt, + }); + + expect(mockSendEvent).not.toHaveBeenCalled(); + expect(mockEmitChunk).toHaveBeenCalledTimes(3); + expect(mockEmitChunk.mock.calls.map(([, event]) => event.event)).toEqual([ + 'on_run_step', + 'on_run_step_delta', + 'on_run_step_completed', + ]); + for (const [emittedStreamId, , options] of mockEmitChunk.mock.calls) { + expect(emittedStreamId).toBe(streamId); + expect(options).toEqual({ expectedCreatedAt: jobCreatedAt }); + } + }); + it('should not expose cached MCP tool definitions when the registry lookup fails', async () => { const serverName = 'private-server'; const mcpTool = `search${Constants.mcp_delimiter}${serverName}`; @@ -968,6 +1029,24 @@ describe('ToolService - Action Capability Gating', () => { const actionToolName = `get_weather${actionDelimiter}api_example_com`; const regularTool = 'calculator'; + it('threads the owning job epoch into web-search attachment callbacks', async () => { + const capabilities = [AgentCapabilities.tools, AgentCapabilities.web_search]; + const req = createMockReq(capabilities); + const res = {}; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + + await loadAgentTools({ + req, + res, + streamId: 'conversation-1', + jobCreatedAt: 1234, + agent: { id: 'agent_123', tools: [Tools.web_search] }, + definitionsOnly: false, + }); + + expect(createOnSearchResults).toHaveBeenCalledWith(res, 'conversation-1', 1234); + }); + it('should not load action sets when actions capability is disabled', async () => { const capabilities = [AgentCapabilities.tools, AgentCapabilities.web_search]; const req = createMockReq(capabilities); @@ -1003,6 +1082,25 @@ describe('ToolService - Action Capability Gating', () => { const actionToolName = `get_weather${actionDelimiter}api_example_com`; const regularTool = Tools.web_search; + it('threads the owning job epoch into web-search attachment callbacks', async () => { + const capabilities = [AgentCapabilities.tools, AgentCapabilities.web_search]; + const req = createMockReq(capabilities); + const res = {}; + mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); + + await loadToolsForExecution({ + req, + res, + streamId: 'conversation-1', + jobCreatedAt: 1234, + agent: { id: 'agent_123', tools: [Tools.web_search] }, + toolNames: [Tools.web_search], + actionsEnabled: false, + }); + + expect(createOnSearchResults).toHaveBeenCalledWith(res, 'conversation-1', 1234); + }); + it('does not load code execution tools that were not registered for the agent', async () => { const capabilities = [ AgentCapabilities.tools, diff --git a/api/typedefs.js b/api/typedefs.js index 6ec36b3573..92839f48d4 100644 --- a/api/typedefs.js +++ b/api/typedefs.js @@ -933,6 +933,7 @@ * signal?: AbortSignal, * memory?: ConversationSummaryBufferMemory, * tool_resources?: AgentToolResources, + * jobCreatedAt?: number, * web_search?: ReturnType, * }} LoadToolOptions * @memberof typedefs diff --git a/client/src/components/Chat/Messages/Content/ApprovalContext.tsx b/client/src/components/Chat/Messages/Content/ApprovalContext.tsx index 7b5bcf9477..8fa66614da 100644 --- a/client/src/components/Chat/Messages/Content/ApprovalContext.tsx +++ b/client/src/components/Chat/Messages/Content/ApprovalContext.tsx @@ -281,6 +281,9 @@ export function useResumeSubmit() { const approvalMutation = useSubmitToolApprovalMutation(); const askMutation = useSubmitAskAnswerMutation(); const { getDecisions, isReady, setStatus } = useApprovalContext(); + /** React state cannot lock a second click in the same browser task. Keep a + * synchronous action-id guard alongside the rendered submission status. */ + const submittingToolActionIdsRef = useRef(new Set()); /** Ask status lives in Recoil so it works from the composer (outside the * provider); tool-approval status stays on the context. */ const { setAskStatus } = useAskSubmitStatus(); @@ -311,12 +314,23 @@ export function useResumeSubmit() { if (!fields || decisions.length === 0 || !isReady(actionId)) { return; } + if (submittingToolActionIdsRef.current.has(actionId)) { + return; + } + submittingToolActionIdsRef.current.add(actionId); setStatus(actionId, 'submitting'); approvalMutation.mutate( { ...fields, actionId, decisions }, { onSuccess: () => setStatus(actionId, 'submitted'), - onError: (error) => setStatus(actionId, isExpiredError(error) ? 'expired' : 'error'), + onError: (error) => { + const expired = isExpiredError(error); + if (!expired) { + // Network/validation failures are retryable; a 409 is terminal. + submittingToolActionIdsRef.current.delete(actionId); + } + setStatus(actionId, expired ? 'expired' : 'error'); + }, }, ); }, diff --git a/client/src/components/Chat/Messages/Content/Part.tsx b/client/src/components/Chat/Messages/Content/Part.tsx index 301d35532e..ef286c44ee 100644 --- a/client/src/components/Chat/Messages/Content/Part.tsx +++ b/client/src/components/Chat/Messages/Content/Part.tsx @@ -327,6 +327,7 @@ const Part = memo(function Part({ diff --git a/client/src/components/Chat/Messages/Content/ToolApproval.tsx b/client/src/components/Chat/Messages/Content/ToolApproval.tsx index 1068eadd1e..31ee82f654 100644 --- a/client/src/components/Chat/Messages/Content/ToolApproval.tsx +++ b/client/src/components/Chat/Messages/Content/ToolApproval.tsx @@ -174,7 +174,11 @@ export default function ToolApproval({ } return ( -
+
{description != null && description.length > 0 && (

{description}

)} diff --git a/client/src/components/Chat/Messages/Content/ToolCall.tsx b/client/src/components/Chat/Messages/Content/ToolCall.tsx index 2efde33a20..4eb0a3d0f9 100644 --- a/client/src/components/Chat/Messages/Content/ToolCall.tsx +++ b/client/src/components/Chat/Messages/Content/ToolCall.tsx @@ -22,6 +22,7 @@ export default function ToolCall({ initialProgress = 0.1, isLast = false, isSubmitting, + toolCallId, name, args: _args = '', output, @@ -33,6 +34,7 @@ export default function ToolCall({ initialProgress: number; isLast?: boolean; isSubmitting: boolean; + toolCallId?: string; name: string; args: string | Record; output?: string | null; @@ -214,7 +216,11 @@ export default function ToolCall({ return getFinishedText(); })()} -
+
-
+
{hasInfo && (
diff --git a/client/src/components/Chat/Messages/Content/ToolCallGroup.tsx b/client/src/components/Chat/Messages/Content/ToolCallGroup.tsx index 955f9b2b47..1d86ac356d 100644 --- a/client/src/components/Chat/Messages/Content/ToolCallGroup.tsx +++ b/client/src/components/Chat/Messages/Content/ToolCallGroup.tsx @@ -24,6 +24,27 @@ interface ToolMeta { hasOutput: boolean; } +type ToolCallWithNestedContent = Agents.ToolCall & { + subagent_content?: TMessageContentParts[]; +}; + +function hasPendingApprovalInPart(part: TMessageContentParts): boolean { + if (part.type !== ContentTypes.TOOL_CALL) { + return false; + } + const toolCall = part[ContentTypes.TOOL_CALL] as ToolCallWithNestedContent | undefined; + if (!toolCall) { + return false; + } + if (toolCall.approval != null && (toolCall.output?.length ?? 0) === 0) { + return true; + } + return ( + Array.isArray(toolCall.subagent_content) && + toolCall.subagent_content.some(hasPendingApprovalInPart) + ); +} + function getToolMeta(part: TMessageContentParts): ToolMeta | null { if (part.type !== ContentTypes.TOOL_CALL) { return null; @@ -107,9 +128,14 @@ export default function ToolCallGroup({ const mcpIconMap = useMCPIconMap(); const rootRef = useRef(null); const cancelLayoutReconcileRef = useRef<(() => void) | null>(null); + const retainedForPendingApprovalRef = useRef(false); const count = parts.length; const toolMetadata = useMemo(() => parts.map((p) => getToolMeta(p.part)), [parts]); + const hasPendingApproval = useMemo( + () => parts.some(({ part }) => hasPendingApprovalInPart(part)), + [parts], + ); const allCompleted = useMemo( () => toolMetadata.every((m) => m?.hasOutput === true), [toolMetadata], @@ -226,12 +252,34 @@ export default function ToolCallGroup({ if (isExpanded) { return; } + if (hasPendingApproval) { + // Approval controls own unsent local form state. Keep unresolved cards + // mounted (the collapsed panel is inert/hidden) so collapsing a batch + // cannot erase decisions the reviewer already made. + retainedForPendingApprovalRef.current = true; + return; + } + retainedForPendingApprovalRef.current = false; setShouldRenderBody(false); notifyLayoutChange(); }, - [isExpanded, notifyLayoutChange], + [hasPendingApproval, isExpanded, notifyLayoutChange], ); + useEffect(() => { + if (isExpanded) { + retainedForPendingApprovalRef.current = false; + return; + } + if (!hasPendingApproval && retainedForPendingApprovalRef.current) { + // A completed collapse transition retained this body only to preserve + // approval form state. Release it once the last approval resolves. + retainedForPendingApprovalRef.current = false; + setShouldRenderBody(false); + notifyLayoutChange(); + } + }, [hasPendingApproval, isExpanded, notifyLayoutChange]); + /** Category-aware header verb: subagents and questions read as their own * category (with tense), everything else is the generic "Used N tools". */ const resolveGroupLabel = (): string => { @@ -310,7 +358,12 @@ export default function ToolCallGroup({ aria-hidden="true" /> -
+
{shouldRenderBody && (
diff --git a/client/src/components/Chat/Messages/Content/__tests__/ApprovalContext.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/ApprovalContext.test.tsx new file mode 100644 index 0000000000..2ce2ae93bd --- /dev/null +++ b/client/src/components/Chat/Messages/Content/__tests__/ApprovalContext.test.tsx @@ -0,0 +1,76 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { act, renderHook } from '@testing-library/react'; +import ApprovalProvider, { useApprovalContext, useResumeSubmit } from '../ApprovalContext'; +import { ChatContext } from '~/Providers/ChatContext'; + +const mockApprovalMutate = jest.fn(); +const mockAskMutate = jest.fn(); + +jest.mock('~/data-provider', () => ({ + useSubmitToolApprovalMutation: () => ({ mutate: mockApprovalMutate }), + useSubmitAskAnswerMutation: () => ({ mutate: mockAskMutate }), +})); + +jest.mock('~/store/agents', () => ({ + useGetEphemeralAgent: () => () => undefined, +})); + +const chatContextValue = { + conversation: { + conversationId: 'conversation-1', + endpoint: 'agents', + agent_id: 'agent-1', + }, +} as unknown as React.ContextType; + +function wrapper({ children }: { children: React.ReactNode }) { + return ( + + + {children} + + + ); +} + +describe('useResumeSubmit', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + test('synchronously deduplicates tool approval submissions and unlocks a retryable error', () => { + const { result } = renderHook( + () => ({ + approval: useApprovalContext(), + resume: useResumeSubmit(), + }), + { wrapper }, + ); + + act(() => { + result.current.approval.registerToolCall('action-1', 'call-1'); + result.current.approval.setDecision('action-1', 'call-1', { + tool_call_id: 'call-1', + decision: 'approve', + }); + }); + + act(() => { + result.current.resume.submitToolApproval('action-1'); + result.current.resume.submitToolApproval('action-1'); + }); + expect(mockApprovalMutate).toHaveBeenCalledTimes(1); + + const firstOptions = mockApprovalMutate.mock.calls[0][1] as { + onError: (error: unknown) => void; + }; + act(() => firstOptions.onError(new Error('temporary failure'))); + + act(() => { + result.current.resume.submitToolApproval('action-1'); + result.current.resume.submitToolApproval('action-1'); + }); + expect(mockApprovalMutate).toHaveBeenCalledTimes(2); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/__tests__/ToolCallGroup.test.tsx b/client/src/components/Chat/Messages/Content/__tests__/ToolCallGroup.test.tsx index 3abff264a0..10d61e4482 100644 --- a/client/src/components/Chat/Messages/Content/__tests__/ToolCallGroup.test.tsx +++ b/client/src/components/Chat/Messages/Content/__tests__/ToolCallGroup.test.tsx @@ -83,6 +83,36 @@ const makePart = ( }, }) as unknown as TMessageContentParts; +const makeApprovalPart = (id: string, output = ''): TMessageContentParts => + ({ + type: ContentTypes.TOOL_CALL, + [ContentTypes.TOOL_CALL]: { + id, + name: 'approval_probe', + args: {}, + output, + approval: { + actionId: 'action-1', + allowed_decisions: ['approve', 'reject'], + }, + }, + }) as unknown as TMessageContentParts; + +const makeSubagentPart = ( + id: string, + subagentContent: TMessageContentParts[], +): TMessageContentParts => + ({ + type: ContentTypes.TOOL_CALL, + [ContentTypes.TOOL_CALL]: { + id, + name: Constants.SUBAGENT, + args: {}, + output: '', + subagent_content: subagentContent, + }, + }) as unknown as TMessageContentParts; + const imageAttachment: TAttachment = { filename: 'foo.png', filepath: '/files/foo.png', @@ -211,6 +241,169 @@ describe('ToolCallGroup image hoisting', () => { expect(screen.queryByTestId('inner-0')).not.toBeInTheDocument(); }); + it('keeps unresolved approval bodies mounted while the group is collapsed', () => { + const approvalParts = [ + { part: makeApprovalPart('t1'), idx: 0 }, + { part: makeApprovalPart('t2'), idx: 1 }, + ]; + renderGroup({ + ...baseProps, + parts: approvalParts, + renderPart: (_p: TMessageContentParts, idx: number) => ( +
+ {'approval'} +
+ ), + }); + + const button = screen.getByRole('button', { name: 'Used 2 tools' }); + const collapsible = button.nextElementSibling as HTMLElement; + expect(screen.getByTestId('approval-0')).toBeInTheDocument(); + + fireEvent.click(button); + fireEvent.transitionEnd(collapsible); + + expect(button).toHaveAttribute('aria-expanded', 'false'); + expect(screen.getByTestId('approval-0')).toBeInTheDocument(); + expect(screen.getByTestId('approval-1')).toBeInTheDocument(); + }); + + it('keeps deeply nested unresolved approval bodies mounted while the group is collapsed', () => { + const nestedApprovalParts = [ + { + part: makeSubagentPart('parent', [ + makeSubagentPart('child', [makeApprovalPart('grandchild')]), + ]), + idx: 0, + }, + { part: makePart('sibling'), idx: 1 }, + ]; + renderGroup({ + ...baseProps, + parts: nestedApprovalParts, + renderPart: (_p: TMessageContentParts, idx: number) => ( +
+ {'nested'} +
+ ), + }); + + const button = screen.getByRole('button', { name: 'Used 2 tools' }); + const collapsible = button.nextElementSibling as HTMLElement; + fireEvent.click(button); + fireEvent.transitionEnd(collapsible); + + expect(button).toHaveAttribute('aria-expanded', 'false'); + expect(screen.getByTestId('nested-0')).toBeInTheDocument(); + expect(screen.getByTestId('nested-1')).toBeInTheDocument(); + }); + + it('does not retain a collapsed group for an already resolved nested approval', () => { + const nestedApprovalParts = [ + { + part: makeSubagentPart('parent', [ + makeSubagentPart('child', [makeApprovalPart('grandchild', 'done')]), + ]), + idx: 0, + }, + { part: makePart('sibling'), idx: 1 }, + ]; + renderGroup({ + ...baseProps, + parts: nestedApprovalParts, + renderPart: (_p: TMessageContentParts, idx: number) => ( +
+ {'nested'} +
+ ), + }); + + const button = screen.getByRole('button', { name: 'Used 2 tools' }); + const collapsible = button.nextElementSibling as HTMLElement; + fireEvent.click(button); + fireEvent.transitionEnd(collapsible); + + expect(screen.queryByTestId('resolved-nested-0')).not.toBeInTheDocument(); + expect(screen.queryByTestId('resolved-nested-1')).not.toBeInTheDocument(); + }); + + it('unmounts retained approval bodies after every approval in a collapsed group resolves', async () => { + const renderPart = (_p: TMessageContentParts, idx: number) => ( +
+ {'approval'} +
+ ); + const propsFor = ( + firstOutput = '', + secondOutput = '', + ): React.ComponentProps => ({ + ...baseProps, + parts: [ + { part: makeApprovalPart('t1', firstOutput), idx: 0 }, + { part: makeApprovalPart('t2', secondOutput), idx: 1 }, + ], + renderPart, + }); + const { rerender } = renderGroup(propsFor()); + + const button = screen.getByRole('button', { name: 'Used 2 tools' }); + const collapsible = button.nextElementSibling as HTMLElement; + fireEvent.click(button); + fireEvent.transitionEnd(collapsible); + expect(screen.getByTestId('retained-0')).toBeInTheDocument(); + + rerender( + + + , + ); + expect(screen.getByTestId('retained-0')).toBeInTheDocument(); + expect(screen.getByTestId('retained-1')).toBeInTheDocument(); + + rerender( + + + , + ); + await waitFor(() => { + expect(screen.queryByTestId('retained-0')).not.toBeInTheDocument(); + expect(screen.queryByTestId('retained-1')).not.toBeInTheDocument(); + }); + }); + + it('waits for an active collapse transition before unmounting resolved approval bodies', () => { + const renderPart = (_p: TMessageContentParts, idx: number) => ( +
+ {'approval'} +
+ ); + const propsFor = (output = ''): React.ComponentProps => ({ + ...baseProps, + parts: [ + { part: makeApprovalPart('t1', output), idx: 0 }, + { part: makeApprovalPart('t2', output), idx: 1 }, + ], + renderPart, + }); + const { rerender } = renderGroup(propsFor()); + + const button = screen.getByRole('button', { name: 'Used 2 tools' }); + const collapsible = button.nextElementSibling as HTMLElement; + fireEvent.click(button); + + rerender( + + + , + ); + expect(screen.getByTestId('transitioning-0')).toBeInTheDocument(); + expect(screen.getByTestId('transitioning-1')).toBeInTheDocument(); + + fireEvent.transitionEnd(collapsible); + expect(screen.queryByTestId('transitioning-0')).not.toBeInTheDocument(); + expect(screen.queryByTestId('transitioning-1')).not.toBeInTheDocument(); + }); + it('reconciles layout after the group collapses from an expanded state', async () => { renderGroup(baseProps); diff --git a/client/src/utils/approval.spec.ts b/client/src/utils/approval.spec.ts index 0f2de67d69..2a7d5cf9cc 100644 --- a/client/src/utils/approval.spec.ts +++ b/client/src/utils/approval.spec.ts @@ -71,6 +71,37 @@ describe('applyPendingAction — tool_approval', () => { }); }); + it('replaces displayed tool args with the matching action request arguments', () => { + const originalArgs = { query: 'original model args' }; + const rewrittenArgs = { query: 'rewritten by policy hook' }; + const message = msg({ content: [toolCallPart('tc1', { args: originalArgs })] }); + const action = toolApprovalAction({ + payload: { + type: 'tool_approval', + action_requests: [ + { + name: 'search', + arguments: rewrittenArgs, + tool_call_id: 'tc1', + description: 'Review rewritten search', + }, + ], + review_configs: [ + { + action_name: 'search', + tool_call_id: 'tc1', + allowed_decisions: ['approve', 'reject', 'edit', 'respond'], + }, + ], + }, + }); + + const result = applyPendingAction(message, action); + + expect(getToolCall(result.content?.[0] as TMessageContentParts)?.args).toEqual(rewrittenArgs); + expect(getToolCall(message.content?.[0] as TMessageContentParts)?.args).toEqual(originalArgs); + }); + it('leaves a completed tool call (with output) untouched and returns the same message reference', () => { const message = msg({ content: [toolCallPart('tc1', { output: 'already ran' })] }); const result = applyPendingAction(message, toolApprovalAction()); diff --git a/client/src/utils/approval.ts b/client/src/utils/approval.ts index eb99abda06..8a56ec44bc 100644 --- a/client/src/utils/approval.ts +++ b/client/src/utils/approval.ts @@ -95,6 +95,10 @@ function tagApprovalOnPart( const reviewConfig = reviewByToolCallId.get(toolCallId); nextToolCall = { ...nextToolCall, + // A PreToolUse hook may replace the model's original args before asking. + // The interrupt payload is authoritative so the reviewer sees, edits, and + // approves the same arguments the resumed tool will actually execute. + args: request.arguments, approval: { actionId, allowed_decisions: reviewConfig?.allowed_decisions ?? [], diff --git a/e2e/config/librechat.e2e.yaml b/e2e/config/librechat.e2e.yaml index 3a8e287f94..07062dce7d 100644 --- a/e2e/config/librechat.e2e.yaml +++ b/e2e/config/librechat.e2e.yaml @@ -60,6 +60,18 @@ endpoints: - chain - ocr - run_in_background + # Keep the shared mock profile non-interactive except for the dedicated + # approval probe. This exercises real HITL pause/resume without wedging the + # existing file-authoring, steering, background-tool, or MCP specs. + toolApproval: + enabled: true + mode: bypass + ask: + - approval_probe_mcp_e2e-memory + reason: E2E approval required before running {tool}. + hooks: + - module: e2e/setup/tool-approval-hook.js + matcher: ^approval_probe_mcp_e2e-memory$ custom: - name: 'Mock Provider A' apiKey: 'e2e-mock-key-a' diff --git a/e2e/setup/fake-mcp-server.js b/e2e/setup/fake-mcp-server.js index 182d46bd80..03c2602147 100644 --- a/e2e/setup/fake-mcp-server.js +++ b/e2e/setup/fake-mcp-server.js @@ -1,9 +1,19 @@ #!/usr/bin/env node +const fs = require('node:fs'); +const path = require('node:path'); const { McpServer } = require('@modelcontextprotocol/sdk/server/mcp.js'); const { StdioServerTransport } = require('@modelcontextprotocol/sdk/server/stdio.js'); const z = require('zod/v4'); +const APPROVAL_AUDIT_DIR = path.join('/tmp', 'librechat-e2e-approval-audit'); + +function recordApprovalInvocation(value) { + fs.mkdirSync(APPROVAL_AUDIT_DIR, { recursive: true }); + const filename = Buffer.from(value).toString('base64url'); + fs.appendFileSync(path.join(APPROVAL_AUDIT_DIR, filename), `${value}\n`); +} + const server = new McpServer({ name: 'e2e-memory', version: '1.0.0', @@ -66,6 +76,29 @@ server.registerTool( }, ); +server.registerTool( + 'approval_probe', + { + description: + 'Echoes reviewed input so LibreChat mock end-to-end tests can verify tool approval decisions.', + inputSchema: { + value: z.string(), + review: z.string().optional(), + }, + }, + async ({ value }) => { + recordApprovalInvocation(value); + return { + content: [ + { + type: 'text', + text: `E2E approval probe executed: ${value}`, + }, + ], + }; + }, +); + async function main() { await server.connect(new StdioServerTransport()); } diff --git a/e2e/setup/fake-model.js b/e2e/setup/fake-model.js index d432e687d0..e284caabc8 100644 --- a/e2e/setup/fake-model.js +++ b/e2e/setup/fake-model.js @@ -34,6 +34,10 @@ const FORCED_ERROR_MARKER = 'E2E_FORCED_ERROR:'; const MARKDOWN_REPLY_MARKER = 'E2E_MARKDOWN_REPLY'; const BACKGROUND_DISPATCH_MARKER = 'E2E_BACKGROUND_DISPATCH:'; const BACKGROUND_COLLECT_MARKER = 'E2E_BACKGROUND_COLLECT:'; +const TOOL_APPROVAL_MARKER = 'E2E_TOOL_APPROVAL:'; +const TOOL_APPROVAL_BATCH_MARKER = 'E2E_TOOL_APPROVAL_BATCH:'; +const TOOL_APPROVAL_RESTRICTED_MARKER = 'E2E_TOOL_APPROVAL_RESTRICTED:'; +const TOOL_APPROVAL_REWRITE_MARKER = 'E2E_TOOL_APPROVAL_REWRITE:'; const CREATE_FILE_AUTHORING_FINAL_TEXT = 'E2E file authoring complete'; const EDIT_FILE_AUTHORING_FINAL_TEXT = 'E2E file edit complete'; const SKILL_ASSERTION_FINAL_TEXT = 'E2E skill assertion passed'; @@ -56,6 +60,8 @@ const CREATE_SKILL_TOOL_CALL_ID = 'call_e2e_create_skill'; const EDIT_SKILL_TOOL_CALL_ID = 'call_e2e_edit_skill'; const BACKGROUND_TOOL_NAME = 'slow_echo_mcp_e2e-memory'; const CHECK_BACKGROUND_TASK_TOOL_NAME = 'check_background_task'; +const APPROVAL_TOOL_NAME = 'approval_probe_mcp_e2e-memory'; +const APPROVAL_TOOL_CALL_PREFIX = 'call_e2e_approval_'; const BACKGROUND_DISPATCH_TOOL_CALL_ID = 'call_e2e_background_dispatch'; const BACKGROUND_COLLECT_TOOL_CALL_ID = 'call_e2e_background_collect'; const MODEL_SPEC_ACCESSIBLE_SKILL = 'e2e-model-spec-allowed'; @@ -832,6 +838,92 @@ function findLastToolMessageText(messages, requiredToken) { return ''; } +function approvalToolResponses(label, toolNames, review) { + if (!toolNames.has(APPROVAL_TOOL_NAME)) { + return { + responses: [`E2E approval unavailable: ${APPROVAL_TOOL_NAME} was not advertised.`], + }; + } + return { + responses: ['', ''], + toolCalls: [ + { + id: `${APPROVAL_TOOL_CALL_PREFIX}${label}`, + name: APPROVAL_TOOL_NAME, + args: { + value: `original-${label}`, + ...(review ? { review } : {}), + }, + type: 'tool_call', + }, + ], + }; +} + +function batchApprovalToolResponses(label, toolNames) { + if (!toolNames.has(APPROVAL_TOOL_NAME)) { + return { + responses: [`E2E approval unavailable: ${APPROVAL_TOOL_NAME} was not advertised.`], + }; + } + return { + responses: ['', ''], + toolCalls: [ + { + id: `${APPROVAL_TOOL_CALL_PREFIX}${label}_first`, + name: APPROVAL_TOOL_NAME, + args: { value: `first-${label}` }, + type: 'tool_call', + }, + { + id: `${APPROVAL_TOOL_CALL_PREFIX}${label}_second`, + name: APPROVAL_TOOL_NAME, + args: { value: `second-${label}` }, + type: 'tool_call', + }, + ], + }; +} + +/** + * Resume rebuilds the fake model without the original prompt in `context.messages`. + * Detect the checkpoint-restored approval tool messages on every model instance + * so the continuation can report the real approve/reject/edit/respond outcome. + */ +function approvalOutcomeResponses(messages) { + let latestHumanIndex = -1; + for (let index = 0; index < (messages ?? []).length; index++) { + const type = messageType(messages[index]); + if (type === 'human' || type === 'user') { + latestHumanIndex = index; + } + } + + const outcomeMessages = (messages ?? []) + .slice(latestHumanIndex + 1) + .filter( + (message) => + messageType(message) === 'tool' && + typeof message?.tool_call_id === 'string' && + message.tool_call_id.startsWith(APPROVAL_TOOL_CALL_PREFIX), + ); + + const isBatch = outcomeMessages.some( + (message) => + message.tool_call_id.endsWith('_first') || message.tool_call_id.endsWith('_second'), + ); + if (isBatch && outcomeMessages.length < 2) { + return null; + } + + const outcomes = outcomeMessages.map((message) => getContentText(message.content)); + + if (outcomes.length === 0) { + return null; + } + return { responses: [`E2E approval outcomes: ${outcomes.join(' | ')}`] }; +} + /** * Turn 1 of the background e2e: emit the MCP tool call with the injected * `run_in_background: true` arg, then (second model invocation, after the @@ -925,6 +1017,26 @@ function backgroundCollectResponses(messages, toolNames) { } function resolveResponses({ graph, messages, text, toolNames }) { + const batchApprovalLabel = getMarkerValue(text, TOOL_APPROVAL_BATCH_MARKER); + if (batchApprovalLabel) { + return batchApprovalToolResponses(batchApprovalLabel, toolNames); + } + + const restrictedApprovalLabel = getMarkerValue(text, TOOL_APPROVAL_RESTRICTED_MARKER); + if (restrictedApprovalLabel) { + return approvalToolResponses(restrictedApprovalLabel, toolNames, 'restricted'); + } + + const rewrittenApprovalLabel = getMarkerValue(text, TOOL_APPROVAL_REWRITE_MARKER); + if (rewrittenApprovalLabel) { + return approvalToolResponses(rewrittenApprovalLabel, toolNames, 'rewrite'); + } + + const approvalLabel = getMarkerValue(text, TOOL_APPROVAL_MARKER); + if (approvalLabel) { + return approvalToolResponses(approvalLabel, toolNames); + } + const reply = replyResponses(text); if (reply) { return reply; @@ -1055,5 +1167,15 @@ module.exports = function fakeModelHook(run, context) { text, toolNames, }); - overrideModel({ graph, responses, sleep, toolCalls, thrownError, resolveOnStream }); + overrideModel({ + graph, + responses, + sleep, + toolCalls, + thrownError, + resolveOnStream: (streamMessages, streamOptions, runManager) => + approvalOutcomeResponses(streamMessages) ?? + resolveOnStream?.(streamMessages, streamOptions, runManager) ?? + null, + }); }; diff --git a/e2e/setup/tool-approval-hook.js b/e2e/setup/tool-approval-hook.js new file mode 100644 index 0000000000..4eebaa751e --- /dev/null +++ b/e2e/setup/tool-approval-hook.js @@ -0,0 +1,29 @@ +/** + * Dynamic approval-policy fixture for the mock Playwright suite. + * + * The `review` argument selects behavior that cannot be expressed by the static + * ask list: a restricted decision set, or an authoritative argument rewrite. + */ +module.exports = () => () => async (input) => { + if (input.toolInput.review === 'restricted') { + return { + decision: 'ask', + reason: 'E2E approval offers approve or reject only.', + allowedDecisions: ['approve', 'reject'], + }; + } + + if (input.toolInput.review === 'rewrite') { + const originalValue = + typeof input.toolInput.value === 'string' ? input.toolInput.value : 'original-missing'; + return { + decision: 'ask', + reason: 'E2E approval reviews rewritten arguments.', + updatedInput: { + value: originalValue.replace(/^original-/, 'rewritten-'), + }, + }; + } + + return {}; +}; diff --git a/e2e/specs/mock/tool-approvals.spec.ts b/e2e/specs/mock/tool-approvals.spec.ts new file mode 100644 index 0000000000..7e38a4059b --- /dev/null +++ b/e2e/specs/mock/tool-approvals.spec.ts @@ -0,0 +1,835 @@ +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { expect, test } from '@playwright/test'; +import type { Locator, Page, Request, Route } from '@playwright/test'; +import type { AgentDetail } from './agents.helpers'; +import { cleanupAgent, openAgentBuilder, uniqueAgentName } from './agents.helpers'; +import { + MOCK_ENDPOINTS, + NEW_CHAT_PATH, + fetchJson, + getAccessToken, + messagesView, + requestJson, + sendMessage, +} from './helpers'; + +const MCP_SERVER_NAME = 'e2e-memory'; +const MCP_SERVER_TOOL_ID = `sys__server__sys_mcp_${MCP_SERVER_NAME}`; +const APPROVAL_TOOL_NAME = 'approval_probe'; +const APPROVAL_TOOL_ID = `${APPROVAL_TOOL_NAME}_mcp_${MCP_SERVER_NAME}`; +const APPROVAL_PROMPT_MARKER = 'E2E_TOOL_APPROVAL:'; +const BATCH_APPROVAL_PROMPT_MARKER = 'E2E_TOOL_APPROVAL_BATCH:'; +const RESTRICTED_APPROVAL_PROMPT_MARKER = 'E2E_TOOL_APPROVAL_RESTRICTED:'; +const REWRITTEN_APPROVAL_PROMPT_MARKER = 'E2E_TOOL_APPROVAL_REWRITE:'; +const APPROVAL_REASON = `E2E approval required before running ${APPROVAL_TOOL_ID}.`; +const APPROVAL_ERROR = 'Something went wrong submitting your decision. Please try again.'; +const APPROVAL_EXPIRED = 'This request expired or was already handled.'; +const DESCRIPTION = 'Verifies human approval behavior for MCP tool calls in mock E2E tests.'; +const APPROVAL_AUDIT_DIR = path.join('/tmp', 'librechat-e2e-approval-audit'); +const uniqueLabel = () => `${Date.now()}-${Math.floor(Math.random() * 1e4)}`; +const approvalInvocationPath = (value: string) => + path.join(APPROVAL_AUDIT_DIR, Buffer.from(value).toString('base64url')); + +function clearApprovalInvocations(...values: string[]) { + values.forEach((value) => fs.rmSync(approvalInvocationPath(value), { force: true })); +} + +function approvalInvocationCount(value: string) { + const filename = approvalInvocationPath(value); + if (!fs.existsSync(filename)) { + return 0; + } + return fs + .readFileSync(filename, 'utf8') + .split('\n') + .filter((line) => line.length > 0).length; +} + +async function expectApprovalInvocationCount(value: string, count: number) { + await expect.poll(() => approvalInvocationCount(value), { timeout: 30000 }).toBe(count); +} + +type MCPToolsResponse = { + servers?: Record }>; +}; + +type ApprovalResumeBody = { + actionId?: string; + agent_id?: string; + conversationId?: string; + endpoint?: string; + decisions?: Array<{ + tool_call_id?: string; + decision?: string; + reason?: string; + responseText?: string; + editedArguments?: Record; + }>; +}; + +type ApprovalResumeResponse = { + conversationId?: string; + status?: string; + streamId?: string; +}; + +const approvalCards = (page: Page) => messagesView(page).getByTestId('tool-approval'); +const approvalCard = (page: Page, toolCallId: string) => + messagesView(page).locator(`[data-testid="tool-approval"][data-tool-call-id="${toolCallId}"]`); + +function isResumeRequest(request: Request) { + return ( + request.method() === 'POST' && new URL(request.url()).pathname === '/api/agents/chat/resume' + ); +} + +async function waitForApprovalTool(page: Page) { + const token = await getAccessToken(page); + let latestTools: MCPToolsResponse | null = null; + + for (let attempt = 0; attempt < 20; attempt++) { + latestTools = await fetchJson(page, '/api/mcp/tools', token); + const tools = latestTools.servers?.[MCP_SERVER_NAME]?.tools ?? []; + if (tools.some((tool) => tool.pluginKey === APPROVAL_TOOL_ID)) { + return; + } + await page.waitForTimeout(500); + } + + expect( + latestTools?.servers?.[MCP_SERVER_NAME]?.tools, + `Expected ${MCP_SERVER_NAME} to expose ${APPROVAL_TOOL_ID}`, + ).toEqual(expect.arrayContaining([expect.objectContaining({ pluginKey: APPROVAL_TOOL_ID })])); +} + +async function createAndSelectApprovalAgent(page: Page): Promise { + await page.goto(NEW_CHAT_PATH, { timeout: 10000 }); + await waitForApprovalTool(page); + + const token = await getAccessToken(page); + const agentName = uniqueAgentName('E2E Tool Approval Agent'); + const agent = await requestJson(page, { + path: '/api/agents', + token, + method: 'POST', + body: { + name: agentName, + description: DESCRIPTION, + instructions: 'Use the requested approval probe tools and report their results.', + provider: MOCK_ENDPOINTS[0].label, + model: MOCK_ENDPOINTS[0].model, + tools: [MCP_SERVER_TOOL_ID, APPROVAL_TOOL_ID], + }, + }); + + const form = await openAgentBuilder(page); + await form.getByRole('combobox', { name: 'Agent', exact: true }).click(); + await page.getByRole('option', { name: agentName }).click(); + await expect(form.getByLabel('Agent name')).toHaveValue(agentName); + await form.getByRole('button', { name: 'Select Agent' }).click(); + return agent.id; +} + +async function startApproval( + page: Page, + label: string, + marker = APPROVAL_PROMPT_MARKER, + expectedReason = APPROVAL_REASON, +): Promise { + const response = await sendMessage(page, `${marker}${label}`); + expect(response.ok()).toBeTruthy(); + await expect(page).toHaveURL(/\/c\/(?!new)/, { timeout: 15000 }); + const card = approvalCards(page).first(); + await expect(card).toBeVisible({ timeout: 30000 }); + await expect(card).toContainText(expectedReason); + return card; +} + +async function submitAndCapture(page: Page, submit: Locator) { + const [request, response] = await Promise.all([ + page.waitForRequest(isResumeRequest), + page.waitForResponse( + (candidate) => isResumeRequest(candidate.request()) && candidate.status() === 200, + ), + submit.click(), + ]); + return { + body: request.postDataJSON() as ApprovalResumeBody, + response, + }; +} + +async function expectCompletedApprovalToolOutput(page: Page, toolCallId: string, output: string) { + const view = messagesView(page); + const groupToggle = view.getByRole('button', { name: /^Used \d+ tools$/ }).last(); + const toolCall = view.locator(`[data-testid="tool-call"][data-tool-call-id="${toolCallId}"]`); + + // On reload, the conversation arrives asynchronously and multi-tool groups + // start collapsed. Wait for either the target card or its group before + // deciding whether expansion is necessary. + await expect(toolCall.or(groupToggle).first()).toBeVisible({ timeout: 30000 }); + if ( + !(await toolCall.isVisible()) && + (await groupToggle.getAttribute('aria-expanded')) !== 'true' + ) { + await groupToggle.click(); + } + + await expect(toolCall).toBeVisible({ timeout: 30000 }); + const toggle = toolCall.getByRole('button', { name: /Ran approval_probe/ }); + await expect(toggle).toBeVisible({ timeout: 30000 }); + if ((await toggle.getAttribute('aria-expanded')) !== 'true') { + await toggle.click(); + } + + // Scope exact output to its stable call id. This catches both a dropped + // completion and an output accidentally attached to a sibling tool card. + await expect( + view.locator(`[data-tool-call-output-id="${toolCallId}"]`).getByText(output, { exact: true }), + ).toBeVisible({ timeout: 30000 }); + // The final model turn is the quiescence barrier: all parallel tool work + // has settled before invocation-count assertions inspect the audit. + await expect(view.getByText(/^E2E approval outcomes:/).last()).toBeVisible({ timeout: 30000 }); +} + +test.describe('tool approvals', () => { + test('approves a paused tool with its original arguments', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + + await expect(card.getByRole('button', { name: 'Approve' })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Reject' })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Edit' })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Respond' })).toBeVisible(); + + const submit = card.getByRole('button', { name: 'Submit' }); + await expect(submit).toBeDisabled(); + await card.getByRole('button', { name: 'Approve' }).click(); + await expect(submit).toBeEnabled(); + + const conversationId = new URL(page.url()).pathname.replace('/c/', ''); + const { body, response } = await submitAndCapture(page, submit); + expect(body.actionId).toBeTruthy(); + expect(body.agent_id).toBe(agentId); + expect(body.conversationId).toBe(conversationId); + expect(body.endpoint).toBe('agents'); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'approve', + tool_call_id: toolCallId, + }), + ]); + await expect(response.json() as Promise).resolves.toEqual( + expect.objectContaining({ + conversationId, + status: 'resuming', + streamId: conversationId, + }), + ); + + await expectCompletedApprovalToolOutput( + page, + toolCallId, + `E2E approval probe executed: ${originalValue}`, + ); + await expectApprovalInvocationCount(originalValue, 1); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('rejects with an optional reason without executing the tool', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const reason = `do not run ${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const submit = card.getByRole('button', { name: 'Submit' }); + + await card.getByRole('button', { name: 'Reject' }).click(); + await card.getByRole('textbox', { name: 'Reject' }).fill(` ${reason} `); + await expect(submit).toBeEnabled(); + + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'reject', + reason, + tool_call_id: toolCallId, + }), + ]); + + await expectCompletedApprovalToolOutput(page, toolCallId, `Blocked: ${reason}`); + await expectApprovalInvocationCount(originalValue, 0); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('requires edited arguments to be a JSON object and executes only the edit', async ({ + page, + }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const editedValue = `edited-${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue, editedValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const submit = card.getByRole('button', { name: 'Submit' }); + + await card.getByRole('button', { name: 'Edit' }).click(); + const editor = card.getByRole('textbox', { name: 'Edit' }); + await expect(editor).toHaveValue(new RegExp(`original-${label}`)); + + for (const invalid of ['{', 'null', '[]', '"text"']) { + await editor.fill(invalid); + await expect(card.getByText('Invalid JSON')).toBeVisible(); + await expect(submit).toBeDisabled(); + } + + await editor.fill(JSON.stringify({ value: editedValue })); + await expect(card.getByText('Invalid JSON')).toHaveCount(0); + await expect(submit).toBeEnabled(); + + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'edit', + editedArguments: { value: editedValue }, + tool_call_id: toolCallId, + }), + ]); + + await expectCompletedApprovalToolOutput( + page, + toolCallId, + `E2E approval probe executed: ${editedValue}`, + ); + await expectApprovalInvocationCount(editedValue, 1); + await expectApprovalInvocationCount(originalValue, 0); + } finally { + clearApprovalInvocations(originalValue, editedValue); + await cleanupAgent(page, agentId); + } + }); + + test('requires a nonblank substitute response and skips tool execution', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const responseText = `manual result ${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const submit = card.getByRole('button', { name: 'Submit' }); + + await card.getByRole('button', { name: 'Respond' }).click(); + const responseInput = card.getByRole('textbox', { name: 'Respond' }); + await responseInput.fill(' '); + await expect(submit).toBeDisabled(); + await responseInput.fill(` ${responseText} `); + await expect(submit).toBeEnabled(); + + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'respond', + responseText, + tool_call_id: toolCallId, + }), + ]); + + await expectCompletedApprovalToolOutput(page, toolCallId, responseText); + await expectApprovalInvocationCount(originalValue, 0); + } finally { + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('honors a hook-restricted decision set', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval( + page, + label, + RESTRICTED_APPROVAL_PROMPT_MARKER, + APPROVAL_REASON, + ); + + await expect(card.getByRole('button', { name: 'Approve' })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Reject' })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Edit' })).toHaveCount(0); + await expect(card.getByRole('button', { name: 'Respond' })).toHaveCount(0); + + const submit = card.getByRole('button', { name: 'Submit' }); + await card.getByRole('button', { name: 'Approve' }).click(); + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'approve', + tool_call_id: toolCallId, + }), + ]); + await expectCompletedApprovalToolOutput( + page, + toolCallId, + `E2E approval probe executed: ${originalValue}`, + ); + await expectApprovalInvocationCount(originalValue, 1); + } finally { + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('reviews and approves the authoritative hook-rewritten arguments', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const rewrittenValue = `rewritten-${label}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue, rewrittenValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval( + page, + label, + REWRITTEN_APPROVAL_PROMPT_MARKER, + APPROVAL_REASON, + ); + + await card.getByRole('button', { name: 'Edit' }).click(); + const editor = card.getByRole('textbox', { name: 'Edit' }); + await expect(editor).toHaveValue(new RegExp(`rewritten-${label}`)); + await expect(editor).not.toHaveValue(new RegExp(`original-${label}`)); + + await card.getByRole('button', { name: 'Edit' }).click(); + const submit = card.getByRole('button', { name: 'Submit' }); + await card.getByRole('button', { name: 'Approve' }).click(); + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toEqual([ + expect.objectContaining({ + decision: 'approve', + tool_call_id: toolCallId, + }), + ]); + + await expectCompletedApprovalToolOutput( + page, + toolCallId, + `E2E approval probe executed: ${rewrittenValue}`, + ); + await expectApprovalInvocationCount(rewrittenValue, 1); + await expectApprovalInvocationCount(originalValue, 0); + } finally { + clearApprovalInvocations(originalValue, rewrittenValue); + await cleanupAgent(page, agentId); + } + }); + + test('submits a mixed batch once and preserves decisions through collapse', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const firstCallId = `call_e2e_approval_${label}_first`; + const secondCallId = `call_e2e_approval_${label}_second`; + const firstValue = `first-${label}`; + const secondValue = `second-${label}`; + const responseText = `manual batch result ${label}`; + let agentId: string | undefined; + clearApprovalInvocations(firstValue, secondValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + await startApproval(page, label, BATCH_APPROVAL_PROMPT_MARKER); + const conversationPath = new URL(page.url()).pathname; + await expect(approvalCards(page)).toHaveCount(2); + + // Reconstruct both pending cards from persisted state before making any + // decisions, not just the simpler one-call resume path. + await page.reload({ waitUntil: 'domcontentloaded' }); + await expect.poll(() => new URL(page.url()).pathname).toBe(conversationPath); + await expect(approvalCards(page)).toHaveCount(2); + + const firstCard = approvalCard(page, firstCallId); + const secondCard = approvalCard(page, secondCallId); + const submit = messagesView(page).getByRole('button', { + name: 'Submit 2 decisions', + exact: true, + }); + + await secondCard.getByRole('button', { name: 'Respond' }).click(); + await secondCard.getByRole('textbox', { name: 'Respond' }).fill(responseText); + await expect(submit).toBeDisabled(); + await firstCard.getByRole('button', { name: 'Approve' }).click(); + await expect(submit).toBeEnabled(); + + const groupToggle = messagesView(page).getByRole('button', { + name: 'Used 2 tools', + exact: true, + }); + const groupPanel = messagesView(page).getByTestId('tool-call-group-panel').last(); + await Promise.all([ + groupPanel.evaluate( + (element) => + new Promise((resolve) => { + const handleTransitionEnd = (event: Event) => { + if ( + event.target === element && + (event as TransitionEvent).propertyName === 'grid-template-rows' + ) { + element.removeEventListener('transitionend', handleTransitionEnd); + resolve(); + } + }; + element.addEventListener('transitionend', handleTransitionEnd); + }), + ), + groupToggle.click(), + ]); + await expect(groupToggle).toHaveAttribute('aria-expanded', 'false'); + await groupToggle.click(); + await expect(groupToggle).toHaveAttribute('aria-expanded', 'true'); + + const reopenedFirstCard = approvalCard(page, firstCallId); + const reopenedSecondCard = approvalCard(page, secondCallId); + await expect(reopenedFirstCard.getByRole('button', { name: 'Approve' })).toHaveAttribute( + 'aria-pressed', + 'true', + ); + await expect(reopenedSecondCard.getByRole('button', { name: 'Respond' })).toHaveAttribute( + 'aria-pressed', + 'true', + ); + await expect(reopenedSecondCard.getByRole('textbox', { name: 'Respond' })).toHaveValue( + responseText, + ); + await expect(submit).toBeEnabled(); + + const { body } = await submitAndCapture(page, submit); + expect(body.decisions).toHaveLength(2); + expect(body.decisions).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + decision: 'approve', + tool_call_id: firstCallId, + }), + expect.objectContaining({ + decision: 'respond', + responseText, + tool_call_id: secondCallId, + }), + ]), + ); + + await expectCompletedApprovalToolOutput( + page, + firstCallId, + `E2E approval probe executed: ${firstValue}`, + ); + await expectCompletedApprovalToolOutput(page, secondCallId, responseText); + await expectApprovalInvocationCount(firstValue, 1); + await expectApprovalInvocationCount(secondValue, 0); + await expect(approvalCards(page)).toHaveCount(0); + + await page.reload({ waitUntil: 'domcontentloaded' }); + await expect.poll(() => new URL(page.url()).pathname).toBe(conversationPath); + await expectCompletedApprovalToolOutput( + page, + firstCallId, + `E2E approval probe executed: ${firstValue}`, + ); + await expectCompletedApprovalToolOutput(page, secondCallId, responseText); + await expectApprovalInvocationCount(firstValue, 1); + await expectApprovalInvocationCount(secondValue, 0); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + clearApprovalInvocations(firstValue, secondValue); + await cleanupAgent(page, agentId); + } + }); + + test('rehydrates a paused approval and its completed result across reloads', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const executedText = `E2E approval probe executed: ${originalValue}`; + let agentId: string | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + await startApproval(page, label); + const conversationPath = new URL(page.url()).pathname; + + await page.reload({ waitUntil: 'domcontentloaded' }); + await expect.poll(() => new URL(page.url()).pathname).toBe(conversationPath); + const rehydratedCard = approvalCard(page, toolCallId); + await expect(rehydratedCard).toBeVisible({ timeout: 30000 }); + await expect(rehydratedCard).toContainText(APPROVAL_REASON); + + await page.goto(NEW_CHAT_PATH, { waitUntil: 'domcontentloaded' }); + await expect(approvalCards(page)).toHaveCount(0); + await page.goto(conversationPath, { waitUntil: 'domcontentloaded' }); + const navigatedCard = approvalCard(page, toolCallId); + await expect(navigatedCard).toBeVisible({ timeout: 30000 }); + await expect(navigatedCard).toContainText(APPROVAL_REASON); + + await navigatedCard.getByRole('button', { name: 'Approve' }).click(); + await submitAndCapture(page, navigatedCard.getByRole('button', { name: 'Submit' })); + await expectCompletedApprovalToolOutput(page, toolCallId, executedText); + await expectApprovalInvocationCount(originalValue, 1); + await expect(approvalCards(page)).toHaveCount(0); + + await page.reload({ waitUntil: 'domcontentloaded' }); + await expect.poll(() => new URL(page.url()).pathname).toBe(conversationPath); + await expectCompletedApprovalToolOutput(page, toolCallId, executedText); + await expectApprovalInvocationCount(originalValue, 1); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('sends only one resume request for two synchronous submit clicks', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const executedText = `E2E approval probe executed: ${originalValue}`; + let agentId: string | undefined; + let releaseResume = () => undefined; + let resumeHandler: ((route: Route) => Promise) | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const submit = card.getByRole('button', { name: 'Submit' }); + await card.getByRole('button', { name: 'Approve' }).click(); + await expect(submit).toBeEnabled(); + + let resumeRequests = 0; + const resumeGate = new Promise((resolve) => { + releaseResume = resolve; + }); + resumeHandler = async (route) => { + resumeRequests++; + if (resumeRequests === 1) { + await resumeGate; + await route.continue(); + return; + } + await route.fulfill({ + status: 409, + contentType: 'application/json', + body: JSON.stringify({ message: 'duplicate resume request' }), + }); + }; + await page.route('**/api/agents/chat/resume', resumeHandler); + + await submit.evaluate((button: HTMLButtonElement) => { + button.click(); + button.click(); + }); + await page.waitForTimeout(250); + expect(resumeRequests).toBe(1); + await expect(card.getByRole('button', { name: 'Submitting' })).toBeDisabled(); + await expect(card.getByRole('button', { name: 'Approve' })).toBeDisabled(); + releaseResume(); + + await expectCompletedApprovalToolOutput(page, toolCallId, executedText); + await expectApprovalInvocationCount(originalValue, 1); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + releaseResume(); + if (resumeHandler) { + await page.unroute('**/api/agents/chat/resume', resumeHandler); + } + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('preserves a decision after a transient resume error and retries successfully', async ({ + page, + }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const responseText = `retry response ${label}`; + let agentId: string | undefined; + let resumeHandler: ((route: Route) => Promise) | undefined; + clearApprovalInvocations(originalValue); + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const submit = card.getByRole('button', { name: 'Submit' }); + await card.getByRole('button', { name: 'Respond' }).click(); + const responseInput = card.getByRole('textbox', { name: 'Respond' }); + await responseInput.fill(responseText); + + let resumeRequests = 0; + resumeHandler = async (route) => { + resumeRequests++; + if (resumeRequests === 1) { + await route.fulfill({ + status: 500, + contentType: 'application/json', + body: JSON.stringify({ message: 'temporary e2e failure' }), + }); + return; + } + await route.continue(); + }; + await page.route('**/api/agents/chat/resume', resumeHandler); + + await Promise.all([ + page.waitForResponse( + (response) => isResumeRequest(response.request()) && response.status() === 500, + ), + submit.click(), + ]); + await expect(card.getByText(APPROVAL_ERROR, { exact: true })).toBeVisible(); + await expect(card.getByRole('button', { name: 'Respond' })).toHaveAttribute( + 'aria-pressed', + 'true', + ); + await expect(responseInput).toHaveValue(responseText); + await expect(submit).toBeEnabled(); + + await Promise.all([ + page.waitForResponse( + (response) => isResumeRequest(response.request()) && response.status() === 200, + ), + submit.click(), + ]); + await expectCompletedApprovalToolOutput(page, toolCallId, responseText); + await expectApprovalInvocationCount(originalValue, 0); + expect(resumeRequests).toBe(2); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + if (resumeHandler) { + await page.unroute('**/api/agents/chat/resume', resumeHandler); + } + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); + + test('locks the approval controls and explains an expired resume action', async ({ page }) => { + test.setTimeout(120000); + const label = uniqueLabel(); + const toolCallId = `call_e2e_approval_${label}`; + const originalValue = `original-${label}`; + const executedText = `E2E approval probe executed: ${originalValue}`; + let agentId: string | undefined; + let capturedResumeBody: Record | undefined; + let backendResolved = false; + let routeInstalled = false; + clearApprovalInvocations(originalValue); + const resumeHandler = async (route: Route) => { + capturedResumeBody = route.request().postDataJSON() as Record; + await route.fulfill({ + status: 409, + contentType: 'application/json', + body: JSON.stringify({ message: 'expired e2e action' }), + }); + }; + + try { + agentId = await createAndSelectApprovalAgent(page); + const card = await startApproval(page, label); + const approve = card.getByRole('button', { name: 'Approve' }); + const submit = card.getByRole('button', { name: 'Submit' }); + await approve.click(); + await page.route('**/api/agents/chat/resume', resumeHandler); + routeInstalled = true; + + await Promise.all([ + page.waitForResponse( + (response) => isResumeRequest(response.request()) && response.status() === 409, + ), + submit.click(), + ]); + await expect(card.getByText(APPROVAL_EXPIRED, { exact: true })).toBeVisible(); + await expect(approve).toHaveAttribute('aria-pressed', 'true'); + await expect(approve).toBeDisabled(); + await expect(card.getByRole('button', { name: 'Reject' })).toBeDisabled(); + await expect(card.getByRole('button', { name: 'Edit' })).toBeDisabled(); + await expect(card.getByRole('button', { name: 'Respond' })).toBeDisabled(); + await expect(submit).toBeDisabled(); + expect(capturedResumeBody).toBeDefined(); + + await page.unroute('**/api/agents/chat/resume', resumeHandler); + routeInstalled = false; + const token = await getAccessToken(page); + await requestJson(page, { + path: '/api/agents/chat/resume', + token, + method: 'POST', + body: capturedResumeBody, + }); + backendResolved = true; + await expectCompletedApprovalToolOutput(page, toolCallId, executedText); + await expectApprovalInvocationCount(originalValue, 1); + await expect(approvalCards(page)).toHaveCount(0); + } finally { + if (routeInstalled) { + await page.unroute('**/api/agents/chat/resume', resumeHandler); + } + if (!backendResolved && capturedResumeBody) { + const token = await getAccessToken(page); + await requestJson(page, { + path: '/api/agents/chat/resume', + token, + method: 'POST', + body: capturedResumeBody, + }).catch(() => undefined); + } + clearApprovalInvocations(originalValue); + await cleanupAgent(page, agentId); + } + }); +}); diff --git a/packages/api/src/agents/hitl/resume.spec.ts b/packages/api/src/agents/hitl/resume.spec.ts index b199cf2eab..93b9d50b1e 100644 --- a/packages/api/src/agents/hitl/resume.spec.ts +++ b/packages/api/src/agents/hitl/resume.spec.ts @@ -1,3 +1,10 @@ +import { + ContentTypes, + GraphEvents, + StepTypes, + createContentAggregator, + type RunStep, +} from '@librechat/agents'; import type { Agents } from 'librechat-data-provider'; import { mapToolApprovalResolutions, @@ -6,6 +13,7 @@ import { findDisallowedDecisions, findIncompleteDecisions, createContentIndexOffsetHandlers, + hydrateResumeRunSteps, attachAskUserQuestionAnswer, attachAskUserQuestionArgs, } from './resume'; @@ -264,6 +272,136 @@ describe('createContentIndexOffsetHandlers', () => { }); }); +describe('hydrateResumeRunSteps', () => { + it('restores step and tool-call identity so a resumed completion updates its seeded card', () => { + const { contentParts, aggregateContent, stepMap } = createContentAggregator(); + contentParts.push( + { type: ContentTypes.TEXT, text: 'Before approval' }, + { + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-approval', + name: 'approval_probe', + args: '{"value":"before"}', + }, + }, + ); + const runStep: RunStep = { + id: 'step-approval', + runId: 'response-1', + type: StepTypes.TOOL_CALLS, + index: 1, + stepDetails: { + type: StepTypes.TOOL_CALLS, + tool_calls: [ + { + id: 'call-approval', + name: 'approval_probe', + args: { value: 'before' }, + }, + ], + }, + usage: null, + }; + const toolCallStepIds = new Map(); + + hydrateResumeRunSteps([runStep], stepMap, { toolCallStepIds }, contentParts); + aggregateContent({ + event: GraphEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id: 'step-approval', + index: 1, + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-approval', + name: 'approval_probe', + args: { value: 'before' }, + output: 'approved output', + }, + }, + }, + }); + + expect(stepMap.get('step-approval')).toBe(runStep); + expect(toolCallStepIds.get('call-approval')).toBe('step-approval'); + expect(contentParts[1]).toMatchObject({ + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-approval', + output: 'approved output', + progress: 1, + }, + }); + }); + + it('realigns a stale persisted index to the seeded tool card by tool-call id', () => { + const { contentParts, aggregateContent, stepMap } = createContentAggregator(); + contentParts.push( + { type: ContentTypes.TEXT, text: 'Prepended after the step was recorded' }, + { + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-shifted', + name: 'approval_probe', + args: '{"value":"shifted"}', + }, + }, + ); + const staleRunStep: RunStep = { + id: 'step-shifted', + runId: 'response-1', + type: StepTypes.TOOL_CALLS, + index: 0, + stepDetails: { + type: StepTypes.TOOL_CALLS, + tool_calls: [ + { + id: 'call-shifted', + name: 'approval_probe', + args: { value: 'shifted' }, + }, + ], + }, + usage: null, + }; + const toolCallStepIds = new Map(); + + hydrateResumeRunSteps([staleRunStep], stepMap, { toolCallStepIds }, contentParts); + aggregateContent({ + event: GraphEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id: 'step-shifted', + index: 0, + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-shifted', + name: 'approval_probe', + args: { value: 'shifted' }, + output: 'shifted output', + }, + }, + }, + }); + + expect(staleRunStep.index).toBe(0); + expect(stepMap.get('step-shifted')?.index).toBe(1); + expect(contentParts[0]).toEqual({ + type: ContentTypes.TEXT, + text: 'Prepended after the step was recorded', + }); + expect(contentParts[1]).toMatchObject({ + type: ContentTypes.TOOL_CALL, + tool_call: { + id: 'call-shifted', + output: 'shifted output', + progress: 1, + }, + }); + }); +}); + describe('attachAskUserQuestionAnswer', () => { const question = { question: 'Which env?', options: [{ label: 'Staging', value: 'staging' }] }; const askPart = (output?: string) => ({ diff --git a/packages/api/src/agents/hitl/resume.ts b/packages/api/src/agents/hitl/resume.ts index 58b465221c..4134ccb510 100644 --- a/packages/api/src/agents/hitl/resume.ts +++ b/packages/api/src/agents/hitl/resume.ts @@ -4,6 +4,7 @@ import type { ToolApprovalDecisionMap, AskUserQuestionResolution, EventHandler, + RunStep, } from '@librechat/agents'; import type { Agents } from 'librechat-data-provider'; import { ASK_USER_QUESTION_TOOL_NAME } from './askUserQuestionTool'; @@ -121,6 +122,82 @@ export function findIncompleteDecisions( .map((r) => r.tool_call_id); } +/** + * Reconcile persisted tool-step indices with the content being seeded into a + * rebuilt aggregator. + * + * A pause-time index is not durable identity: hosts can prepend content after + * the step was emitted, and persisted reconstruction can compact sparse + * content. Tool-call ids are stable across both operations, so use them to + * relocate a step without mutating the stored object. + */ +type ResumableRunStep = { + id: string; + index: number; + stepDetails: { + type: string; + tool_calls?: readonly { id?: string }[]; + }; +}; + +export function normalizeResumeRunStepIndices( + runSteps: readonly T[], + seedContent: readonly { type?: string; tool_call?: { id?: string } }[] = [], +): T[] { + const toolCallIndices = new Map(); + seedContent.forEach((part, index) => { + const toolCallId = part?.tool_call?.id; + if (part?.type === 'tool_call' && typeof toolCallId === 'string') { + toolCallIndices.set(toolCallId, index); + } + }); + + return runSteps.map((runStep) => { + if (runStep.stepDetails.type !== 'tool_calls') { + return runStep; + } + const contentIndex = runStep.stepDetails.tool_calls + ?.map((toolCall) => (toolCall.id ? toolCallIndices.get(toolCall.id) : undefined)) + .find((index) => index != null); + return contentIndex != null && contentIndex !== runStep.index + ? { ...runStep, index: contentIndex } + : runStep; + }); +} + +/** + * Restore the streamed run-step sidecars that a fresh SDK Run cannot recover + * from the LangGraph checkpoint by itself. + * + * Human-review resume can happen in a later request or process. The checkpoint + * restarts directly inside ToolNode, so it does not replay ON_RUN_STEP before + * dispatching ON_RUN_STEP_COMPLETED. Seeding both maps lets the ToolNode emit + * the original step id and lets the content aggregator resolve that id back to + * the already-rendered tool card. + */ +export function hydrateResumeRunSteps( + runSteps: readonly RunStep[], + stepMap: Map | undefined, + graph: { toolCallStepIds?: Map } | null | undefined, + seedContent: readonly { type?: string; tool_call?: { id?: string } }[] = [], +): void { + for (const runStep of normalizeResumeRunStepIndices(runSteps, seedContent)) { + if (!runStep?.id) { + continue; + } + stepMap?.set(runStep.id, runStep); + const stepDetails: ResumableRunStep['stepDetails'] = runStep.stepDetails; + if (stepDetails.type !== 'tool_calls') { + continue; + } + for (const toolCall of stepDetails.tool_calls ?? []) { + if (toolCall.id) { + graph?.toolCallStepIds?.set(toolCall.id, runStep.id); + } + } + } +} + /** * Wrap a resume run's event handlers so every content index the rebuilt graph * emits is shifted past the pre-pause content. diff --git a/packages/api/src/agents/memory.spec.ts b/packages/api/src/agents/memory.spec.ts index 463a115d00..22f24380ac 100644 --- a/packages/api/src/agents/memory.spec.ts +++ b/packages/api/src/agents/memory.spec.ts @@ -1,10 +1,11 @@ import { Types } from 'mongoose'; -import { Run, Providers } from '@librechat/agents'; -import { MemoryScope } from 'librechat-data-provider'; +import { Tools, MemoryScope } from 'librechat-data-provider'; +import { Run, Providers, GraphEvents } from '@librechat/agents'; import type { IUser } from '@librechat/data-schemas'; import type { Response } from 'express'; import { processMemory, + createMemoryProcessor, createMemoryTool, getMemoryAgentId, getRequestMemories, @@ -12,6 +13,7 @@ import { invalidateRequestMemories, agentHasInlineMemoryTools, } from './memory'; +import { GenerationJobManager } from '~/stream/GenerationJobManager'; jest.mock('~/stream/GenerationJobManager'); @@ -98,6 +100,72 @@ function createTestUser(overrides: Partial = {}): IUser { } as IUser; } +describe('Memory attachment generation fencing', () => { + it('emits artifacts with the generation epoch that started memory processing', async () => { + const memoryArtifact = { + type: 'update' as const, + key: 'response_style', + value: 'concise', + }; + const processStream = jest.fn(async () => { + const runConfig = (Run.create as jest.Mock).mock.calls[0][0]; + runConfig.customHandlers[GraphEvents.TOOL_END].handle( + GraphEvents.TOOL_END, + { + output: { + tool_call_id: 'memory-call-1', + artifact: { [Tools.memory]: memoryArtifact }, + }, + }, + { + run_id: 'response-1', + thread_id: 'conversation-1', + }, + ); + return 'success'; + }); + (Run.create as jest.Mock).mockReturnValueOnce({ processStream }); + + const [, runMemory] = await createMemoryProcessor({ + res: { + headersSent: true, + write: jest.fn(), + } as unknown as Response, + userId: 'user-1', + messageId: 'response-1', + conversationId: 'conversation-1', + streamId: 'conversation-1', + jobCreatedAt: 1234, + memoryMethods: { + setMemory: jest.fn(), + deleteMemory: jest.fn(), + getFormattedMemories: jest.fn().mockResolvedValue({ + withKeys: '', + withoutKeys: '', + totalTokens: 0, + }), + }, + }); + + await runMemory([]); + + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'conversation-1', + { + event: 'attachment', + data: { + type: Tools.memory, + toolCallId: 'memory-call-1', + messageId: 'response-1', + conversationId: 'conversation-1', + [Tools.memory]: memoryArtifact, + }, + }, + { expectedCreatedAt: 1234 }, + ); + }); +}); + describe('Memory Agent Header Resolution', () => { let testUser: IUser; let mockRes: Response; diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index 1f4c362be8..b7922ec7dd 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -701,6 +701,7 @@ export async function processMemory({ tokenLimit, totalTokens = 0, streamId = null, + jobCreatedAt, user, }: { res: ServerResponse; @@ -719,6 +720,7 @@ export async function processMemory({ totalTokens?: number; llmConfig?: Partial; streamId?: string | null; + jobCreatedAt?: number; user?: IUser; }): Promise<(TAttachment | null)[] | undefined> { try { @@ -826,7 +828,12 @@ ${memory ?? 'No existing memories'}`; }); const artifactPromises: Promise[] = []; - const memoryCallback = createMemoryCallback({ res, artifactPromises, streamId }); + const memoryCallback = createMemoryCallback({ + res, + artifactPromises, + streamId, + jobCreatedAt, + }); const customHandlers = { [GraphEvents.TOOL_END]: new BasicToolEndHandler(memoryCallback), }; @@ -926,6 +933,7 @@ export async function createMemoryProcessor({ conversationId, config = {}, streamId = null, + jobCreatedAt, user, }: { res: ServerResponse; @@ -937,6 +945,7 @@ export async function createMemoryProcessor({ memoryMethods: RequiredMemoryMethods; config?: MemoryConfig; streamId?: string | null; + jobCreatedAt?: number; user?: IUser; }): Promise<[string, (messages: BaseMessage[]) => Promise<(TAttachment | null)[] | undefined>]> { const { validKeys, instructions, llmConfig, tokenLimit } = config; @@ -961,6 +970,7 @@ export async function createMemoryProcessor({ messageId, tokenLimit, streamId, + jobCreatedAt, conversationId, memory: withKeys, totalTokens: totalTokens || 0, @@ -981,11 +991,13 @@ async function handleMemoryArtifact({ data, metadata, streamId = null, + jobCreatedAt, }: { res: ServerResponse; data: ToolEndData; metadata?: ToolEndMetadata; streamId?: string | null; + jobCreatedAt?: number; }) { const output = data?.output as ToolMessage | undefined; if (!output) { @@ -1012,7 +1024,11 @@ async function handleMemoryArtifact({ return attachment; } if (streamId) { - GenerationJobManager.emitChunk(streamId, { event: 'attachment', data: attachment }); + GenerationJobManager.emitChunk( + streamId, + { event: 'attachment', data: attachment }, + { expectedCreatedAt: jobCreatedAt }, + ); } else { res.write(`event: attachment\ndata: ${JSON.stringify(attachment)}\n\n`); } @@ -1025,16 +1041,19 @@ async function handleMemoryArtifact({ * @param params.res - The server response object * @param params.artifactPromises - Array to collect artifact promises * @param params.streamId - The stream ID for resumable mode, or null for standard mode + * @param params.jobCreatedAt - The generation epoch that owns emitted artifacts * @returns The memory callback function */ export function createMemoryCallback({ res, artifactPromises, streamId = null, + jobCreatedAt, }: { res: ServerResponse; artifactPromises: Promise | null>[]; streamId?: string | null; + jobCreatedAt?: number; }): ToolEndCallback { return async (data: ToolEndData, metadata?: Record) => { const output = data?.output as ToolMessage | undefined; @@ -1043,7 +1062,7 @@ export function createMemoryCallback({ return; } artifactPromises.push( - handleMemoryArtifact({ res, data, metadata, streamId }).catch((error) => { + handleMemoryArtifact({ res, data, metadata, streamId, jobCreatedAt }).catch((error) => { logger.error('Error processing memory artifact content:', error); return null; }), diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 272b9d77a0..262ec93f95 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -40,6 +40,7 @@ import { import { isPendingActionStale, isPendingActionExpired } from './interfaces/IJobStore'; import { InMemoryEventTransport } from './implementations/InMemoryEventTransport'; import { InMemoryJobStore } from './implementations/InMemoryJobStore'; +import { normalizeResumeRunStepIndices } from '~/agents/hitl/resume'; import { emitChunkWithReceipt } from './internal/chunkPublication'; import { filterPersistableAbortContent } from './abortContent'; import { toClientPendingAction } from '~/agents/hitl/policy'; @@ -132,6 +133,34 @@ function isOAuthReplayEvent(event: t.ServerSentEvent): boolean { return false; } +function normalizeRunStepReplayIndices( + replayEvents: t.ResumeState['replayEvents'], + runSteps: readonly Agents.RunStep[], +): t.ResumeState['replayEvents'] { + if (!replayEvents) { + return replayEvents; + } + const runStepsById = new Map(runSteps.map((runStep) => [runStep.id, runStep])); + return replayEvents.map((event) => { + if (event.event !== 'on_run_step' || event.data == null || typeof event.data !== 'object') { + return event; + } + const stepId = 'id' in event.data ? event.data.id : undefined; + const normalizedRunStep = typeof stepId === 'string' ? runStepsById.get(stepId) : undefined; + const eventIndex = 'index' in event.data ? event.data.index : undefined; + if (!normalizedRunStep || normalizedRunStep.index === eventIndex) { + return event; + } + return { + ...event, + data: { + ...event.data, + index: normalizedRunStep.index, + }, + }; + }); +} + /** * Configuration options for GenerationJobManager */ @@ -2323,10 +2352,14 @@ class GenerationJobManagerClass { async emitChunk( streamId: string, event: t.ServerSentEvent, - options?: { durable?: boolean }, + options?: { durable?: boolean; expectedCreatedAt?: number }, ): Promise { const runtime = this.runtimeState.get(streamId); - if (!runtime || !this.isCurrentRuntime(streamId, runtime)) { + if ( + !runtime || + (options?.expectedCreatedAt != null && runtime.createdAt !== options.expectedCreatedAt) || + !this.isCurrentRuntime(streamId, runtime) + ) { return; } @@ -2414,14 +2447,24 @@ class GenerationJobManagerClass { } markSnapshotReady(); + // Retain run-step identity independently of the live graph. Paused in-memory + // runs release that graph before a later request rebuilds the run, but the + // resume path still needs the original step ids to correlate tool results. + const eventObj = event as Record; + const eventType = eventObj.event as string | undefined; + const eventData = eventObj.data; + if ( + (eventType === 'on_run_step' || eventType === 'on_run_step_completed') && + eventData != null && + typeof eventData === 'object' + ) { + this.saveRunStepFromEvent(streamId, eventData as Record, runtime.createdAt); + } + // For Redis mode, persist chunk for later reconstruction (fire-and-forget for resumability) if (this._isRedis) { // The SSE event structure is { event: string, data: unknown, ... } // The aggregator expects { event: string, data: unknown } where data is the payload - const eventObj = event as Record; - const eventType = eventObj.event as string | undefined; - const eventData = eventObj.data; - if (eventType && eventData !== undefined) { // Store in format expected by aggregateContent: { event, data } const appendPromise = this.jobStore @@ -2430,15 +2473,6 @@ class GenerationJobManagerClass { logger.error(`[GenerationJobManager] Failed to append chunk:`, err); }); - // For run step events, also save to run steps key for quick retrieval - if (eventType === 'on_run_step' || eventType === 'on_run_step_completed') { - this.saveRunStepFromEvent( - streamId, - eventData as Record, - runtime.createdAt, - ); - } - if (options?.durable === true) { await appendPromise; if (!this.isCurrentRuntime(streamId, runtime)) { @@ -2567,9 +2601,9 @@ class GenerationJobManagerClass { } /** - * Accumulate run steps for a stream (Redis mode only). - * Uses a simple in-memory buffer that gets flushed to Redis. - * Not used in in-memory mode - run steps come from live graph via WeakRef. + * Accumulate run steps for a stream. + * Redis stores flush this buffer for cross-replica recovery; in-memory stores + * retain it as a fallback after a paused run's live graph has been released. */ private runStepBuffers: Map | null = null; @@ -2578,7 +2612,7 @@ class GenerationJobManagerClass { runStep: Agents.RunStep, expectedCreatedAt: number, ): void { - // Lazy initialization - only create map when first used (Redis mode) + // Lazy initialization keeps the per-stream allocation off non-agent paths. if (!this.runStepBuffers) { this.runStepBuffers = new Map(); } @@ -2970,6 +3004,16 @@ class GenerationJobManagerClass { this.jobStore.peekSteers(streamId, jobData.createdAt), ]); const aggregatedContent = result?.content ?? []; + const bufferState = this.runStepBuffers?.get(streamId); + const bufferedRunSteps = bufferState?.createdAt === jobData.createdAt ? bufferState.steps : []; + const runStepsById = new Map(runSteps.map((runStep) => [runStep.id, runStep])); + for (const runStep of bufferedRunSteps) { + runStepsById.set(runStep.id, runStep); + } + const effectiveRunSteps = normalizeResumeRunStepIndices( + [...runStepsById.values()], + aggregatedContent, + ); let titleEvent: t.ResumeState['titleEvent']; if (jobData.titleEvent) { try { @@ -2982,6 +3026,7 @@ class GenerationJobManagerClass { if (jobData.replayEvents) { try { replayEvents = JSON.parse(jobData.replayEvents) as t.ResumeState['replayEvents']; + replayEvents = normalizeRunStepReplayIndices(replayEvents, effectiveRunSteps); } catch { // Ignore malformed persisted replay events. } @@ -3013,13 +3058,13 @@ class GenerationJobManagerClass { logger.debug(`[GenerationJobManager] getResumeState:`, { streamId, - runStepsLength: runSteps.length, + runStepsLength: effectiveRunSteps.length, aggregatedContentLength: aggregatedContent.length, collectedUsageLength: collectedUsage?.length ?? 0, }); return { - runSteps, + runSteps: effectiveRunSteps, aggregatedContent, userMessage: jobData.userMessage, responseMessageId: jobData.responseMessageId, @@ -3334,7 +3379,7 @@ class GenerationJobManagerClass { this.eventTransport.cleanup(streamId); } - // Also check runStepBuffers for any orphaned entries (Redis mode only) + // Also check runStepBuffers for any orphaned entries. if (this.runStepBuffers) { for (const streamId of this.runStepBuffers.keys()) { if (!(await this.jobStore.hasJob(streamId))) { diff --git a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts index c9721efc00..82a726caae 100644 --- a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts +++ b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts @@ -1,3 +1,5 @@ +import type { StandardGraph } from '@librechat/agents'; +import type { Agents } from 'librechat-data-provider'; import type { ServerSentEvent } from '~/types'; import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; @@ -6,9 +8,13 @@ import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; jest.spyOn(console, 'log').mockImplementation(); function createInMemoryManager(): GenerationJobManagerClass { + return createManagerWithStore(new InMemoryJobStore({ ttlAfterComplete: 60000 })); +} + +function createManagerWithStore(store: InMemoryJobStore): GenerationJobManagerClass { const manager = new GenerationJobManagerClass(); manager.configure({ - jobStore: new InMemoryJobStore({ ttlAfterComplete: 60000 }), + jobStore: store, eventTransport: new InMemoryEventTransport(), isRedis: false, }); @@ -30,6 +36,18 @@ class SnapshotReplayJobStore extends InMemoryJobStore { } } +class PartialRunStepJobStore extends InMemoryJobStore { + private persistedRunSteps: Agents.RunStep[] = []; + + setPersistedRunSteps(runSteps: Agents.RunStep[]): void { + this.persistedRunSteps = runSteps; + } + + async getRunSteps(): Promise { + return this.persistedRunSteps; + } +} + function createSnapshotReplayManager(): GenerationJobManagerClass { const manager = new GenerationJobManagerClass(); manager.configure({ @@ -109,6 +127,226 @@ describe('GenerationJobManager resume replay events', () => { expect(resumeState?.replayEvents).toEqual([runStepEvent, authEvent]); }); + test('retains emitted run steps when the live graph is unavailable during resume', async () => { + manager = createInMemoryManager(); + const streamId = `run-step-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const runStep = { + id: 'step-approval', + runId: 'response-1', + index: 1, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-approval', name: 'approval_probe', args: '{}' }], + }, + }; + + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: runStep, + }); + + const resumeState = await manager.getResumeState(streamId); + + expect(resumeState?.runSteps).toEqual([runStep]); + }); + + test('realigns a stale run-step index to the aggregated tool card by tool-call id', async () => { + manager = createInMemoryManager(); + const streamId = `run-step-index-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + manager.setContentParts(streamId, [ + { type: 'text', text: 'Prepended content' }, + { + type: 'tool_call', + tool_call: { + id: 'call-shifted', + name: 'approval_probe', + args: '{"value":"shifted"}', + }, + }, + ]); + + const staleRunStep = { + id: 'step-shifted', + runId: 'response-1', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-shifted', name: 'approval_probe', args: '{}' }], + }, + }; + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: staleRunStep, + }); + + const resumeState = await manager.getResumeState(streamId); + + expect(staleRunStep.index).toBe(0); + expect(resumeState?.runSteps).toEqual([{ ...staleRunStep, index: 1 }]); + }); + + test('realigns the persisted OAuth start replay with its normalized run-step index', async () => { + manager = createInMemoryManager(); + const streamId = `oauth-index-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + manager.setContentParts(streamId, [ + { type: 'text', text: 'Prepended content' }, + { + type: 'tool_call', + tool_call: { + id: 'call-oauth-shifted', + name: 'oauth_mcp_Google-Workspace', + args: '', + }, + }, + ]); + const staleReplayEvent = { + event: 'on_run_step', + data: { + id: 'step-oauth-shifted', + runId: 'USE_PRELIM_RESPONSE_MESSAGE_ID', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [ + { + id: 'call-oauth-shifted', + name: 'oauth_mcp_Google-Workspace', + args: '', + }, + ], + }, + }, + } satisfies ServerSentEvent; + + await manager.emitChunk(streamId, staleReplayEvent); + const resumeState = await manager.getResumeState(streamId); + + expect(staleReplayEvent.data.index).toBe(0); + expect(resumeState?.runSteps[0]?.index).toBe(1); + expect(resumeState?.replayEvents).toEqual([ + { + ...staleReplayEvent, + data: { ...staleReplayEvent.data, index: 1 }, + }, + ]); + }); + + test('merges persisted and buffered run steps, preferring the buffered version by id', async () => { + const store = new PartialRunStepJobStore({ ttlAfterComplete: 60000 }); + manager = createManagerWithStore(store); + const streamId = `run-step-merge-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + + const persistedStep = { + id: 'step-persisted', + runId: 'response-1', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-persisted', name: 'approval_probe', args: '{}' }], + }, + } as Agents.RunStep; + const updatedPersistedStep = { ...persistedStep, index: 3 }; + const bufferedOnlyStep = { + ...persistedStep, + id: 'step-buffered', + index: 4, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-buffered', name: 'approval_probe', args: '{}' }], + }, + } as Agents.RunStep; + store.setPersistedRunSteps([persistedStep]); + + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: updatedPersistedStep, + }); + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: bufferedOnlyStep, + }); + + const resumeState = await manager.getResumeState(streamId); + + expect(resumeState?.runSteps).toEqual([updatedPersistedStep, bufferedOnlyStep]); + }); + + test('does not carry live content or run steps into a replacement job with the same stream id', async () => { + const store = new InMemoryJobStore({ ttlAfterComplete: 60000 }); + manager = createManagerWithStore(store); + const streamId = `run-step-replacement-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + const oldRunStep = { + id: 'step-old-job', + runId: 'response-old', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-old-job', name: 'approval_probe', args: '{}' }], + }, + } as Agents.RunStep; + await manager.emitChunk(streamId, { + event: 'on_run_step', + data: oldRunStep, + }); + store.setGraph(streamId, { + contentData: [oldRunStep], + } as unknown as StandardGraph); + store.setContentParts(streamId, [{ type: 'text', text: 'old content' }]); + store.setCollectedUsage(streamId, [{ input_tokens: 1, output_tokens: 2 }]); + + await manager.createJob(streamId, 'user-1', streamId); + + const resumeState = await manager.getResumeState(streamId); + expect(resumeState?.runSteps).toEqual([]); + expect(resumeState?.aggregatedContent).toEqual([]); + expect(store.getCollectedUsage(streamId)).toEqual([]); + }); + + test('rejects a delayed predecessor run step after the stream id is replaced', async () => { + manager = createInMemoryManager(); + const streamId = `run-step-delayed-predecessor-${Date.now()}`; + const predecessor = await manager.createJob(streamId, 'user-1', streamId); + const replacement = await manager.createJob(streamId, 'user-1', streamId); + const predecessorRunStep = { + id: 'step-predecessor', + runId: 'response-predecessor', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-predecessor', name: 'approval_probe', args: '{}' }], + }, + }; + const replacementRunStep = { + id: 'step-replacement', + runId: 'response-replacement', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-replacement', name: 'approval_probe', args: '{}' }], + }, + }; + + await manager.emitChunk( + streamId, + { event: 'on_run_step', data: predecessorRunStep }, + { expectedCreatedAt: predecessor.createdAt }, + ); + await manager.emitChunk( + streamId, + { event: 'on_run_step', data: replacementRunStep }, + { expectedCreatedAt: replacement.createdAt }, + ); + + const resumeState = await manager.getResumeState(streamId); + expect(resumeState?.runSteps).toEqual([replacementRunStep]); + }); + test('replaces OAuth replay event for the same step id', async () => { manager = createInMemoryManager(); const streamId = `oauth-delta-replace-${Date.now()}`; diff --git a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts index e48fb3bf34..3019268d51 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -682,6 +682,142 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('createJob clears persisted and live content state when a stream id is reused', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `stale-run-steps-${Date.now()}`; + await store.createJob(streamId, 'user-1', streamId); + const oldRunStep = { + id: 'step-old-job', + runId: 'response-old', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-old-job', name: 'approval_probe', args: '{}' }], + }, + } as Agents.RunStep; + await store.saveRunSteps(streamId, [oldRunStep]); + await store.appendChunk(streamId, { + event: 'on_run_step', + data: { + id: 'old-message-step', + runId: 'old-run', + index: 0, + stepDetails: { type: 'message_creation' }, + }, + }); + await store.appendChunk(streamId, { + event: 'on_message_delta', + data: { + id: 'old-message-step', + delta: { content: { type: 'text', text: 'old durable content' } }, + }, + }); + await store.updateJob(streamId, { + completedAt: Date.now(), + error: 'old error', + userMessage: { + messageId: 'old-user-message', + text: 'old user message', + }, + responseMessageId: 'old-response-message', + discoveredTools: ['old_tool'], + createdEventEmitted: true, + sender: 'Old sender', + finalEvent: '{"event":"old-final"}', + titleEvent: '{"event":"old-title"}', + replayEvents: '[{"event":"old-replay"}]', + contextUsage: '{"usedTokens":10}', + tokenUsage: '[{"input_tokens":1,"output_tokens":2}]', + endpoint: 'old-endpoint', + iconURL: 'https://example.com/old.png', + model: 'old-model', + promptTokens: 10, + agent_id: 'old-agent', + isTemporary: true, + }); + store.setGraph(streamId, { + getContentParts: () => [{ type: 'text', text: 'old graph content' }], + getRunSteps: () => [oldRunStep], + } as unknown as StandardGraph); + store.setContentParts(streamId, [{ type: 'text', text: 'old host content' }]); + store.setCollectedUsage(streamId, [{ input_tokens: 1, output_tokens: 2 }]); + expect(await store.getRunSteps(streamId)).toHaveLength(1); + + await store.createJob(streamId, 'user-1', streamId); + + expect(await store.getRunSteps(streamId)).toEqual([]); + expect(await store.getContentParts(streamId)).toBeNull(); + expect(store.getCollectedUsage(streamId)).toEqual([]); + expect(await store.getJob(streamId)).toEqual( + expect.objectContaining({ + streamId, + userId: 'user-1', + conversationId: streamId, + status: 'running', + syncSent: false, + }), + ); + expect(await store.getJob(streamId)).not.toEqual( + expect.objectContaining({ + responseMessageId: 'old-response-message', + }), + ); + await store.destroy(); + }); + + test('createJob preserves the prior live content state when replacement persistence fails', async () => { + if (!ioredisClient) { + return; + } + + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const store = new RedisJobStore(ioredisClient); + await store.initialize(); + + const streamId = `failed-replacement-${Date.now()}`; + const oldRunStep = { + id: 'step-old-job', + runId: 'response-old', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-old-job', name: 'approval_probe', args: '{}' }], + }, + } as Agents.RunStep; + await store.createJob(streamId, 'user-1', streamId); + store.setGraph(streamId, { + getContentParts: () => [{ type: 'text', text: 'old graph content' }], + getRunSteps: () => [oldRunStep], + } as unknown as StandardGraph); + store.setContentParts(streamId, [{ type: 'text', text: 'old host content' }]); + store.setCollectedUsage(streamId, [{ input_tokens: 1, output_tokens: 2 }]); + + const evalSpy = jest + .spyOn(ioredisClient, 'eval') + .mockRejectedValueOnce(new Error('replacement write failed')); + try { + await expect(store.createJob(streamId, 'user-1', streamId)).rejects.toThrow( + 'replacement write failed', + ); + } finally { + evalSpy.mockRestore(); + } + + expect(await store.getRunSteps(streamId)).toEqual([oldRunStep]); + expect(await store.getContentParts(streamId)).toEqual({ + content: [{ type: 'text', text: 'old host content' }], + }); + expect(store.getCollectedUsage(streamId)).toEqual([{ input_tokens: 1, output_tokens: 2 }]); + await store.destroy(); + }); + test('should not drop paused jobs from user tracking when cleanup sees a stale running index', async () => { if (!ioredisClient) { return;