feat: Enhance Stream Management with Abort Functionality

- Updated the abort endpoint to support aborting ongoing generation streams using either streamId or conversationId.
- Introduced a new mutation hook `useAbortStreamMutation` for client-side integration.
- Added `useStreamStatus` query to monitor stream status and facilitate resuming conversations.
- Enhanced `useChatHelpers` to incorporate abort functionality when stopping generation.
- Improved `useResumableSSE` to handle stream errors and token refresh seamlessly.
- Updated `useResumeOnLoad` to check for active streams and resume conversations appropriately.
This commit is contained in:
Danny Avila 2025-12-11 21:19:43 -05:00
parent 0a517a2b8f
commit 1985e53b80
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
11 changed files with 295 additions and 136 deletions

View file

@ -1,10 +1,10 @@
import { useCallback, useState } from 'react';
import { QueryKeys } from 'librechat-data-provider';
import { QueryKeys, isAssistantsEndpoint } from 'librechat-data-provider';
import { useQueryClient } from '@tanstack/react-query';
import { useRecoilState, useResetRecoilState, useSetRecoilState } from 'recoil';
import type { TMessage } from 'librechat-data-provider';
import { useAbortStreamMutation, useGetMessagesByConvoId } from '~/data-provider';
import useChatFunctions from '~/hooks/Chat/useChatFunctions';
import { useGetMessagesByConvoId } from '~/data-provider';
import { useAuthContext } from '~/hooks/AuthContext';
import useNewConvo from '~/hooks/useNewConvo';
import store from '~/store';
@ -17,11 +17,12 @@ export default function useChatHelpers(index = 0, paramId?: string) {
const queryClient = useQueryClient();
const { isAuthenticated } = useAuthContext();
const abortMutation = useAbortStreamMutation();
const { newConversation } = useNewConvo(index);
const { useCreateConversationAtom } = store;
const { conversation, setConversation } = useCreateConversationAtom(index);
const { conversationId } = conversation ?? {};
const { conversationId, endpoint, endpointType } = conversation ?? {};
const queryParam = paramId === 'new' ? paramId : (conversationId ?? paramId ?? '');
@ -107,7 +108,43 @@ export default function useChatHelpers(index = 0, paramId?: string) {
}
};
const stopGenerating = () => clearAllSubmissions();
/**
* Stop generation - for non-assistants endpoints, calls abort endpoint first.
* The abort endpoint will cause the backend to emit a `done` event with `aborted: true`,
* which will be handled by the SSE event handler to clean up UI.
* Assistants endpoint has its own abort mechanism via useEventHandlers.abortConversation.
*/
const stopGenerating = useCallback(async () => {
const actualEndpoint = endpointType ?? endpoint;
const isAssistants = isAssistantsEndpoint(actualEndpoint);
console.log('[useChatHelpers] stopGenerating called', {
conversationId,
endpoint,
endpointType,
actualEndpoint,
isAssistants,
});
// For non-assistants endpoints (using resumable streams), call abort endpoint first
if (conversationId && !isAssistants) {
try {
console.log('[useChatHelpers] Calling abort mutation for:', conversationId);
await abortMutation.mutateAsync({ conversationId });
console.log('[useChatHelpers] Abort mutation succeeded');
// The SSE will receive a `done` event with `aborted: true` and clean up
// We still clear submissions as a fallback
clearAllSubmissions();
} catch (error) {
console.error('[useChatHelpers] Abort failed:', error);
// Fall back to clearing submissions
clearAllSubmissions();
}
} else {
// For assistants endpoints, just clear submissions (existing behavior)
console.log('[useChatHelpers] Assistants endpoint, just clearing submissions');
clearAllSubmissions();
}
}, [conversationId, endpoint, endpointType, abortMutation, clearAllSubmissions]);
const handleStopGenerating = (e: React.MouseEvent<HTMLButtonElement>) => {
e.preventDefault();