update switching logic

web/src/app/chat/ChatPage.tsx

@@ -2614,6 +2614,13 @@ export function ChatPage({
                                     {message.sub_questions &&
                                     message.sub_questions.length > 0 ? (
                                       <AgenticMessage
+                                        docSidebarToggled={
+                                          documentSidebarToggled &&
+                                          (selectedMessageForDocDisplay ==
+                                            message.messageId ||
+                                            selectedMessageForDocDisplay ==
+                                              secondLevelMessage?.messageId)
+                                        }
                                         isImprovement={message.isImprovement}
                                         secondLevelGenerating={
                                           (message.second_level_generating &&
@@ -2707,18 +2714,33 @@ export function ChatPage({
                                           messageHistory.length - 1 == i
                                         }
                                         selectedDocuments={selectedDocuments}
-                                        toggleDocumentSelection={() => {
+                                        toggleDocumentSelection={(
+                                          second: boolean
+                                        ) => {
                                           if (
-                                            !documentSidebarToggled ||
+                                            (!second &&
+                                              !documentSidebarToggled) ||
                                             (documentSidebarToggled &&
                                               selectedMessageForDocDisplay ===
                                                 message.messageId)
                                           ) {
                                             toggleDocumentSidebar();
                                           }
+                                          if (
+                                            (second &&
+                                              !documentSidebarToggled) ||
+                                            (documentSidebarToggled &&
+                                              selectedMessageForDocDisplay ===
+                                                secondLevelMessage?.messageId)
+                                          ) {
+                                            toggleDocumentSidebar();
+                                          }
 
                                           setSelectedMessageForDocDisplay(
-                                            message.messageId
+                                            second
+                                              ? secondLevelMessage?.messageId ||
+                                                  null
+                                              : message.messageId
                                           );
                                         }}
                                         currentPersona={liveAssistant}

web/src/app/chat/message/AgenticMessage.tsx

@@ -50,6 +50,7 @@ import SubQuestionsDisplay from "./SubQuestionsDisplay";
 import { StatusRefinement } from "../Refinement";
 
 export const AgenticMessage = ({
+  docSidebarToggled,
   isImprovement,
   secondLevelAssistantMessage,
   secondLevelGenerating,
@@ -79,6 +80,7 @@ export const AgenticMessage = ({
   secondLevelSubquestions,
   toggleDocDisplay,
 }: {
+  docSidebarToggled?: boolean;
   isImprovement?: boolean | null;
   secondLevelSubquestions?: SubQuestionDetail[] | null;
   agenticDocs?: OnyxDocument[] | null;
@@ -91,7 +93,7 @@ export const AgenticMessage = ({
   otherMessagesCanSwitchTo?: number[];
   onMessageSelection?: (messageId: number) => void;
   selectedDocuments?: OnyxDocument[] | null;
-  toggleDocumentSelection?: () => void;
+  toggleDocumentSelection?: (second: boolean) => void;
   docs?: OnyxDocument[] | null;
   alternativeAssistant?: Persona | null;
   currentPersona: Persona;
@@ -395,6 +397,7 @@ export const AgenticMessage = ({
                 <div className="w-full desktop:ml-4">
                   {subQuestions && subQuestions.length > 0 && (
                     <SubQuestionsDisplay
+                      docSidebarToggled={docSidebarToggled || false}
                       finishedGenerating={finalContent.length > 8}
                       overallAnswerGenerating={finalContent.length < 8}
                       showSecondLevel={!isViewingInitialAnswer}
@@ -403,7 +406,9 @@ export const AgenticMessage = ({
                       subQuestions={subQuestions}
                       secondLevelQuestions={secondLevelSubquestions || []}
                       documents={isViewingInitialAnswer ? docs! : agenticDocs!}
-                      toggleDocumentSelection={toggleDocumentSelection!}
+                      toggleDocumentSelection={() => {
+                        toggleDocumentSelection!(!isViewingInitialAnswer);
+                      }}
                       setPresentingDocument={setPresentingDocument!}
                       unToggle={false}
                     />

web/src/app/chat/message/SourcesDisplay.tsx

@@ -14,6 +14,7 @@ interface SourcesDisplayProps {
   animateEntrance?: boolean;
   threeCols?: boolean;
   hideDocumentDisplay?: boolean;
+  docSidebarToggled?: boolean;
 }
 
 const SourceCard: React.FC<{
@@ -54,6 +55,7 @@ export const SourcesDisplay: React.FC<SourcesDisplayProps> = ({
   animateEntrance = false,
   threeCols = false,
   hideDocumentDisplay = false,
+  docSidebarToggled = false,
 }) => {
   const displayedDocuments = documents.slice(0, 5);
   const hasMoreDocuments = documents.length > 3;
@@ -108,7 +110,7 @@ export const SourcesDisplay: React.FC<SourcesDisplayProps> = ({
 
         {hasMoreDocuments && (
           <SeeMoreBlock
-            toggled={false}
+            toggled={docSidebarToggled}
             toggleDocumentSelection={toggleDocumentSelection}
             uniqueSources={
               Array.from(

web/src/app/chat/message/SubQuestionsDisplay.tsx

@@ -65,6 +65,7 @@ export interface TemporaryDisplay {
   tinyQuestion: string;
 }
 interface SubQuestionsDisplayProps {
+  docSidebarToggled: boolean;
   finishedGenerating: boolean;
   currentlyOpenQuestion?: BaseQuestionIdentifier | null;
   subQuestions: SubQuestionDetail[];
@@ -472,6 +473,7 @@ const SubQuestionsDisplay: React.FC<SubQuestionsDisplayProps> = ({
   toggleDocumentSelection,
   setPresentingDocument,
   secondLevelQuestions,
+  docSidebarToggled,
   showSecondLevel,
   overallAnswerGenerating,
 }) => {
@@ -745,6 +747,7 @@ const SubQuestionsDisplay: React.FC<SubQuestionsDisplayProps> = ({
 
       {shownDocuments && shownDocuments.length > 0 && (
         <SourcesDisplay
+          docSidebarToggled={docSidebarToggled}
           hideDocumentDisplay
           animateEntrance={true}
           toggleDocumentSelection={toggleDocumentSelection}