From 65cd136f52377ac4c8f50fa46bc4e3aa16b8b660 Mon Sep 17 00:00:00 2001
From: Henry Fontanier <h.fontanier@gmail.com>
Date: Thu, 1 Aug 2024 19:12:15 +0200
Subject: [PATCH] enh: make viz responsive, remove extra renders, handle err in
 main window (#6629)

* enh: improve viz responsiveness, remove extra renders, handle error in main

* mark last viz as complete on agent message success

* naming

* nit

* remove outdated comment

* r

---------

Co-authored-by: Henry Fontanier <henry@dust.tt>
---
 .../assistant/conversation/AgentMessage.tsx   |  6 ++
 .../actions/VisualizationActionIframe.tsx     | 79 ++++++++++++-------
 front/lib/api/assistant/visualization.ts      | 12 ++-
 types/src/front/assistant/visualization.ts    | 39 ++++-----
 viz/app/components/Components.tsx             | 49 ------------
 viz/app/components/ErrorBoundary.tsx          | 61 +-------------
 viz/app/components/VisualizationWrapper.tsx   | 16 +---
 viz/package-lock.json                         | 15 ++--
 viz/package.json                              |  2 +-
 9 files changed, 96 insertions(+), 183 deletions(-)

diff --git a/front/components/assistant/conversation/AgentMessage.tsx b/front/components/assistant/conversation/AgentMessage.tsx
index a8fbf9cb79bd..a7b3afe7ddfe 100644
--- a/front/components/assistant/conversation/AgentMessage.tsx
+++ b/front/components/assistant/conversation/AgentMessage.tsx
@@ -211,6 +211,12 @@ export function AgentMessage({
             ...event.message,
           };
         });
+        // Mark the last viz as complete if it is not already.
+        setVisualizations((v) =>
+          v.map((item, index) =>
+            index === v.length - 1 ? { ...item, complete: true } : item
+          )
+        );
         break;
       }
 
diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
index 0e27c296c686..47091f19bb7b 100644
--- a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
+++ b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -1,4 +1,4 @@
-import { Spinner } from "@dust-tt/sparkle";
+import { Button, Spinner } from "@dust-tt/sparkle";
 import type {
   CommandResultMap,
   VisualizationRPCCommand,
@@ -38,14 +38,14 @@ const sendResponseToIframe = <T extends VisualizationRPCCommand>(
 // Custom hook to encapsulate the logic for handling visualization messages.
 function useVisualizationDataHandler({
   visualization,
-  onRetry,
   setContentHeight,
+  setIsErrored,
   vizIframeRef,
   workspaceId,
 }: {
   visualization: Visualization;
-  onRetry: () => void;
   setContentHeight: (v: SetStateAction<number>) => void;
+  setIsErrored: (v: SetStateAction<boolean>) => void;
   vizIframeRef: React.MutableRefObject<HTMLIFrameElement | null>;
   workspaceId: string;
 }) {
@@ -98,14 +98,14 @@ function useVisualizationDataHandler({
 
           break;
 
-        case "retry":
-          onRetry();
-          break;
-
         case "setContentHeight":
           setContentHeight(data.params.height);
           break;
 
+        case "setErrored":
+          setIsErrored(true);
+          break;
+
         default:
           assertNever(data);
       }
@@ -117,8 +117,8 @@ function useVisualizationDataHandler({
     visualization.identifier,
     code,
     getFileBlob,
-    onRetry,
     setContentHeight,
+    setIsErrored,
     vizIframeRef,
   ]);
 }
@@ -130,45 +130,41 @@ export function VisualizationActionIframe({
 }: {
   owner: WorkspaceType;
   visualization: Visualization;
-
   onRetry: () => void;
 }) {
-  const [contentHeight, setContentHeight] = useState(0);
-  const [iframeLoaded, setIframeLoaded] = useState(false);
-  const [showSpinner, setShowSpinner] = useState(true);
+  const [contentHeight, setContentHeight] = useState<number>(0);
+  const [isErrored, setIsErrored] = useState(false);
   const [activeIndex, setActiveIndex] = useState(1);
 
   const vizIframeRef = useRef<HTMLIFrameElement>(null);
   const containerRef = useRef<HTMLDivElement>(null);
   const codeRef = useRef<HTMLDivElement>(null);
+  const errorRef = useRef<HTMLDivElement>(null);
 
   const workspaceId = owner.sId;
 
   useVisualizationDataHandler({
     visualization,
     workspaceId,
-    onRetry,
     setContentHeight,
+    setIsErrored,
     vizIframeRef,
   });
 
   const { code, complete: codeFullyGenerated } = visualization;
 
+  const iframeLoaded = contentHeight > 0;
+  const showSpinner =
+    ((!codeFullyGenerated && !code) || (codeFullyGenerated && !iframeLoaded)) &&
+    !isErrored;
+
   useEffect(() => {
     if (!codeFullyGenerated) {
-      // Display spinner over the code block while waiting for code generation.
-      setShowSpinner(!code);
       setActiveIndex(0);
-    } else if (iframeLoaded) {
-      // Display iframe if code is generated and iframe has loaded.
-      setShowSpinner(false);
-      setActiveIndex(1);
     } else {
-      // Show spinner while iframe is loading.
-      setShowSpinner(true);
       setActiveIndex(1);
     }
-  }, [codeFullyGenerated, code, iframeLoaded]);
+  }, [codeFullyGenerated, code]);
 
   useEffect(() => {
     if (!containerRef.current) {
@@ -182,9 +178,13 @@ export function VisualizationActionIframe({
         ? `${codeRef.current?.scrollHeight}px`
         : "100%";
     } else if (activeIndex === 1) {
-      containerRef.current.style.height = `${contentHeight}px`;
+      if (isErrored && errorRef.current) {
+        containerRef.current.style.height = `${errorRef.current.scrollHeight}px`;
+      } else if (!isErrored) {
+        containerRef.current.style.height = `${contentHeight}px`;
+      }
     }
-  }, [activeIndex, contentHeight, codeFullyGenerated]);
+  }, [activeIndex, contentHeight, codeFullyGenerated, isErrored]);
 
   return (
     <div className="relative flex flex-col">
@@ -202,7 +202,9 @@ export function VisualizationActionIframe({
       <div
         className={classNames(
           "transition-height relative w-full overflow-hidden duration-500 ease-in-out",
-          codeFullyGenerated ? "min-h-96" : ""
+          codeFullyGenerated && !isErrored ? "min-h-96" : "",
+          isErrored ? "h-full" : "",
+          activeIndex === 1 ? "max-h-[60vh]" : ""
         )}
         ref={containerRef}
       >
@@ -219,21 +221,38 @@ export function VisualizationActionIframe({
             />
           </div>
           <div className="relative flex h-full w-full shrink-0 items-center justify-center">
-            {codeFullyGenerated && (
+            {codeFullyGenerated && !isErrored && (
               <div
-                style={{ height: `${contentHeight}px` }}
+                style={{
+                  height: !isErrored ? `${contentHeight}px` : "100%",
+                  minHeight: !isErrored ? "96" : undefined,
+                }}
                 className={classNames("max-h-[60vh] w-full")}
               >
                 <iframe
                   ref={vizIframeRef}
-                  // Set a min height so iframe can display error.
-                  className="h-full min-h-96 w-full"
+                  className={classNames(
+                    "h-full w-full",
+                    !isErrored ? "min-h-96" : ""
+                  )}
                   src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?identifier=${visualization.identifier}`}
                   sandbox="allow-scripts"
-                  onLoad={() => setIframeLoaded(true)}
                 />
               </div>
             )}
+            {isErrored && (
+              <div
+                className="flex h-full w-full flex-col items-center gap-4 py-8"
+                ref={errorRef}
+              >
+                <div className="text-sm text-element-800">
+                  An error occured while rendering the visualization.
+                </div>
+                <div>
+                  <Button label="Retry" onClick={onRetry} size="sm" />
+                </div>
+              </div>
+            )}
           </div>
         </div>
       </div>
diff --git a/front/lib/api/assistant/visualization.ts b/front/lib/api/assistant/visualization.ts
index ca0d915ae510..bf106c2bfc46 100644
--- a/front/lib/api/assistant/visualization.ts
+++ b/front/lib/api/assistant/visualization.ts
@@ -119,11 +119,15 @@ The generated component should not have any required props / parameters.
 
    
 
-### Outermost div height and width    
+### Responsiveness
 
-The component's outermost JSX tag should have a fixed height and width in pixels, set using the \`style\` prop, e.g. \`<div style={{ height: '600px', width: '600px' }}>...</div>\`.
+The content should be responsive and should not have fixed widths or heights. The component should be able to adapt to different screen sizes.
+The content should never overflow the viewport and should never have horizontal or vertical scrollbars.
 
-The height and width should be set to a fixed value, not a percentage. This style should not use tailwind CSS or any type of custom class. There should be a few pixels of horizontal padding to ensure the content is fully visible by the user.
+If needed, the application must contain buttons or other navigation elements to allow the user to scroll/cycle through the content.
+Never use tailwind's specific values like \`h-[600px]\`.
+
+Always add padding to the content (both horizontal and vertical) to make it look better and make sure the labels are fully visible.
 
     
 
@@ -173,7 +177,7 @@ if (file) {
 
 - Base React is available to be imported. In order to use hooks, they have to be imported at the top of the script, e.g. \`import { useState } from "react"\`
 
-- The recharts charting library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`. Support for defaultProps will be removed from function components in a future major release. JavaScript default parameters should be used instead.
+- The recharts charting library is available to be imported, e.g. \`import { LineChart, XAxis, ... } from "recharts"\` & \`<LineChart ...><XAxis dataKey="name"> ...\`.
 
 - The papaparse library is available to be imported, e.g. \`import Papa from "papaparse"\` & \`const parsed = Papa.parse(fileContent, {header:true, skipEmptyLines: "greedy"});\`. The \`skipEmptyLines:"greedy"\` configuration should always be used.
 
diff --git a/types/src/front/assistant/visualization.ts b/types/src/front/assistant/visualization.ts
index 226f28156b65..47dd8ae8ecfe 100644
--- a/types/src/front/assistant/visualization.ts
+++ b/types/src/front/assistant/visualization.ts
@@ -12,10 +12,6 @@ interface GetFileParams {
   fileId: string;
 }
 
-interface RetryParams {
-  errorMessage: string;
-}
-
 interface SetContentHeightParams {
   height: number;
 }
@@ -24,8 +20,8 @@ interface SetContentHeightParams {
 export type VisualizationRPCRequestMap = {
   getFile: GetFileParams;
   getCodeToExecute: null;
-  retry: RetryParams;
   setContentHeight: SetContentHeightParams;
+  setErrored: void;
 };
 
 // Derive the command type from the keys of the request map
@@ -42,8 +38,8 @@ export type VisualizationRPCRequest = {
 export const validCommands: VisualizationRPCCommand[] = [
   "getFile",
   "getCodeToExecute",
-  "retry",
   "setContentHeight",
+  "setErrored",
 ];
 
 // Command results.
@@ -51,8 +47,8 @@ export const validCommands: VisualizationRPCCommand[] = [
 export interface CommandResultMap {
   getFile: { fileBlob: Blob | null };
   getCodeToExecute: { code: string };
-  retry: void;
   setContentHeight: void;
+  setErrored: void;
 }
 
 // TODO(@fontanierh): refactor all these guards to use io-ts instead of manual checks.
@@ -100,12 +96,12 @@ export function isGetCodeToExecuteRequest(
   );
 }
 
-// Type guard for retry.
-export function isRetryRequest(
+// Type guard for setContentHeight.
+export function isSetContentHeightRequest(
   value: unknown
 ): value is VisualizationRPCRequest & {
-  command: "retry";
-  params: RetryParams;
+  command: "setContentHeight";
+  params: SetContentHeightParams;
 } {
   if (typeof value !== "object" || value === null) {
     return false;
@@ -114,21 +110,19 @@ export function isRetryRequest(
   const v = value as Partial<VisualizationRPCRequest>;
 
   return (
-    v.command === "retry" &&
+    v.command === "setContentHeight" &&
     typeof v.identifier === "string" &&
     typeof v.messageUniqueId === "string" &&
     typeof v.params === "object" &&
     v.params !== null &&
-    typeof (v.params as RetryParams).errorMessage === "string"
+    typeof (v.params as SetContentHeightParams).height === "number"
   );
 }
 
-// Type guard for setContentHeight.
-export function isSetContentHeightRequest(
+export function isSetErroredRequest(
   value: unknown
 ): value is VisualizationRPCRequest & {
-  command: "setContentHeight";
-  params: SetContentHeightParams;
+  command: "setErrored";
 } {
   if (typeof value !== "object" || value === null) {
     return false;
@@ -137,12 +131,9 @@ export function isSetContentHeightRequest(
   const v = value as Partial<VisualizationRPCRequest>;
 
   return (
-    v.command === "setContentHeight" &&
+    v.command === "setErrored" &&
     typeof v.identifier === "string" &&
-    typeof v.messageUniqueId === "string" &&
-    typeof v.params === "object" &&
-    v.params !== null &&
-    typeof (v.params as SetContentHeightParams).height === "number"
+    typeof v.messageUniqueId === "string"
   );
 }
 
@@ -156,7 +147,7 @@ export function isVisualizationRPCRequest(
   return (
     isGetCodeToExecuteRequest(value) ||
     isGetFileRequest(value) ||
-    isRetryRequest(value) ||
-    isSetContentHeightRequest(value)
+    isSetContentHeightRequest(value) ||
+    isSetErroredRequest(value)
   );
 }
diff --git a/viz/app/components/Components.tsx b/viz/app/components/Components.tsx
index 44f0653e5d77..6580f8b00fae 100644
--- a/viz/app/components/Components.tsx
+++ b/viz/app/components/Components.tsx
@@ -1,54 +1,5 @@
 "use client";
 
-// We can't use Sparkle components in the viz app,
-// because of client-side rendering issue.
-// So we define the components here.
-
-export const Button = ({
-  label,
-  onClick,
-}: {
-  label: string;
-  onClick: () => void;
-}) => {
-  return (
-    <button
-      onClick={onClick}
-      className="px-4 py-2 text-sm font-medium text-white bg-blue-500 rounded hover:bg-blue-600"
-    >
-      {label}
-    </button>
-  );
-};
-
-export const ErrorMessage = ({
-  children,
-  title,
-}: {
-  children: React.ReactNode;
-  title: string;
-}) => {
-  return (
-    <div className="bg-pink-100 border-l-4 border-pink-500 rounded-lg p-4 max-w-md">
-      <div className="flex items-center mb-2">
-        <svg
-          className="w-6 h-6 text-pink-500 mr-2"
-          fill="currentColor"
-          viewBox="0 0 20 20"
-        >
-          <path
-            fillRule="evenodd"
-            d="M18 10a8 8 0 11-16 0 8 8 0 0116 0zm-7-4a1 1 0 11-2 0 1 1 0 012 0zM9 9a1 1 0 000 2v3a1 1 0 001 1h1a1 1 0 100-2v-3a1 1 0 00-1-1H9z"
-            clipRule="evenodd"
-          />
-        </svg>
-        <h3 className="text-lg font-semibold text-pink-800">{title}</h3>
-      </div>
-      <div className="text-pink-700">{children}</div>
-    </div>
-  );
-};
-
 export const Spinner = () => {
   return (
     <div className="flex items-center justify-center">
diff --git a/viz/app/components/ErrorBoundary.tsx b/viz/app/components/ErrorBoundary.tsx
index 5bd8392f853b..ac8395753e56 100644
--- a/viz/app/components/ErrorBoundary.tsx
+++ b/viz/app/components/ErrorBoundary.tsx
@@ -1,5 +1,4 @@
-import { Button, ErrorMessage } from "@viz/app/components/Components";
-import React, { useState } from "react";
+import React from "react";
 
 interface ErrorBoundaryState {
   error: unknown;
@@ -8,8 +7,7 @@ interface ErrorBoundaryState {
 
 interface ErrorBoundaryProps {
   children: React.ReactNode;
-  errorMessage: string;
-  onRetryClick: (errorMessage: string) => void;
+  onErrored: () => void;
 }
 
 export class ErrorBoundary extends React.Component<
@@ -28,65 +26,14 @@ export class ErrorBoundary extends React.Component<
 
   componentDidCatch(error: unknown) {
     this.setState({ hasError: true, error });
+    this.props.onErrored();
   }
 
   render() {
     if (this.state.hasError) {
-      let error: Error;
-      if (this.state.error instanceof Error) {
-        error = this.state.error;
-      } else {
-        error = new Error("Unknown error.");
-      }
-
-      return (
-        <RenderError
-          error={error}
-          onRetryClick={this.props.onRetryClick}
-          message={this.props.errorMessage}
-        />
-      );
+      return null;
     }
 
     return <>{this.props.children}</>;
   }
 }
-
-// This is the component to render when an error occurs.
-export function RenderError({
-  error,
-  message,
-  onRetryClick,
-}: {
-  error: Error;
-  message: string;
-  onRetryClick: (errorMessage: string) => void;
-}) {
-  const [showDetails, setShowDetails] = useState(false);
-
-  return (
-    <div className="flex w-full flex-col items-center justify-center gap-4">
-      <ErrorMessage title="Error">
-        <>
-          {message}
-          <div className="mt-2">
-            <button
-              onClick={() => setShowDetails(!showDetails)}
-              className="text-pink-700 underline focus:outline-none"
-            >
-              {showDetails ? "Hide details" : "Show details"}
-            </button>
-            {showDetails && (
-              <div className="mt-2 p-2 bg-pink-50 rounded">
-                <strong>Error message:</strong> {error.message}
-              </div>
-            )}
-          </div>
-        </>
-      </ErrorMessage>
-      <div>
-        <Button label="Retry" onClick={() => onRetryClick(error.message)} />
-      </div>
-    </div>
-  );
-}
diff --git a/viz/app/components/VisualizationWrapper.tsx b/viz/app/components/VisualizationWrapper.tsx
index c6eaab43bc70..506b785bb78b 100644
--- a/viz/app/components/VisualizationWrapper.tsx
+++ b/viz/app/components/VisualizationWrapper.tsx
@@ -61,15 +61,6 @@ export function useVisualizationAPI(
     [sendCrossDocumentMessage]
   );
 
-  // This retry function sends a command to the host window requesting a retry of a previous
-  // operation, typically if the generated code fails.
-  const retry = useCallback(
-    async (errorMessage: string): Promise<void> => {
-      await sendCrossDocumentMessage("retry", { errorMessage });
-    },
-    [sendCrossDocumentMessage]
-  );
-
   const sendHeightToParent = useCallback(
     async ({ height }: { height: number | null }) => {
       if (height === null) {
@@ -83,7 +74,7 @@ export function useVisualizationAPI(
     [sendCrossDocumentMessage]
   );
 
-  return { fetchCode, fetchFile, error, retry, sendHeightToParent };
+  return { fetchCode, fetchFile, error, sendHeightToParent };
 }
 
 const useFile = (
@@ -134,9 +125,8 @@ export function VisualizationWrapperWithErrorBoundary({
 
   return (
     <ErrorBoundary
-      errorMessage="We encountered an error while running the code generated above. You can try again by clicking the button below."
-      onRetryClick={(errorMessage: string) => {
-        sendCrossDocumentMessage("retry", { errorMessage });
+      onErrored={() => {
+        sendCrossDocumentMessage("setErrored", undefined);
       }}
     >
       <VisualizationWrapper api={api} />
diff --git a/viz/package-lock.json b/viz/package-lock.json
index 609f8991358b..a6223223b559 100644
--- a/viz/package-lock.json
+++ b/viz/package-lock.json
@@ -16,7 +16,7 @@
         "react-dom": "^18",
         "react-resize-detector": "^11.0.1",
         "react-runner": "^1.0.5",
-        "recharts": "^2.12.7"
+        "recharts": "^2.13.0-alpha.4"
       },
       "devDependencies": {
         "@types/node": "^20",
@@ -4819,14 +4819,14 @@
       }
     },
     "node_modules/recharts": {
-      "version": "2.12.7",
-      "resolved": "https://registry.npmjs.org/recharts/-/recharts-2.12.7.tgz",
-      "integrity": "sha512-hlLJMhPQfv4/3NBSAyq3gzGg4h2v69RJh6KU7b3pXYNNAELs9kEoXOjbkxdXpALqKBoVmVptGfLpxdaVYqjmXQ==",
+      "version": "2.13.0-alpha.4",
+      "resolved": "https://registry.npmjs.org/recharts/-/recharts-2.13.0-alpha.4.tgz",
+      "integrity": "sha512-K9naL6F7pEcDYJE6yFQASSCQecSLPP0JagnvQ9hPtA/aHgsxsnIOjouLP5yrFZehxzfCkV5TEORr7/uNtSr7Qw==",
       "dependencies": {
         "clsx": "^2.0.0",
         "eventemitter3": "^4.0.1",
         "lodash": "^4.17.21",
-        "react-is": "^16.10.2",
+        "react-is": "^18.3.1",
         "react-smooth": "^4.0.0",
         "recharts-scale": "^0.4.4",
         "tiny-invariant": "^1.3.1",
@@ -4848,6 +4848,11 @@
         "decimal.js-light": "^2.4.1"
       }
     },
+    "node_modules/recharts/node_modules/react-is": {
+      "version": "18.3.1",
+      "resolved": "https://registry.npmjs.org/react-is/-/react-is-18.3.1.tgz",
+      "integrity": "sha512-/LLMVyas0ljjAtoYiPqYiL8VWXzUUdThrmU5+n20DZv+a+ClRoevUzw5JxU+Ieh5/c87ytoTBV9G1FiKfNJdmg=="
+    },
     "node_modules/reflect.getprototypeof": {
       "version": "1.0.6",
       "resolved": "https://registry.npmjs.org/reflect.getprototypeof/-/reflect.getprototypeof-1.0.6.tgz",
diff --git a/viz/package.json b/viz/package.json
index d3f568bde9ba..14ab17811f05 100644
--- a/viz/package.json
+++ b/viz/package.json
@@ -17,7 +17,7 @@
     "react-dom": "^18",
     "react-resize-detector": "^11.0.1",
     "react-runner": "^1.0.5",
-    "recharts": "^2.12.7"
+    "recharts": "^2.13.0-alpha.4"
   },
   "devDependencies": {
     "@types/node": "^20",