diff --git a/api/server/routes/agents/__tests__/abort.spec.js b/api/server/routes/agents/__tests__/abort.spec.js index 15efb19a39..c9cc6b5b0f 100644 --- a/api/server/routes/agents/__tests__/abort.spec.js +++ b/api/server/routes/agents/__tests__/abort.spec.js @@ -121,7 +121,10 @@ describe('Agent Abort Endpoint', () => { expect(response.status).toBe(200); expect(response.body).toEqual({ success: true, aborted: jobStreamId }); - expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith(jobStreamId); + expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith( + jobStreamId, + expect.objectContaining({ transformAbortContent: expect.any(Function) }), + ); }); it('should allow abort when job has no userId metadata (backwards compatibility)', async () => { @@ -327,6 +330,56 @@ describe('Agent Abort Endpoint', () => { ); }); + it('stamps a paused ask_user_question via transformAbortContent, before the final SSE emits', async () => { + const jobStreamId = 'test-stream-123'; + const question = { question: 'Deploy where?', options: [{ label: 'Prod', value: 'prod' }] }; + + mockGenerationJobManager.getJob.mockResolvedValue({ + metadata: { + userId: 'test-user-123', + pendingAction: { payload: { type: 'ask_user_question', question } }, + }, + }); + + // abortJob applies the transform; capture it and echo the transformed + // content back as the result, mirroring the real (Redis) reconstruction + // where the ask tool_call arrives with empty args. + let capturedTransform; + mockGenerationJobManager.abortJob.mockImplementation(async (_streamId, options) => { + capturedTransform = options?.transformAbortContent; + const rawContent = [ + { type: 'tool_call', tool_call: { id: 'tc1', name: 'ask_user_question', args: '' } }, + ]; + const content = capturedTransform ? capturedTransform(rawContent) : rawContent; + return { + success: true, + jobData: { + userMessage: { messageId: 'user-msg-123' }, + responseMessageId: 'response-msg-456', + conversationId: jobStreamId, + }, + content, + text: '', + }; + }); + + mockSaveMessage.mockResolvedValue(); + + const response = await request(app) + .post('/api/agents/chat/abort') + .send({ conversationId: jobStreamId }); + + expect(response.status).toBe(200); + expect(capturedTransform).toEqual(expect.any(Function)); + // The saved (and, in prod, emitted) content carries the stamped args. + // saveMessage(reqLike, responseMessage, opts) — the message is arg #2. + const savedMessage = mockSaveMessage.mock.calls[0][1]; + const askPart = savedMessage.content.find( + (p) => p?.tool_call?.name === 'ask_user_question', + ); + expect(JSON.parse(askPart.tool_call.args)).toMatchObject({ question: 'Deploy where?' }); + }); + it('should handle saveMessage errors gracefully', async () => { const jobStreamId = 'test-stream-123'; @@ -388,7 +441,10 @@ describe('Agent Abort Endpoint', () => { expect(response.status).toBe(200); expect(response.body).toEqual({ success: true, aborted: 'running-stream' }); - expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith('running-stream'); + expect(mockGenerationJobManager.abortJob).toHaveBeenCalledWith( + 'running-stream', + expect.objectContaining({ transformAbortContent: expect.any(Function) }), + ); }); it('should not abort paused fallback jobs', async () => { diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index a4221a3c37..81612197e8 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -289,7 +289,18 @@ router.post('/chat/abort', configMiddleware, async (req, res) => { } logger.debug(`[AgentStream] Job found, aborting: ${jobStreamId}`); - const abortResult = await GenerationJobManager.abortJob(jobStreamId); + // Re-attach a paused ask_user_question's args to the abort content BEFORE + // abortJob emits the final SSE. Redis reconstructs abort content from the + // chunk log, which never saw the pause-time stamp applied to the in-process + // contentParts — stamping inside abortJob (not after) means the LIVE client + // gets the question too, not just the saved message on reload. + const abortedAskPayload = job.metadata?.pendingAction?.payload; + const abortResult = await GenerationJobManager.abortJob(jobStreamId, { + transformAbortContent: (content) => + abortedAskPayload?.type === 'ask_user_question' && Array.isArray(content) + ? attachAskUserQuestionArgs(content, abortedAskPayload.question) + : content, + }); logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`, { abortResultSuccess: abortResult.success, abortResultUserMessageId: abortResult.jobData?.userMessage?.messageId, @@ -320,14 +331,10 @@ router.post('/chat/abort', configMiddleware, async (req, res) => { hasPersistableAbortContent(abortResult.content) ) { const { jobData, text } = abortResult; - let { content } = abortResult; - // Redis reconstructs abort content from the graph/chunk log, which never saw - // the pause-time args stamp applied to the in-process contentParts — re-stamp - // here so a question abandoned via Stop persists with its question intact. - const abortedPayload = job.metadata?.pendingAction?.payload; - if (abortedPayload?.type === 'ask_user_question' && Array.isArray(content)) { - content = attachAskUserQuestionArgs(content, abortedPayload.question); - } + // `abortResult.content` is already stamped by `transformAbortContent` + // above (same content the final SSE carried), so the saved message and + // the live client agree. + const { content } = abortResult; const responseMessage = { messageId: jobData.responseMessageId, parentMessageId: jobData.userMessage.messageId, diff --git a/client/src/components/Chat/Messages/Content/ApprovalContext.tsx b/client/src/components/Chat/Messages/Content/ApprovalContext.tsx index f3714ece52..cfc301f2bb 100644 --- a/client/src/components/Chat/Messages/Content/ApprovalContext.tsx +++ b/client/src/components/Chat/Messages/Content/ApprovalContext.tsx @@ -1,4 +1,5 @@ import { createContext, useCallback, useContext, useMemo, useRef, useState } from 'react'; +import { atom, useRecoilState } from 'recoil'; import { Constants } from 'librechat-data-provider'; import type { Agents } from 'librechat-data-provider'; import { @@ -14,6 +15,40 @@ import { useGetEphemeralAgent } from '~/store/agents'; * controls and explain a terminal outcome. */ type ActionStatus = 'idle' | 'submitting' | 'submitted' | 'expired' | 'error'; +/** + * Ask-answer submit status keyed by `actionId`. This lives in Recoil, NOT the + * {@link ApprovalContext} React state, because the PRIMARY answer surface is + * the composer in `ChatForm` — which renders OUTSIDE `ApprovalProvider` (that + * only wraps message content). A React context read there degrades to the + * inert {@link FALLBACK}, so an in-flight lock / expired status tracked in the + * context would never engage for the composer. A global atom is visible to + * both the composer (`useAskAnswerMode`) and the message-content card + * (`AskUserQuestion`), so a fast double-submit is actually blocked and a + * terminal (expired/error) state surfaces on either surface. + */ +const askSubmitStatusAtom = atom>({ + key: 'askAnswerSubmitStatus', + default: {}, +}); + +/** Shared read/write for the ask-answer submit status. */ +export function useAskSubmitStatus(): { + getAskStatus: (actionId: string) => ActionStatus; + setAskStatus: (actionId: string, status: ActionStatus) => void; +} { + const [statusMap, setStatusMap] = useRecoilState(askSubmitStatusAtom); + const getAskStatus = useCallback( + (actionId: string): ActionStatus => statusMap[actionId] ?? 'idle', + [statusMap], + ); + const setAskStatus = useCallback( + (actionId: string, status: ActionStatus) => + setStatusMap((prev) => ({ ...prev, [actionId]: status })), + [setStatusMap], + ); + return { getAskStatus, setAskStatus }; +} + interface ApprovalContextValue { /** Record (or clear) a card's decision for its tool_call within an action. */ setDecision: ( @@ -233,6 +268,9 @@ export function useResumeSubmit() { const approvalMutation = useSubmitToolApprovalMutation(); const askMutation = useSubmitAskAnswerMutation(); const { getDecisions, isReady, setStatus } = useApprovalContext(); + /** Ask status lives in Recoil so it works from the composer (outside the + * provider); tool-approval status stays on the context. */ + const { setAskStatus } = useAskSubmitStatus(); const buildResumeFields = useCallback((): ResumeAgentFields | null => { const conversationId = conversation?.conversationId; @@ -278,12 +316,12 @@ export function useResumeSubmit() { if (!fields || answer.length === 0) { return; } - setStatus(actionId, 'submitting'); + setAskStatus(actionId, 'submitting'); askMutation.mutate( { ...fields, actionId, answer }, { onSuccess: () => { - setStatus(actionId, 'submitted'); + setAskStatus(actionId, 'submitted'); /** * Drop the synthetic question part now that the run is resuming: the * server streams the resumed segment at ABSOLUTE content indices @@ -314,11 +352,11 @@ export function useResumeSubmit() { */ opts?.onSuccess?.(); }, - onError: (error) => setStatus(actionId, isExpiredError(error) ? 'expired' : 'error'), + onError: (error) => setAskStatus(actionId, isExpiredError(error) ? 'expired' : 'error'), }, ); }, - [askMutation, buildResumeFields, setStatus, chatContext], + [askMutation, buildResumeFields, setAskStatus, chatContext], ); return { submitToolApproval, submitAskAnswer }; diff --git a/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx b/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx index b1b8e468f6..f1d3794d93 100644 --- a/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx +++ b/client/src/components/Chat/Messages/Content/AskUserQuestion.tsx @@ -2,7 +2,7 @@ import { useContext, useMemo, useState } from 'react'; import { ChevronUp, TriangleAlert } from 'lucide-react'; import { Button, TextareaAutosize } from '@librechat/client'; import type { Agents } from 'librechat-data-provider'; -import { useApprovalContext, useResumeSubmit } from './ApprovalContext'; +import { useAskSubmitStatus, useResumeSubmit } from './ApprovalContext'; import useAskAnswerMode from '~/hooks/Input/useAskAnswerMode'; import { ChatContext } from '~/Providers/ChatContext'; import { splitOtherOption } from '~/utils/approval'; @@ -25,7 +25,7 @@ export default function AskUserQuestion({ question: Agents.AskUserQuestionRequest; }) { const localize = useLocalize(); - const { getStatus } = useApprovalContext(); + const { getAskStatus } = useAskSubmitStatus(); const { submitAskAnswer } = useResumeSubmit(); const [answer, setAnswer] = useState(''); const [localChecked, setLocalChecked] = useState([]); @@ -48,7 +48,7 @@ export default function AskUserQuestion({ [question.options], ); - const status = getStatus(actionId); + const status = getAskStatus(actionId); if (popoverVisible && isLivePause) { return null; } diff --git a/client/src/hooks/Input/useAskAnswerMode.ts b/client/src/hooks/Input/useAskAnswerMode.ts index 75e0f294d2..43aead74e8 100644 --- a/client/src/hooks/Input/useAskAnswerMode.ts +++ b/client/src/hooks/Input/useAskAnswerMode.ts @@ -1,7 +1,7 @@ import { useCallback, useEffect, useMemo } from 'react'; import { atom, useRecoilState, useRecoilValue } from 'recoil'; import { - useApprovalContext, + useAskSubmitStatus, useResumeSubmit, } from '~/components/Chat/Messages/Content/ApprovalContext'; import { @@ -74,14 +74,16 @@ export default function useAskAnswerMode(conversationId?: string | null) { const [checked, setChecked] = useRecoilState(askAnswerCheckedAtom); const saveDrafts = useRecoilValue(store.saveDrafts); const { submitAskAnswer } = useResumeSubmit(); - const { getStatus } = useApprovalContext(); + /** Recoil-backed so the lock/status works from the composer, which renders + * outside `ApprovalProvider` (where the context status would be inert). */ + const { getAskStatus } = useAskSubmitStatus(); /** Absent outside ChatView (Share/search render the answer card without the * composer form) — resets are simply skipped there. */ const resetComposer = useOptionalChatFormContext()?.reset; /** The answer is in flight (or terminal): every submit path must become a * no-op so a double-click or a stray Skip can't race a second resume. */ - const status = liveAsk != null ? getStatus(liveAsk.actionId) : 'idle'; + const status = liveAsk != null ? getAskStatus(liveAsk.actionId) : 'idle'; const locked = status === 'submitting' || status === 'submitted' || status === 'expired'; const dismissed = liveAsk != null && dismissedIds.includes(liveAsk.actionId); /** diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index d6d34d8304..d8005e8c6b 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -763,8 +763,19 @@ class GenerationJobManagerClass { * Cross-replica support (Redis mode): * - Emits abort signal via Redis pub/sub * - The replica running generation receives signal and aborts its AbortController + * + * `options.transformAbortContent` rewrites the persistable content BEFORE the + * final SSE is emitted (and before it is returned for the DB save), so a + * host-side stamp — e.g. re-attaching a paused `ask_user_question`'s args + * that the Redis chunk-log reconstruction dropped — reaches the LIVE client + * too, not just the saved message. Pure/optional; identity when omitted. */ - async abortJob(streamId: string): Promise { + async abortJob( + streamId: string, + options?: { + transformAbortContent?: (content: TMessageContentParts[]) => TMessageContentParts[]; + }, + ): Promise { const jobData = await this.jobStore.getJob(streamId); const runtime = this.runtimeState.get(streamId); @@ -795,7 +806,12 @@ class GenerationJobManagerClass { /** Content before clearing state */ const result = await this.jobStore.getContentParts(streamId); const content = result?.content ?? []; - const abortContent = filterPersistableAbortContent(content); + let abortContent = filterPersistableAbortContent(content); + if (options?.transformAbortContent) { + abortContent = options.transformAbortContent( + abortContent as TMessageContentParts[], + ) as typeof abortContent; + } const shouldPersistAbortContent = abortContent.length > 0; /** Collected usage for all models */