Thinking block persist toggle state (#4916)

* Thinking block persist toggle state * dev build
2026-01-28 13:27:00 -08:00 · 2026-01-28 13:27:00 -08:00 · 88459ce2d2
commit 88459ce2d2
parent 9584ebcd2c
5 changed files with 127 additions and 52 deletions
--- a/.github/workflows/dev-build.yaml
+++ b/.github/workflows/dev-build.yaml
@ -6,7 +6,7 @@ concurrency:

 on:
  push:
-    branches: ['docker-model-runner-download-from-ui'] # put your current branch to create a build. Core team only.
+    branches: ['4855-thinking-block-toggle'] # put your current branch to create a build. Core team only.
    paths-ignore:
      - '**.md'
      - 'cloud-deployments/*'
--- a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/HistoricalMessage/index.jsx
+++ b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/HistoricalMessage/index.jsx
@ -115,7 +115,7 @@ const HistoricalMessage = ({
              <RenderChatContent
                role={role}
                message={message}
-                expanded={isLastMessage}
+                messageId={uuid}
              />
              {isRefusalMessage && (
                <Link
@ -212,7 +212,7 @@ function ChatAttachments({ attachments = [] }) {
 }

 const RenderChatContent = memo(
-  ({ role, message, expanded = false }) => {
+  ({ role, message, messageId }) => {
    // If the message is not from the assistant, we can render it directly
    // as normal since the user cannot think (lol)
    if (role !== "assistant")
@ -240,18 +240,16 @@ const RenderChatContent = memo(
    // This can occur when the assistant starts with <thinking> and then <response>'s later.
    if (
      message.match(THOUGHT_REGEX_OPEN) &&
-      message.match(THOUGHT_REGEX_CLOSE)
+      !message.match(THOUGHT_REGEX_CLOSE)
    ) {
-      const closingTag = message.match(THOUGHT_REGEX_CLOSE)?.[0];
-      const splitMessage = message.split(closingTag);
-      thoughtChain = splitMessage[0] + closingTag;
-      msgToRender = splitMessage[1];
+      thoughtChain = message;
+      msgToRender = "";
    }

    return (
      <>
        {thoughtChain && (
-          <ThoughtChainComponent content={thoughtChain} expanded={expanded} />
+          <ThoughtChainComponent content={thoughtChain} messageId={messageId} />
        )}
        <span
          className="flex flex-col gap-y-1"
@ -266,7 +264,7 @@ const RenderChatContent = memo(
    return (
      prevProps.role === nextProps.role &&
      prevProps.message === nextProps.message &&
-      prevProps.expanded === nextProps.expanded
+      prevProps.messageId === nextProps.messageId
    );
  }
 );
--- a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/PromptReply/index.jsx
+++ b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/PromptReply/index.jsx
@ -70,6 +70,7 @@ const PromptReply = ({
          <RenderAssistantChatContent
            key={`${uuid}-prompt-reply-content`}
            message={reply}
+            messageId={uuid}
          />
        </div>
        <Citations sources={sources} />
@ -94,7 +95,7 @@ export function WorkspaceProfileImage({ workspace }) {
  return <UserIcon user={{ uid: workspace.slug }} role="assistant" />;
 }

-function RenderAssistantChatContent({ message }) {
+function RenderAssistantChatContent({ message, messageId }) {
  const contentRef = useRef("");
  const thoughtChainRef = useRef(null);

@ -121,7 +122,11 @@ function RenderAssistantChatContent({ message }) {
    message.match(THOUGHT_REGEX_OPEN) && !message.match(THOUGHT_REGEX_CLOSE);
  if (thinking)
    return (
-      <ThoughtChainComponent ref={thoughtChainRef} content="" expanded={true} />
+      <ThoughtChainComponent
+        ref={thoughtChainRef}
+        content=""
+        messageId={messageId}
+      />
    );

  return (
@ -130,7 +135,7 @@ function RenderAssistantChatContent({ message }) {
        <ThoughtChainComponent
          ref={thoughtChainRef}
          content=""
-          expanded={true}
+          messageId={messageId}
        />
      )}
      <span
--- a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/ThoughtContainer/index.jsx
+++ b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/ThoughtContainer/index.jsx
@ -1,4 +1,12 @@
-import { useState, forwardRef, useImperativeHandle } from "react";
+import {
+  useState,
+  useEffect,
+  forwardRef,
+  useImperativeHandle,
+  createContext,
+  useContext,
+  useCallback,
+} from "react";
 import renderMarkdown from "@/utils/chat/markdown";
 import { CaretDown } from "@phosphor-icons/react";
 import DOMPurify from "dompurify";
@ -6,6 +14,50 @@ import { isMobile } from "react-device-detect";
 import ThinkingAnimation from "@/media/animations/thinking-animation.webm";
 import ThinkingStatic from "@/media/animations/thinking-static.png";

+/**
+ * Context to persist thought expansion state across component transitions
+ * (e.g., from PromptReply to HistoricalMessage)
+ */
+const ThoughtExpansionContext = createContext(null);
+
+export function ThoughtExpansionProvider({ children }) {
+  const [expansionStates, setExpansionStates] = useState({});
+
+  const getExpanded = useCallback(
+    (messageId) => {
+      if (!messageId) return false;
+      return expansionStates[messageId] ?? false;
+    },
+    [expansionStates]
+  );
+
+  const setExpanded = useCallback((messageId, expanded) => {
+    if (!messageId) return;
+    setExpansionStates((prev) => ({
+      ...prev,
+      [messageId]: expanded,
+    }));
+  }, []);
+
+  return (
+    <ThoughtExpansionContext.Provider value={{ getExpanded, setExpanded }}>
+      {children}
+    </ThoughtExpansionContext.Provider>
+  );
+}
+
+export function useThoughtExpansion(messageId) {
+  const context = useContext(ThoughtExpansionContext);
+  if (!context) {
+    // Fallback when used outside provider - use local state only
+    return { expanded: false, setExpanded: () => {} };
+  }
+  return {
+    expanded: context.getExpanded(messageId),
+    setExpanded: (value) => context.setExpanded(messageId, value),
+  };
+}
+
 const THOUGHT_KEYWORDS = ["thought", "thinking", "think", "thought_chain"];
 const CLOSING_TAGS = [...THOUGHT_KEYWORDS, "response", "answer"];
 export const THOUGHT_REGEX_OPEN = new RegExp(
@ -40,16 +92,31 @@ function contentIsNotEmpty(content = "") {
 /**
 * Component to render a thought chain.
 * @param {string} content - The content of the thought chain.
- * @param {boolean} expanded - Whether the thought chain is expanded.
+ * @param {string} messageId - The unique ID for this message (used to persist expansion state).
 * @returns {JSX.Element}
 */
 export const ThoughtChainComponent = forwardRef(
-  ({ content: initialContent, expanded }, ref) => {
+  ({ content: initialContent, messageId }, ref) => {
    const [content, setContent] = useState(initialContent);
    const [hasReadableContent, setHasReadableContent] = useState(
      contentIsNotEmpty(initialContent)
    );
-    const [isExpanded, setIsExpanded] = useState(expanded);
+    const { expanded: persistedExpanded, setExpanded: setPersistedExpanded } =
+      useThoughtExpansion(messageId);
+    const [localExpanded, setLocalExpanded] = useState(false);
+
+    // Use persisted state if messageId is provided, otherwise use local state
+    const isExpanded = messageId ? persistedExpanded : localExpanded;
+    const setIsExpanded = messageId ? setPersistedExpanded : setLocalExpanded;
+
+    // Sync content state with prop changes (for streaming through HistoricalMessage)
+    useEffect(() => {
+      if (initialContent !== content) {
+        setContent(initialContent);
+        setHasReadableContent(contentIsNotEmpty(initialContent));
+      }
+    }, [initialContent]);
+
    useImperativeHandle(ref, () => ({
      updateContent: (newContent) => {
        setContent(newContent);
@ -65,8 +132,6 @@ export const ThoughtChainComponent = forwardRef(
    const tagStrippedContent = content
      .replace(THOUGHT_REGEX_OPEN, "")
      .replace(THOUGHT_REGEX_CLOSE, "");
-    const autoExpand =
-      isThinking && tagStrippedContent.length > THOUGHT_PREVIEW_LENGTH;
    const canExpand = tagStrippedContent.length > THOUGHT_PREVIEW_LENGTH;
    if (!content || !content.length || !hasReadableContent) return null;

@ -83,10 +148,10 @@ export const ThoughtChainComponent = forwardRef(
              transition: "all 0.1s ease-in-out",
              borderRadius: "6px",
            }}
-            className={`${isExpanded || autoExpand ? "" : `${canExpand ? "hover:bg-theme-sidebar-item-hover" : ""}`} items-start bg-theme-bg-chat-input py-2 px-4 flex gap-x-2`}
+            className={`${isExpanded ? "" : `${canExpand ? "hover:bg-theme-sidebar-item-hover" : ""}`} items-start bg-theme-bg-chat-input py-2 px-4 flex gap-x-2`}
          >
            <div
-              className={`w-7 h-7 flex justify-center flex-shrink-0 ${!isExpanded && !autoExpand ? "items-center" : "items-start pt-[2px]"}`}
+              className={`w-7 h-7 flex justify-center flex-shrink-0 ${!isExpanded ? "items-center" : "items-start pt-[2px]"}`}
            >
              {isThinking || isComplete ? (
                <>
@ -115,16 +180,16 @@ export const ThoughtChainComponent = forwardRef(
            </div>
            <div className="flex-1 min-w-0">
              <div
-                className={`overflow-hidden transition-all transform duration-300 ease-in-out origin-top ${isExpanded || autoExpand ? "" : "max-h-6"}`}
+                className={`overflow-hidden transition-all transform duration-300 ease-in-out origin-top ${isExpanded ? "" : "max-h-6"}`}
              >
                <div
-                  className={`text-theme-text-secondary font-mono leading-6 ${isExpanded || autoExpand ? "-ml-[5.5px] -mt-[4px]" : "mt-[2px]"}`}
+                  className={`text-theme-text-secondary font-mono leading-6 ${isExpanded ? "-ml-[5.5px] -mt-[4px]" : "mt-[2px]"}`}
                >
                  <span
-                    className={`block w-full ${!isExpanded && !autoExpand ? "truncate" : ""}`}
+                    className={`block w-full ${!isExpanded ? "truncate" : ""}`}
                    dangerouslySetInnerHTML={{
                      __html: DOMPurify.sanitize(
-                        isExpanded || autoExpand
+                        isExpanded
                          ? renderMarkdown(tagStrippedContent)
                          : tagStrippedContent
                      ),
@ -134,7 +199,7 @@ export const ThoughtChainComponent = forwardRef(
              </div>
            </div>
            <div className="flex items-center gap-x-2">
-              {!autoExpand && canExpand ? (
+              {canExpand ? (
                <button
                  onClick={handleExpandClick}
                  data-tooltip-id="expand-cot"
--- a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/index.jsx
+++ b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/index.jsx
@ -24,6 +24,7 @@ import useChatHistoryScrollHandle from "@/hooks/useChatHistoryScrollHandle";
 import { v4 } from "uuid";
 import { useTranslation } from "react-i18next";
 import { useChatMessageAlignment } from "@/hooks/useChatMessageAlignment";
+import { ThoughtExpansionProvider } from "./ThoughtContainer";

 export default forwardRef(function (
  {
@ -238,34 +239,39 @@ export default forwardRef(function (
  }

  return (
-    <div
-      className={`markdown text-white/80 light:text-theme-text-primary font-light ${textSizeClass} h-full md:h-[83%] pb-[100px] pt-6 md:pt-0 md:pb-20 md:mx-0 overflow-y-scroll flex flex-col justify-start ${showScrollbar ? "show-scrollbar" : "no-scroll"}`}
-      id="chat-history"
-      ref={chatHistoryRef}
-      onScroll={handleScroll}
-    >
-      {compiledHistory.map((item, index) =>
-        Array.isArray(item) ? renderStatusResponse(item, index) : item
-      )}
-      {showing && (
-        <ManageWorkspace hideModal={hideModal} providedSlug={workspace.slug} />
-      )}
-      {!isAtBottom && (
-        <div className="fixed bottom-40 right-10 md:right-20 z-50 cursor-pointer animate-pulse">
-          <div className="flex flex-col items-center">
-            <div
-              className="p-1 rounded-full border border-white/10 bg-white/10 hover:bg-white/20 hover:text-white"
-              onClick={() => {
-                scrollToBottom(isStreaming ? false : true);
-                setIsUserScrolling(false);
-              }}
-            >
-              <ArrowDown weight="bold" className="text-white/60 w-5 h-5" />
+    <ThoughtExpansionProvider>
+      <div
+        className={`markdown text-white/80 light:text-theme-text-primary font-light ${textSizeClass} h-full md:h-[83%] pb-[100px] pt-6 md:pt-0 md:pb-20 md:mx-0 overflow-y-scroll flex flex-col justify-start ${showScrollbar ? "show-scrollbar" : "no-scroll"}`}
+        id="chat-history"
+        ref={chatHistoryRef}
+        onScroll={handleScroll}
+      >
+        {compiledHistory.map((item, index) =>
+          Array.isArray(item) ? renderStatusResponse(item, index) : item
+        )}
+        {showing && (
+          <ManageWorkspace
+            hideModal={hideModal}
+            providedSlug={workspace.slug}
+          />
+        )}
+        {!isAtBottom && (
+          <div className="fixed bottom-40 right-10 md:right-20 z-50 cursor-pointer animate-pulse">
+            <div className="flex flex-col items-center">
+              <div
+                className="p-1 rounded-full border border-white/10 bg-white/10 hover:bg-white/20 hover:text-white"
+                onClick={() => {
+                  scrollToBottom(isStreaming ? false : true);
+                  setIsUserScrolling(false);
+                }}
+              >
+                <ArrowDown weight="bold" className="text-white/60 w-5 h-5" />
+              </div>
            </div>
          </div>
-        </div>
-      )}
-    </div>
+        )}
+      </div>
+    </ThoughtExpansionProvider>
  );
 });

@ -352,6 +358,7 @@ function buildMessages({
      acc.push(
        <HistoricalMessage
          key={index}
+          uuid={props.uuid}
          message={props.content}
          role={props.role}
          workspace={workspace}