diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx
index c23681ddb..9765fb858 100644
--- a/web/src/app/chat/ChatPage.tsx
+++ b/web/src/app/chat/ChatPage.tsx
@@ -2176,7 +2176,12 @@ export function ChatPage({
           >
             <DocumentResults
               agenticMessage={
-                aiMessage?.sub_questions?.length! > 0 ? true : false
+                aiMessage?.sub_questions?.length! > 0 ||
+                messageHistory.find(
+                  (m) => m.messageId === aiMessage?.parentMessageId
+                )?.sub_questions?.length! > 0
+                  ? true
+                  : false
               }
               setPresentingDocument={setPresentingDocument}
               modal={true}
@@ -2331,7 +2336,12 @@ export function ChatPage({
           >
             <DocumentResults
               agenticMessage={
-                aiMessage?.sub_questions?.length! > 0 ? true : false
+                aiMessage?.sub_questions?.length! > 0 ||
+                messageHistory.find(
+                  (m) => m.messageId === aiMessage?.parentMessageId
+                )?.sub_questions?.length! > 0
+                  ? true
+                  : false
               }
               setPresentingDocument={setPresentingDocument}
               modal={false}
diff --git a/web/src/app/chat/Refinement.tsx b/web/src/app/chat/Refinement.tsx
index b1f8ae94b..ba762508c 100644
--- a/web/src/app/chat/Refinement.tsx
+++ b/web/src/app/chat/Refinement.tsx
@@ -312,11 +312,7 @@ export function StatusRefinement({
                   onClick={() => {
                     const viewInitialAnswer = !isViewingInitialAnswer;
                     setIsViewingInitialAnswer(viewInitialAnswer);
-                    toggleDocDisplay &&
-                      toggleDocDisplay(!isViewingInitialAnswer);
-                    if (viewInitialAnswer) {
-                      setIsViewingInitialAnswer(true);
-                    }
+                    toggleDocDisplay && toggleDocDisplay(viewInitialAnswer);
                   }}
                 >
                   {isViewingInitialAnswer