diff --git a/src/components/chat/DyadThink.tsx b/src/components/chat/DyadThink.tsx
index d975527..4123add 100644
--- a/src/components/chat/DyadThink.tsx
+++ b/src/components/chat/DyadThink.tsx
@@ -22,6 +22,13 @@ export const DyadThink: React.FC<DyadThinkProps> = ({ children, node }) => {
         )
       : null;
 
+  // Collapse when transitioning from in-progress to not-in-progress
+  useEffect(() => {
+    if (!inProgress && isExpanded) {
+      setIsExpanded(false);
+    }
+  }, [inProgress]);
+
   // If it's token savings format, render DyadTokenSavings component
   if (tokenSavingsMatch) {
     const originalTokens = parseFloat(tokenSavingsMatch[1]);
@@ -34,13 +41,6 @@ export const DyadThink: React.FC<DyadThinkProps> = ({ children, node }) => {
     );
   }
 
-  // Collapse when transitioning from in-progress to not-in-progress
-  useEffect(() => {
-    if (!inProgress && isExpanded) {
-      setIsExpanded(false);
-    }
-  }, [inProgress]);
-
   return (
     <div
       className={`relative bg-(--background-lightest) dark:bg-zinc-900 hover:bg-(--background-lighter) rounded-lg px-4 py-2 border my-2 cursor-pointer ${