From 21e4eadadb806460184068c50bab8fa077a428cb Mon Sep 17 00:00:00 2001
From: Hailong Cui <ihailong@amazon.com>
Date: Mon, 11 Dec 2023 16:47:57 +0800
Subject: [PATCH] Address review comments

Signed-off-by: Hailong Cui <ihailong@amazon.com>
---
 public/chat_header_button.tsx     | 4 +---
 public/hooks/use_chat_actions.tsx | 4 +++-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/public/chat_header_button.tsx b/public/chat_header_button.tsx
index b45d98b4..9b3d1503 100644
--- a/public/chat_header_button.tsx
+++ b/public/chat_header_button.tsx
@@ -65,9 +65,7 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
       selectedTabId,
       preSelectedTabId,
       setSelectedTabId: (tabId: TabId) => {
-        if (selectedTabId !== tabId) {
-          setPreSelectedTabId(selectedTabId);
-        }
+        setPreSelectedTabId(selectedTabId);
         setSelectedTabId(tabId);
       },
       flyoutVisible,
diff --git a/public/hooks/use_chat_actions.tsx b/public/hooks/use_chat_actions.tsx
index e1975b46..c78bff31 100644
--- a/public/hooks/use_chat_actions.tsx
+++ b/public/hooks/use_chat_actions.tsx
@@ -138,7 +138,9 @@ export const useChatActions = (): AssistantActions => {
 
       case 'view_trace':
         if ('traceId' in message) {
-          chatContext.setSelectedTabId('trace');
+          if (chatContext.selectedTabId !== TAB_ID.TRACE) {
+            chatContext.setSelectedTabId(TAB_ID.TRACE);
+          }
           chatContext.setTraceId(message.traceId);
         }
         break;