From 05a4a631c8507450e9b84932e5a24a840b89a6e6 Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Mon, 14 Apr 2025 14:22:12 -0600
Subject: [PATCH] avoid duplicate version checks for the metric

---
 .../api/buster_rest/metrics/queryRequests.ts  |  8 +++
 .../ChatLayout/ChatContext/ChatContext.tsx    |  2 -
 .../ChatContext/useAutoChangeLayout.ts        | 60 +++++++++++++------
 3 files changed, 50 insertions(+), 20 deletions(-)

diff --git a/web/src/api/buster_rest/metrics/queryRequests.ts b/web/src/api/buster_rest/metrics/queryRequests.ts
index 79f4e55c4..15d1f108b 100644
--- a/web/src/api/buster_rest/metrics/queryRequests.ts
+++ b/web/src/api/buster_rest/metrics/queryRequests.ts
@@ -93,6 +93,14 @@ export const useGetMetric = <TData = IBusterMetric>(
     const isLatestVersion =
       updatedMetric.version_number === last(updatedMetric.versions)?.version_number;
     if (isLatestVersion) setOriginalMetric(updatedMetric);
+
+    if (!versionNumber && result?.version_number) {
+      queryClient.setQueryData(
+        metricsQueryKeys.metricsGetMetric(result.id, result.version_number).queryKey,
+        updatedMetric
+      );
+    }
+
     return updatedMetric;
   });
 
diff --git a/web/src/layouts/ChatLayout/ChatContext/ChatContext.tsx b/web/src/layouts/ChatLayout/ChatContext/ChatContext.tsx
index b66a289d0..dc62599b0 100644
--- a/web/src/layouts/ChatLayout/ChatContext/ChatContext.tsx
+++ b/web/src/layouts/ChatLayout/ChatContext/ChatContext.tsx
@@ -94,8 +94,6 @@ export const ChatContextProvider = React.memo(({ children }: PropsWithChildren<{
     onSetSelectedFile
   });
 
-  console.log(selectedFile);
-
   return (
     <IndividualChatContext.Provider value={useChatContextValue}>
       {children}
diff --git a/web/src/layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts b/web/src/layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts
index e71b3b08b..d8e51341f 100644
--- a/web/src/layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts
+++ b/web/src/layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts
@@ -1,6 +1,10 @@
 'use client';
 
-import { useGetChatMessageMemoized, useGetChatMessage } from '@/api/buster_rest/chats';
+import {
+  useGetChatMessageMemoized,
+  useGetChatMessage,
+  useGetChatMemoized
+} from '@/api/buster_rest/chats';
 import type { SelectedFile } from '../interfaces';
 import { MutableRefObject, useEffect, useRef } from 'react';
 import findLast from 'lodash/findLast';
@@ -20,6 +24,8 @@ export const useAutoChangeLayout = ({
   selectedFileId: string | undefined;
   chatId: string | undefined;
 }) => {
+  const getChatMemoized = useGetChatMemoized();
+  const getChatMessageMemoized = useGetChatMessageMemoized();
   const onSetSelectedFile = useChatLayoutContextSelector((x) => x.onSetSelectedFile);
   const messageId = useChatLayoutContextSelector((x) => x.messageId);
   const metricId = useChatLayoutContextSelector((x) => x.metricId);
@@ -35,7 +41,6 @@ export const useAutoChangeLayout = ({
     lastMessageId,
     (x) => x?.reasoning_message_ids?.length || 0
   );
-  const getChatMessageMemoized = useGetChatMessageMemoized();
   const { getFileLinkMeta } = useGetFileLink();
 
   const previousIsCompletedStream = usePrevious(isCompletedStream);
@@ -100,11 +105,11 @@ export const useAutoChangeLayout = ({
         return;
       }
 
-      const chatMessage = getChatMessageMemoized(lastMessageId);
+      const chat = getChatMemoized(chatId);
 
       //reasoning_message_mode
       if (messageId) {
-        const messageExists = !!chatMessage?.reasoning_message_ids.find((id) => id === messageId);
+        const messageExists = !!chat?.message_ids.some((id) => id === messageId);
         if (messageExists) {
           return;
         } else {
@@ -119,16 +124,26 @@ export const useAutoChangeLayout = ({
 
       //dashboard_mode
       if (dashboardId) {
+        console.log('dashboardId', dashboardId);
         if (!dashboardVersionNumber) {
-          const lastMatchingDashboardInChat = chatMessage?.response_message_ids.reduce<
+          const lastMatchingDashboardInChat = chat?.message_ids.reduce<
             BusterChatResponseMessage_file | undefined
-          >((acc, messageId) => {
-            const message = chatMessage?.response_messages[messageId]!;
-            const isFile =
-              message.type === 'file' &&
-              message.file_type === 'dashboard' &&
-              message.id === dashboardId;
-            if (isFile) return message;
+          >((acc, chatMessageId) => {
+            const chatMessage = getChatMessageMemoized(chatMessageId);
+
+            console.log('chatMessage', chatMessage);
+
+            chatMessage?.response_message_ids.forEach((responseMessageId) => {
+              const message = chatMessage?.response_messages[responseMessageId]!;
+              const isFile =
+                message.type === 'file' &&
+                message.file_type === 'dashboard' &&
+                message.id === dashboardId;
+
+              if (isFile) {
+                acc = message;
+              }
+            });
 
             return acc;
           }, undefined);
@@ -151,13 +166,22 @@ export const useAutoChangeLayout = ({
       //metric_mode
       if (metricId) {
         if (!metricVersionNumber) {
-          const lastMatchingMetricInChat = chatMessage?.response_message_ids.reduce<
+          const lastMatchingMetricInChat = chat?.message_ids.reduce<
             BusterChatResponseMessage_file | undefined
-          >((acc, messageId) => {
-            const message = chatMessage?.response_messages[messageId]!;
-            const isFile =
-              message.type === 'file' && message.file_type === 'metric' && message.id === metricId;
-            if (isFile) return message;
+          >((acc, chatMessageId) => {
+            const chatMessage = getChatMessageMemoized(chatMessageId);
+
+            chatMessage?.response_message_ids.forEach((responseMessageId) => {
+              const message = chatMessage?.response_messages[responseMessageId]!;
+              const isFile =
+                message.type === 'file' &&
+                message.file_type === 'metric' &&
+                message.id === metricId;
+
+              if (isFile) {
+                acc = message;
+              }
+            });
 
             return acc;
           }, undefined);