From 5f5d5d4747e6278adec52dcffc304afeba1d8ce6 Mon Sep 17 00:00:00 2001 From: Marco Beretta <81851188+berry-13@users.noreply.github.com> Date: Sun, 26 Jul 2026 15:21:09 +0200 Subject: [PATCH] fix: handle steer acks that arrive after the run ends --- .../Parts/__tests__/PendingSteers.test.tsx | 25 ++- .../Chat/__tests__/useSteerRecovery.spec.tsx | 165 ++++++++++++++++++ client/src/hooks/Chat/index.ts | 1 - client/src/hooks/Chat/useSteerRecovery.ts | 120 ++++++++----- client/src/hooks/Chat/useSteering.ts | 52 +++++- client/src/hooks/SSE/useResumableSSE.ts | 13 +- 6 files changed, 327 insertions(+), 49 deletions(-) create mode 100644 client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx diff --git a/client/src/components/Chat/Messages/Content/Parts/__tests__/PendingSteers.test.tsx b/client/src/components/Chat/Messages/Content/Parts/__tests__/PendingSteers.test.tsx index 661d1f4ff9..32bda0e51f 100644 --- a/client/src/components/Chat/Messages/Content/Parts/__tests__/PendingSteers.test.tsx +++ b/client/src/components/Chat/Messages/Content/Parts/__tests__/PendingSteers.test.tsx @@ -1,17 +1,20 @@ import React from 'react'; import { RecoilRoot } from 'recoil'; -import { render, screen } from '@testing-library/react'; +import { render, screen, fireEvent } from '@testing-library/react'; import type { PendingSteer } from '~/store/families'; import PendingSteers from '../PendingSteers'; import store from '~/store'; +const mockRetry = jest.fn(); +const mockSendAsNew = jest.fn(); + jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key, })); jest.mock('~/hooks/Chat/useSteerRecovery', () => ({ __esModule: true, - default: () => ({ retry: jest.fn(), sendAsNew: jest.fn(), remove: jest.fn() }), + default: () => ({ retry: mockRetry, sendAsNew: mockSendAsNew }), })); jest.mock('../SteerPart', () => ({ @@ -38,6 +41,10 @@ function renderPending(steers: PendingSteer[]) { } describe('PendingSteers', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + it('renders nothing with no pending steers', () => { const { container } = renderPending([]); expect(container).toBeEmptyDOMElement(); @@ -55,4 +62,18 @@ describe('PendingSteers', () => { expect(screen.getByText('com_ui_retry')).toBeInTheDocument(); expect(screen.getByText('com_ui_send_as_new')).toBeInTheDocument(); }); + + it('retries the failed steer by id', () => { + renderPending([pending({ status: 'failed', steerId: 's-failed' })]); + fireEvent.click(screen.getByText('com_ui_retry')); + expect(mockRetry).toHaveBeenCalledWith('s-failed'); + expect(mockSendAsNew).not.toHaveBeenCalled(); + }); + + it('sends the failed steer as new by id', () => { + renderPending([pending({ status: 'failed', steerId: 's-failed' })]); + fireEvent.click(screen.getByText('com_ui_send_as_new')); + expect(mockSendAsNew).toHaveBeenCalledWith('s-failed'); + expect(mockRetry).not.toHaveBeenCalled(); + }); }); diff --git a/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx b/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx new file mode 100644 index 0000000000..82a7ac2e97 --- /dev/null +++ b/client/src/hooks/Chat/__tests__/useSteerRecovery.spec.tsx @@ -0,0 +1,165 @@ +import React from 'react'; +import { act, renderHook } from '@testing-library/react'; +import { RecoilRoot, useRecoilValue, type MutableSnapshot } from 'recoil'; +import useSteerRecovery from '../useSteerRecovery'; +import store from '~/store'; + +const mockMutate = jest.fn(); +const mockFetchStreamStatus = jest.fn(); + +jest.mock('~/data-provider', () => ({ + useSteerMessageMutation: () => ({ mutate: mockMutate }), + fetchStreamStatus: (...args: unknown[]) => mockFetchStreamStatus(...args), +})); + +const CONVO_ID = 'convo-steer-recovery'; + +function setup(initialize?: (snapshot: MutableSnapshot) => void) { + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ); + return renderHook( + () => ({ + recovery: useSteerRecovery(CONVO_ID), + chips: useRecoilValue(store.pendingSteersByConvoId(CONVO_ID)), + queue: useRecoilValue(store.queuedMessagesByConvoId(CONVO_ID)), + applied: useRecoilValue(store.appliedSteerIdsByConvoId(CONVO_ID)), + }), + { wrapper }, + ); +} + +describe('useSteerRecovery', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('retry', () => { + it('marks the chip sending immediately', () => { + mockMutate.mockImplementation(() => undefined); + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 'local-1', text: 'redo this', status: 'failed', createdAt: 5 }, + ]); + }); + act(() => { + result.current.recovery.retry('local-1'); + }); + expect(result.current.chips).toEqual([ + expect.objectContaining({ steerId: 'local-1', status: 'sending' }), + ]); + }); + + it('swaps the local id for the server id on success, keeping it pending', () => { + mockMutate.mockImplementation((_params, { onSuccess }) => { + onSuccess({ steerId: 'srv-9', status: 'queued', position: 1, conversationId: CONVO_ID }); + }); + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 'local-1', text: 'redo this', status: 'failed', createdAt: 5 }, + ]); + }); + act(() => { + result.current.recovery.retry('local-1'); + }); + // The old local id must be gone entirely — leaving it behind is what let + // the applied SteerPart and a stale pending copy render together. + expect(result.current.chips).toEqual([ + expect.objectContaining({ steerId: 'srv-9', text: 'redo this', status: 'pending' }), + ]); + }); + + it('routes to the queue on NO_ACTIVE_RUN instead of marking it failed again', () => { + mockMutate.mockImplementation((_params, { onError }) => { + onError({ response: { data: { code: 'NO_ACTIVE_RUN' } } }); + }); + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 'local-2', text: 'too late', status: 'failed', createdAt: 7 }, + ]); + }); + act(() => { + result.current.recovery.retry('local-2'); + }); + expect(result.current.chips).toEqual([]); + expect(result.current.queue).toEqual([ + expect.objectContaining({ id: 'local-2', text: 'too late' }), + ]); + }); + + it('also routes to the queue on RUN_PAUSED / STEER_UNSUPPORTED / STEER_QUEUE_FULL', () => { + for (const code of ['RUN_PAUSED', 'STEER_UNSUPPORTED', 'STEER_QUEUE_FULL']) { + mockMutate.mockImplementation((_params, { onError }) => { + onError({ response: { data: { code } } }); + }); + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: `local-${code}`, text: code, status: 'failed', createdAt: 1 }, + ]); + }); + act(() => { + result.current.recovery.retry(`local-${code}`); + }); + expect(result.current.queue).toEqual([expect.objectContaining({ id: `local-${code}` })]); + } + }); + + it('marks it failed again on an unrecognized error', () => { + mockMutate.mockImplementation((_params, { onError }) => { + onError(new Error('network')); + }); + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 'local-3', text: 'network flake', status: 'failed', createdAt: 1 }, + ]); + }); + act(() => { + result.current.recovery.retry('local-3'); + }); + expect(result.current.chips).toEqual([ + expect.objectContaining({ steerId: 'local-3', status: 'failed' }), + ]); + expect(result.current.queue).toEqual([]); + }); + + it('no-ops when the steer id is no longer pending', () => { + const { result } = setup(); + act(() => { + result.current.recovery.retry('missing'); + }); + expect(mockMutate).not.toHaveBeenCalled(); + }); + }); + + describe('sendAsNew', () => { + it('moves the chip to the queue, merged chronologically with existing items', () => { + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 's-old', text: 'older failed steer', status: 'failed', createdAt: 0 }, + ]); + set(store.queuedMessagesByConvoId(CONVO_ID), [ + { id: 'existing', text: 'first', createdAt: 1 }, + ]); + }); + act(() => { + result.current.recovery.sendAsNew('s-old'); + }); + expect(result.current.chips).toEqual([]); + // Chronological merge (not blind-append): the older steer sorts ahead + // of the already-queued, later item. + expect(result.current.queue.map((item) => item.id)).toEqual(['s-old', 'existing']); + }); + + it('records the id in appliedSteerIds so a late ACK cannot re-mint the chip', () => { + const { result } = setup(({ set }) => { + set(store.pendingSteersByConvoId(CONVO_ID), [ + { steerId: 's-late', text: 'send me', status: 'failed', createdAt: 1 }, + ]); + }); + act(() => { + result.current.recovery.sendAsNew('s-late'); + }); + expect(result.current.applied).toEqual(['s-late']); + }); + }); +}); diff --git a/client/src/hooks/Chat/index.ts b/client/src/hooks/Chat/index.ts index 052f14dee6..3dc2a8fd65 100644 --- a/client/src/hooks/Chat/index.ts +++ b/client/src/hooks/Chat/index.ts @@ -10,4 +10,3 @@ export { default as useQueueDrain } from './useQueueDrain'; export { default as useSteering } from './useSteering'; export { default as useSteerCancel, useSteerReclaim } from './useSteerCancel'; export { default as useSteerConvert } from './useSteerConvert'; -export { default as useSteerRecovery } from './useSteerRecovery'; diff --git a/client/src/hooks/Chat/useSteerRecovery.ts b/client/src/hooks/Chat/useSteerRecovery.ts index e7755efe2f..52ec52af6a 100644 --- a/client/src/hooks/Chat/useSteerRecovery.ts +++ b/client/src/hooks/Chat/useSteerRecovery.ts @@ -1,17 +1,34 @@ +import { useRef, useEffect, useCallback } from 'react'; import { useRecoilCallback } from 'recoil'; import type { PendingSteer } from '~/store/families'; +import { getSteerErrorCode, resolveAcknowledgedSteer } from '~/hooks/Chat/useSteering'; +import useSteerConvert from '~/hooks/Chat/useSteerConvert'; import { useSteerMessageMutation } from '~/data-provider'; +import { carriedSteerContext } from '~/utils'; import store from '~/store'; /** * Retry or re-route a pending steer from OUTSIDE the composer. The thread's - * pending block needs these two actions without dragging the whole + * pending block needs these actions without dragging the whole * `SteeringControls` object through the message tree. */ export default function useSteerRecovery(conversationId: string) { const { mutate: steerMessage } = useSteerMessageMutation(); + const convertSteersToQueued = useSteerConvert(); - const markSending = useRecoilCallback( + /** `PendingSteers` only renders while `isLast && isSubmitting` is true, so + * its unmount IS the run ending. A retry's ack can resolve afterward; + * mirrors `composerMountedRef` in `ChatForm`, which guards the equivalent + * race for a reclaimed steer whose restore resolves post-unmount. */ + const mountedRef = useRef(true); + useEffect( + () => () => { + mountedRef.current = false; + }, + [], + ); + + const markStatus = useRecoilCallback( ({ set }) => (steerId: string, status: PendingSteer['status']) => { set(store.pendingSteersByConvoId(conversationId), (prev) => @@ -21,66 +38,91 @@ export default function useSteerRecovery(conversationId: string) { [conversationId], ); + const acknowledgeRetry = useRecoilCallback( + (cbInterface) => (localId: string, steer: PendingSteer, runOver: boolean) => { + resolveAcknowledgedSteer(cbInterface, conversationId, localId, steer, runOver); + }, + [conversationId], + ); + + /** Routes a steer straight into the queue: reused for a retry that degrades + * (no active run / paused / unsupported / queue full) and for `sendAsNew`, + * so both get the same id-dedup, chronological merge, and applied-id + * bookkeeping as every other steer-to-queue conversion — instead of a + * blind append that a late ACK could still re-mint a chip for. */ + const queueSteer = useCallback( + (steer: PendingSteer) => { + convertSteersToQueued(conversationId, [ + { + steerId: steer.steerId, + text: steer.text, + createdAt: steer.createdAt, + ...(steer.files && steer.files.length > 0 && { files: steer.files }), + ...carriedSteerContext(steer), + }, + ]); + }, + [conversationId, convertSteersToQueued], + ); + const retry = useRecoilCallback( ({ snapshot }) => (steerId: string) => { - const steers = snapshot + const steer = snapshot .getLoadable(store.pendingSteersByConvoId(conversationId)) - .getValue(); - const steer = steers.find((item) => item.steerId === steerId); + .getValue() + .find((item) => item.steerId === steerId); if (!steer) { return; } - markSending(steerId, 'sending'); + markStatus(steerId, 'sending'); steerMessage( { conversationId, text: steer.text, files: steer.files }, { - onSuccess: () => markSending(steerId, 'pending'), - onError: () => markSending(steerId, 'failed'), + onSuccess: (response) => { + acknowledgeRetry( + steerId, + { ...steer, steerId: response.steerId, status: 'pending' }, + !mountedRef.current, + ); + }, + onError: (error) => { + const code = getSteerErrorCode(error); + // The run ended, is paused, or can't accept a steer right now — + // none of that means the words are lost, just that a queued + // follow-up is the only way left to send them. + if ( + code === 'NO_ACTIVE_RUN' || + code === 'RUN_PAUSED' || + code === 'STEER_UNSUPPORTED' || + code === 'STEER_QUEUE_FULL' + ) { + queueSteer(steer); + return; + } + markStatus(steerId, 'failed'); + }, }, ); }, - [conversationId, steerMessage, markSending], + [conversationId, steerMessage, markStatus, acknowledgeRetry, queueSteer], ); /** Move a failed steer into the queue: it sends when the reply finishes. */ const sendAsNew = useRecoilCallback( - ({ snapshot, set }) => + ({ snapshot }) => (steerId: string) => { - const steers = snapshot + const steer = snapshot .getLoadable(store.pendingSteersByConvoId(conversationId)) - .getValue(); - const steer = steers.find((item) => item.steerId === steerId); + .getValue() + .find((item) => item.steerId === steerId); if (!steer) { return; } - set(store.pendingSteersByConvoId(conversationId), (prev) => - prev.filter((item) => item.steerId !== steerId), - ); - set(store.queuedMessagesByConvoId(conversationId), (prev) => [ - ...prev, - { - id: steer.steerId, - text: steer.text, - createdAt: steer.createdAt, - files: steer.files, - quotes: steer.quotes, - manualSkills: steer.manualSkills, - }, - ]); + queueSteer(steer); }, - [conversationId], + [conversationId, queueSteer], ); - const remove = useRecoilCallback( - ({ set }) => - (steerId: string) => { - set(store.pendingSteersByConvoId(conversationId), (prev) => - prev.filter((item) => item.steerId !== steerId), - ); - }, - [conversationId], - ); - - return { retry, sendAsNew, remove }; + return { retry, sendAsNew }; } diff --git a/client/src/hooks/Chat/useSteering.ts b/client/src/hooks/Chat/useSteering.ts index c348a7376f..41b397e814 100644 --- a/client/src/hooks/Chat/useSteering.ts +++ b/client/src/hooks/Chat/useSteering.ts @@ -4,6 +4,7 @@ import { useToastContext } from '@librechat/client'; import { useRecoilValue, useSetRecoilState, useRecoilCallback } from 'recoil'; import { Constants, ContentTypes, isAssistantsEndpoint } from 'librechat-data-provider'; import type { TMessage, TConversation, TMessageContentParts } from 'librechat-data-provider'; +import type { CallbackInterface } from 'recoil'; import type { RunEnd, PendingSteer, QueuedMessage, QueuedMessageOrigin } from '~/store/families'; import type { GenerationProtocolVersion } from '~/data-provider'; import type { ExtendedFile, FileSetter } from '~/common'; @@ -41,7 +42,7 @@ export interface QueuedMessageContext { /** Server-side cap on a usage touch (mirrors `FILES_USAGE_MAX_IDS`). */ const QUEUE_USAGE_MAX_FILES = 10; -type SteerErrorCode = +export type SteerErrorCode = | 'NO_ACTIVE_RUN' | 'RUN_PAUSED' | 'RUN_REPLACED' @@ -65,7 +66,7 @@ type SubmitSteerOptions = { generationProtocolVersion?: GenerationProtocolVersion; }; -function getSteerErrorCode(error: unknown): SteerErrorCode | undefined { +export function getSteerErrorCode(error: unknown): SteerErrorCode | undefined { const response = (error as { response?: { data?: { code?: string } } } | undefined)?.response; return response?.data?.code; } @@ -93,6 +94,53 @@ function isSameRunEpoch(a: RunEnd | null, b: RunEnd): boolean { return a.endedAt === b.endedAt; } +/** + * Resolves a steer's 202 ACK against the applied-id set and the run's live + * state: on the common path the local chip is swapped for its server-assigned + * id (still `pending`, awaiting `on_steer_applied`). If an `on_steer_applied` + * event already beat the ACK, or the run itself ended before the ACK landed, + * no later SSE event will ever resolve a `pending` chip — so the words are + * routed straight into the queue instead of stranding one. + * + * Exported so an ACK originating OUTSIDE the composer (`useSteerRecovery`'s + * retry, for a steer that already failed once) resolves identically instead + * of approximating this logic. + */ +export function resolveAcknowledgedSteer( + { snapshot, set }: Pick, + conversationId: string, + localId: string, + steer: PendingSteer, + runOver: boolean, +): void { + const applied = snapshot.getLoadable(store.appliedSteerIdsByConvoId(conversationId)).getValue(); + const alreadyApplied = applied.includes(steer.steerId); + set(store.pendingSteersByConvoId(conversationId), (prev) => { + const next = prev.filter((item) => item.steerId !== localId); + // Upsert: an SSE reconnect may have reseeded the chip under the server id + // already — appending again would duplicate it. + const alreadySeeded = next.some((item) => item.steerId === steer.steerId); + return alreadyApplied || runOver || alreadySeeded ? next : [...next, steer]; + }); + if (alreadyApplied || !runOver) { + return; + } + set(store.queuedMessagesByConvoId(conversationId), (prev) => + prev.some((queued) => queued.id === steer.steerId) + ? prev + : [ + ...prev, + { + id: steer.steerId, + text: steer.text, + createdAt: steer.createdAt, + ...(steer.files && steer.files.length > 0 && { files: steer.files }), + ...carriedSteerContext(steer), + }, + ], + ); +} + /** True when the latest assistant message carries an unresolved tool approval — * the run is (or is about to be) paused, so a steer POST would 409. */ function hasLiveToolApproval(messages: TMessage[] | undefined): boolean { diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 6f29a92599..1c46de44ae 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -939,10 +939,13 @@ export default function useResumableSSE( const convertSteersToQueued = useSteerConvert(); /** Error events carry no `pendingSteers` payload (the server drops its copy - * on failure), but every acknowledged chip's text is local — convert them - * to queued follow-ups so the user's words survive a failed run. `sending` - * chips settle through their own POST callbacks (404 falls back to - * queue/send) and `failed` chips keep their manual controls. */ + * on failure), but every acknowledged OR failed chip's text is local — + * convert both to queued follow-ups so the user's words survive a failed + * run. `sending` chips settle through their own POST callbacks (404 falls + * back to queue/send); `pending` and `failed` chips have no such callback + * waiting on them, so leaving either behind would strand it — and worse, + * `PendingSteers` reads this same atom on the NEXT run's reply, so a + * stranded chip would leak into a message it was never part of. */ const convertLocalSteersToQueued = useRecoilCallback( ({ snapshot }) => ( @@ -958,7 +961,7 @@ export default function useResumableSSE( const settled = chips .filter( (steer) => - steer.status === 'pending' && + (steer.status === 'pending' || steer.status === 'failed') && !excluded.has(steer.steerId) && (steer.clientSteerId == null || !excluded.has(steer.clientSteerId)), )