feat: add queryhook layer and invalidation so we keep tokens fresh during conversation turns

client/src/data-provider/queries.ts

@@ -32,6 +32,20 @@ import type {
 } from 'librechat-data-provider';
 import type { ConversationCursorData } from '~/utils/convos';
 import { findConversationInInfinite } from '~/utils';
+type TConversationCosts = {
+  conversationId: string;
+  totals: {
+    prompt: { usd: number; tokenCount: number };
+    completion: { usd: number; tokenCount: number };
+    total: { usd: number; tokenCount: number };
+  };
+  perMessage: Array<{
+    messageId: string;
+    tokenType: 'prompt' | 'completion';
+    tokenCount: number;
+    usd: number;
+  }>;
+};
 
 export const useGetPresetsQuery = (
   config?: UseQueryOptions<TPreset[]>,
@@ -76,6 +90,25 @@ export const useGetConvoIdQuery = (
   );
 };
 
+export const useGetConversationCosts = (
+  conversationId: string,
+  config?: UseQueryOptions<TConversationCosts>,
+): QueryObserverResult<TConversationCosts> => {
+  return useQuery<TConversationCosts>(
+    [QueryKeys.conversation, conversationId, 'costs'],
+    () => {
+      return dataService.getConversationCosts(conversationId) as unknown as TConversationCosts;
+    },
+    {
+      refetchOnWindowFocus: false,
+      refetchOnReconnect: false,
+      refetchOnMount: false,
+      enabled: !!conversationId && conversationId !== 'new',
+      ...config,
+    },
+  );
+};
+
 export const useConversationsInfiniteQuery = (
   params: ConversationListParams,
   config?: UseInfiniteQueryOptions<ConversationListResponse, unknown>,

client/src/hooks/SSE/useEventHandlers.ts

@@ -232,8 +232,14 @@ export default function useEventHandlers({
           },
         ]);
       }
+
+      if (userMessage?.conversationId) {
+        queryClient.invalidateQueries({
+          queryKey: [QueryKeys.conversation, userMessage.conversationId, 'costs'],
+        });
+      }
     },
-    [setMessages, announcePolite, setIsSubmitting],
+    [setMessages, announcePolite, setIsSubmitting, queryClient],
   );
 
   const cancelHandler = useCallback(
@@ -275,6 +281,12 @@ export default function useEventHandlers({
         });
       }
 
+      if (convoUpdate?.conversationId) {
+        queryClient.invalidateQueries({
+          queryKey: [QueryKeys.conversation, convoUpdate.conversationId, 'costs'],
+        });
+      }
+
       setIsSubmitting(false);
     },
     [setMessages, setConversation, genTitle, isAddedRequest, queryClient, setIsSubmitting],
@@ -341,6 +353,12 @@ export default function useEventHandlers({
       if (resetLatestMessage) {
         resetLatestMessage();
       }
+
+      if (conversationId) {
+        queryClient.invalidateQueries({
+          queryKey: [QueryKeys.conversation, conversationId, 'costs'],
+        });
+      }
     },
     [
       queryClient,
@@ -527,6 +545,12 @@ export default function useEventHandlers({
         );
       }
 
+      if (conversation.conversationId) {
+        queryClient.invalidateQueries({
+          queryKey: [QueryKeys.conversation, conversation.conversationId, 'costs'],
+        });
+      }
+
       if (isNewConvo && submissionConvo.conversationId) {
         removeConvoFromAllQueries(queryClient, submissionConvo.conversationId);
       }