From 6aebccd20ad7322db39619b23b298f9cc5bb6295 Mon Sep 17 00:00:00 2001
From: Danny Avila <danny@librechat.ai>
Date: Thu, 18 Dec 2025 09:33:00 -0500
Subject: [PATCH] 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.
---
 client/src/data-provider/SSE/queries.ts |  8 ++++++--
 client/src/hooks/Chat/useChatHelpers.ts |  9 +++++++--
 client/src/hooks/SSE/useResumableSSE.ts | 14 +++++---------
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/client/src/data-provider/SSE/queries.ts b/client/src/data-provider/SSE/queries.ts
index 01907d198a..b6039d1529 100644
--- a/client/src/data-provider/SSE/queries.ts
+++ b/client/src/data-provider/SSE/queries.ts
@@ -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,
diff --git a/client/src/hooks/Chat/useChatHelpers.ts b/client/src/hooks/Chat/useChatHelpers.ts
index cea7bcfe17..46d38d3a4d 100644
--- a/client/src/hooks/Chat/useChatHelpers.ts
+++ b/client/src/hooks/Chat/useChatHelpers.ts
@@ -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();
diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts
index 6aeb181dcd..3eefcb7d40 100644
--- a/client/src/hooks/SSE/useResumableSSE.ts
+++ b/client/src/hooks/SSE/useResumableSSE.ts
@@ -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),
       }));
     },