diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
index 66048b422646..031738dd2c69 100644
--- a/front/components/assistant/conversation/AgentMessage.tsx
+++ b/front/components/assistant/conversation/AgentMessage.tsx
@@ -509,6 +509,27 @@ export function AgentMessage({
     return (
       <div className="flex flex-col gap-y-4">
         <AgentMessageActions agentMessage={agentMessage} size={size} />
+        <>
+          {agentMessage.actions
+            .filter((a) => isVisualizationActionType(a))
+            .map((a, i) => {
+              const streamingViz = streamedVisualizations.find(
+                (sv) => sv.actionId === a.id
+              );
+              assert(isVisualizationActionType(a));
+              return (
+                <VisualizationActionIframe
+                  action={a}
+                  conversationId={conversationId}
+                  isStreaming={!!streamingViz}
+                  key={i}
+                  onRetry={() => retryHandler(agentMessage)}
+                  owner={owner}
+                  streamedCode={streamingViz?.visualization || null}
+                />
+              );
+            })}
+        </>
 
         {agentMessage.chainOfThought?.length ? (
           <div className="flex w-full flex-col gap-2 rounded-2xl border border-slate-200 bg-slate-100 p-4 text-sm text-slate-800">
@@ -556,27 +577,6 @@ export function AgentMessage({
             )}
           </div>
         )}
-        <>
-          {agentMessage.actions
-            .filter((a) => isVisualizationActionType(a))
-            .map((a, i) => {
-              const streamingViz = streamedVisualizations.find(
-                (sv) => sv.actionId === a.id
-              );
-              assert(isVisualizationActionType(a));
-              return (
-                <VisualizationActionIframe
-                  action={a}
-                  conversationId={conversationId}
-                  isStreaming={!!streamingViz}
-                  key={i}
-                  onRetry={() => retryHandler(agentMessage)}
-                  owner={owner}
-                  streamedCode={streamingViz?.visualization || null}
-                />
-              );
-            })}
-        </>
         {agentMessage.status === "cancelled" && (
           <Chip
             label="Message generation was interrupted"
diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
index 830d4cb1566f..21a6b708964a 100644
--- a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
+++ b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -1,4 +1,4 @@
-import { BracesIcon, PlayIcon, Tab } from "@dust-tt/sparkle";
+import { BracesIcon, IconToggleButton } from "@dust-tt/sparkle";
 import type {
   CommandResultMap,
   VisualizationActionType,
@@ -15,6 +15,7 @@ import type { SetStateAction } from "react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 
 import { RenderMessageMarkdown } from "@app/components/assistant/RenderMessageMarkdown";
+import { classNames } from "@app/lib/utils";
 
 const sendResponseToIframe = <T extends VisualizationRPCCommand>(
   request: { command: T } & VisualizationRPCRequest,
@@ -135,8 +136,7 @@ export function VisualizationActionIframe({
   isStreaming: boolean;
   onRetry: () => void;
 }) {
-  const [activeTab, setActiveTab] = useState<"code" | "runtime">("code");
-  const [tabManuallyChanged, setTabManuallyChanged] = useState(false);
+  const [showIframe, setShowIframe] = useState<boolean | null>(null);
   const [contentHeight, setContentHeight] = useState(0);
 
   const workspaceId = owner.sId;
@@ -148,11 +148,10 @@ export function VisualizationActionIframe({
   });
 
   useEffect(() => {
-    if (activeTab === "code" && action.generation && !tabManuallyChanged) {
-      setActiveTab("runtime");
-      setTabManuallyChanged(true);
+    if (showIframe === null && action.generation) {
+      setShowIframe(true);
     }
-  }, [action.generation, activeTab, tabManuallyChanged]);
+  }, [action.generation, showIframe]);
 
   let extractedCode: string | null = null;
 
@@ -161,48 +160,48 @@ export function VisualizationActionIframe({
   );
 
   return (
-    <>
-      <Tab
-        tabs={[
-          {
-            label: "Code",
-            id: "code",
-            current: activeTab === "code",
-            icon: BracesIcon,
-            sizing: "expand",
-          },
-          {
-            label: "Run",
-            id: "runtime",
-            current: activeTab === "runtime",
-            icon: PlayIcon,
-            sizing: "expand",
-            hasSeparator: true,
-          },
-        ]}
-        setCurrentTab={(tabId, event) => {
-          event.preventDefault();
-          setActiveTab(tabId as "code" | "runtime");
-        }}
-      />
-      {activeTab === "code" && extractedCode && extractedCode.length > 0 && (
-        <RenderMessageMarkdown
-          content={"```javascript\n" + extractedCode + "\n```"}
-          isStreaming={isStreaming}
+    <div className="relative">
+      {showIframe && (
+        <div
+          style={{
+            height: `${contentHeight}px`,
+          }}
         />
       )}
-      {activeTab === "runtime" && (
-        <div
-          style={{ height: `${contentHeight}px` }}
-          className="max-h-[40vh] w-full"
-        >
-          <iframe
-            className="h-full w-full"
-            src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
-            sandbox="allow-scripts"
+      <div>
+        {!(showIframe && contentHeight > 0) && extractedCode && (
+          <RenderMessageMarkdown
+            content={"```javascript\n" + extractedCode + "\n```"}
+            isStreaming={isStreaming}
+          />
+        )}
+
+        {!!action.generation && (
+          <div
+            style={{ height: `${contentHeight}px` }}
+            className={classNames(
+              "absolute left-0 top-0 max-h-[60vh] w-full",
+              !showIframe && contentHeight > 0 ? "opacity-0" : "opacity-100"
+            )}
+          >
+            <iframe
+              className="h-full w-full"
+              src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
+              sandbox="allow-scripts"
+            />
+          </div>
+        )}
+      </div>
+
+      {action.generation && contentHeight > 0 && (
+        <div className="absolute left-4 top-4">
+          <IconToggleButton
+            icon={BracesIcon}
+            selected={!showIframe}
+            onClick={() => setShowIframe((prev) => !prev)}
           />
         </div>
       )}
-    </>
+    </div>
   );
 }
diff --git a/types/src/front/lib/actions/registry.ts b/types/src/front/lib/actions/registry.ts
index 6d712ab18e00..5c800b598e25 100644
--- a/types/src/front/lib/actions/registry.ts
+++ b/types/src/front/lib/actions/registry.ts
@@ -238,7 +238,7 @@ export const DustProdActionRegistry = createActionRegistry({
       workspaceId: PRODUCTION_DUST_APPS_WORKSPACE_ID,
       appId: "tWcuYDj1OE",
       appHash:
-        "e404a057516d5d1bb7b1807b74a20484c3f60bf0a36c510cb799127f060f032b",
+        "8298c6543759d1d11db0e360a8b7aa7b8ec0fa71ed274f2667678302073e4f8d",
     },
     config: {
       MODEL: {