refactor: Update active jobs query integration for optimistic updates on abort

- Introduced a new interface for active jobs response to standardize data handling.
- Updated query keys for active jobs to ensure consistency across components.
- Enhanced job management logic in hooks to properly reflect active job states, improving overall application responsiveness.
This commit is contained in:
Danny Avila 2025-12-18 09:33:00 -05:00
parent be3b6d21d7
commit 6aebccd20a
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
3 changed files with 18 additions and 13 deletions

View file

@ -31,9 +31,13 @@ export function useStreamStatus(conversationId: string | undefined, enabled = tr
});
}
export const activeJobsQueryKey = [QueryKeys.activeJobs] as const;
export const genTitleQueryKey = (conversationId: string) => ['genTitle', conversationId] as const;
/** Response type for active jobs query */
export interface ActiveJobsResponse {
activeJobIds: string[];
}
// Module-level queue for title generation (survives re-renders)
// Stores conversationIds that need title generation once their job completes
const titleQueue = new Set<string>();
@ -123,7 +127,7 @@ export function useTitleGeneration(enabled = true) {
*/
export function useActiveJobs(enabled = true) {
return useQuery({
queryKey: activeJobsQueryKey,
queryKey: [QueryKeys.activeJobs],
queryFn: () => dataService.getActiveJobs(),
enabled,
staleTime: 5_000,

View file

@ -3,7 +3,8 @@ 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 type { ActiveJobsResponse } from '~/data-provider';
import { useGetMessagesByConvoId, useAbortStreamMutation } from '~/data-provider';
import useChatFunctions from '~/hooks/Chat/useChatFunctions';
import { useAuthContext } from '~/hooks/AuthContext';
import useNewConvo from '~/hooks/useNewConvo';
@ -129,6 +130,10 @@ export default function useChatHelpers(index = 0, paramId?: string) {
// For non-assistants endpoints (using resumable streams), call abort endpoint first
if (conversationId && !isAssistants) {
queryClient.setQueryData<ActiveJobsResponse>([QueryKeys.activeJobs], (old) => ({
activeJobIds: (old?.activeJobIds ?? []).filter((id) => id !== conversationId),
}));
try {
console.log('[useChatHelpers] Calling abort mutation for:', conversationId);
await abortMutation.mutateAsync({ conversationId });
@ -146,7 +151,7 @@ export default function useChatHelpers(index = 0, paramId?: string) {
console.log('[useChatHelpers] Assistants endpoint, just clearing submissions');
clearAllSubmissions();
}
}, [conversationId, endpoint, endpointType, abortMutation, clearAllSubmissions]);
}, [conversationId, endpoint, endpointType, abortMutation, clearAllSubmissions, queryClient]);
const handleStopGenerating = (e: React.MouseEvent<HTMLButtonElement>) => {
e.preventDefault();

View file

@ -6,23 +6,19 @@ import { useQueryClient } from '@tanstack/react-query';
import {
request,
Constants,
QueryKeys,
createPayload,
LocalStorageKeys,
removeNullishValues,
} from 'librechat-data-provider';
import type { TMessage, TPayload, TSubmission, EventSubmission } from 'librechat-data-provider';
import type { EventHandlerParams } from './useEventHandlers';
import { useGetStartupConfig, useGetUserBalance } from '~/data-provider';
import { activeJobsQueryKey, queueTitleGeneration } from '~/data-provider/SSE/queries';
import { useGetStartupConfig, useGetUserBalance, queueTitleGeneration } from '~/data-provider';
import type { ActiveJobsResponse } from '~/data-provider';
import { useAuthContext } from '~/hooks/AuthContext';
import useEventHandlers from './useEventHandlers';
import store from '~/store';
/** Response type for active jobs query */
interface ActiveJobsResponse {
activeJobIds: string[];
}
const clearDraft = (conversationId?: string | null) => {
if (conversationId) {
localStorage.removeItem(`${LocalStorageKeys.TEXT_DRAFT}${conversationId}`);
@ -72,7 +68,7 @@ export default function useResumableSSE(
*/
const addActiveJob = useCallback(
(jobId: string) => {
queryClient.setQueryData<ActiveJobsResponse>(activeJobsQueryKey, (old) => ({
queryClient.setQueryData<ActiveJobsResponse>([QueryKeys.activeJobs], (old) => ({
activeJobIds: [...new Set([...(old?.activeJobIds ?? []), jobId])],
}));
},
@ -85,7 +81,7 @@ export default function useResumableSSE(
*/
const removeActiveJob = useCallback(
(jobId: string) => {
queryClient.setQueryData<ActiveJobsResponse>(activeJobsQueryKey, (old) => ({
queryClient.setQueryData<ActiveJobsResponse>([QueryKeys.activeJobs], (old) => ({
activeJobIds: (old?.activeJobIds ?? []).filter((id) => id !== jobId),
}));
},