feat: automatically set iframe height (#6512)

* feat: automatically resize iframe based on content size and fix re-render issues * tw styles * throttle instead of debounce --------- Co-authored-by: Henry Fontanier <henry@dust.tt>
dust-tt · Jul 25, 2024 · eba7c3b · eba7c3b
1 parent e02a2bf
commit eba7c3b
Show file tree

Hide file tree

Showing 7 changed files with 145 additions and 72 deletions.
diff --git a/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx b/front/components/assistant/conversation/actions/VisualizationActionIframe.tsx
@@ -9,6 +9,7 @@ import {
   isVisualizationRPCRequest,
   visualizationExtractCode,
 } from "@dust-tt/types";
+import type { SetStateAction } from "react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 
 import { RenderMessageMarkdown } from "@app/components/assistant/RenderMessageMarkdown";
@@ -33,8 +34,15 @@ const sendResponseToIframe = (
 // Custom hook to encapsulate the logic for handling visualization messages.
 function useVisualizationDataHandler(
   action: VisualizationActionType,
-  workspaceId: string,
-  onRetry: () => void
+  {
+    workspaceId,
+    onRetry,
+    setContentHeight,
+  }: {
+    workspaceId: string;
+    onRetry: () => void;
+    setContentHeight: (v: SetStateAction<number>) => void;
+  }
 ) {
   const extractedCode = useMemo(
     () => visualizationExtractCode(action.generation ?? ""),
@@ -87,14 +95,25 @@ function useVisualizationDataHandler(
           onRetry();
           break;
 
+        case "setContentHeight":
+          setContentHeight(data.params.height);
+          break;
+
         default:
           assertNever(data);
       }
     };
 
     window.addEventListener("message", listener);
     return () => window.removeEventListener("message", listener);
-  }, [action.generation, action.id, extractedCode, getFile, onRetry]);
+  }, [
+    action.generation,
+    action.id,
+    extractedCode,
+    getFile,
+    onRetry,
+    setContentHeight,
+  ]);
 
   return { getFile };
 }
@@ -115,10 +134,15 @@ export function VisualizationActionIframe({
 }) {
   const [activeTab, setActiveTab] = useState<"code" | "runtime">("code");
   const [tabManuallyChanged, setTabManuallyChanged] = useState(false);
+  const [contentHeight, setContentHeight] = useState(0);
 
   const workspaceId = owner.sId;
 
-  useVisualizationDataHandler(action, workspaceId, onRetry);
+  useVisualizationDataHandler(action, {
+    workspaceId,
+    onRetry,
+    setContentHeight,
+  });
 
   useEffect(() => {
     if (activeTab === "code" && action.generation && !tabManuallyChanged) {
@@ -165,11 +189,16 @@ export function VisualizationActionIframe({
         />
       )}
       {activeTab === "runtime" && (
-        <iframe
-          style={{ width: "100%", height: "600px" }}
-          src={`${process.env.NEXT_PUBLIC_VIZ_URL}/content?aId=${action.id}`}
-          sandbox="allow-scripts"
-        />
+        <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>
       )}
     </>
   );

diff --git a/front/lib/api/assistant/actions/visualization.ts b/front/lib/api/assistant/actions/visualization.ts
@@ -232,7 +232,7 @@ export class VisualizationConfigurationServerRunner extends BaseActionConfigurat
     const modelConversationRes = await renderConversationForModelMultiActions({
       conversation,
       model: agentModelConfig,
-      prompt: prompt,
+      prompt,
       allowedTokenCount: agentModelConfig.contextSize - MIN_GENERATION_TOKENS,
       excludeActions: false,
       excludeImages: true,

diff --git a/types/src/front/assistant/actions/visualization.ts b/types/src/front/assistant/actions/visualization.ts
@@ -61,6 +61,7 @@ export type VisualizationRPCRequestMap = {
   getFile: GetFileParams;
   getCodeToExecute: null;
   retry: RetryParams;
+  setContentHeight: { height: number };
 };
 
 // Derive the command type from the keys of the request map
@@ -78,6 +79,7 @@ export const validCommands: VisualizationRPCCommand[] = [
   "getFile",
   "getCodeToExecute",
   "retry",
+  "setContentHeight",
 ];
 
 // Command results.
@@ -86,6 +88,7 @@ export interface CommandResultMap {
   getFile: { file: File };
   getCodeToExecute: { code: string };
   retry: void;
+  setContentHeight: void;
 }
 
 // Type guard for getFile.
@@ -154,6 +157,29 @@ export function isRetryRequest(
   );
 }
 
+// Type guard for setContentHeight.
+export function isSetContentHeightRequest(
+  value: unknown
+): value is VisualizationRPCRequest & {
+  command: "setContentHeight";
+  params: { height: number };
+} {
+  if (typeof value !== "object" || value === null) {
+    return false;
+  }
+
+  const v = value as Partial<VisualizationRPCRequest>;
+
+  return (
+    v.command === "setContentHeight" &&
+    typeof v.actionId === "number" &&
+    typeof v.messageUniqueId === "string" &&
+    typeof v.params === "object" &&
+    v.params !== null &&
+    typeof (v.params as { height: number }).height === "number"
+  );
+}
+
 export function isVisualizationRPCRequest(
   value: unknown
 ): value is VisualizationRPCRequest {
@@ -164,6 +190,7 @@ export function isVisualizationRPCRequest(
   return (
     isGetCodeToExecuteRequest(value) ||
     isGetFileRequest(value) ||
-    isRetryRequest(value)
+    isRetryRequest(value) ||
+    isSetContentHeightRequest(value)
   );
 }
diff --git a/viz/app/components/VisualizationWrapper.tsx b/viz/app/components/VisualizationWrapper.tsx
@@ -11,8 +11,8 @@ import * as reactAll from "react";
 import React, { useCallback } from "react";
 import { useEffect, useState } from "react";
 import { importCode, Runner } from "react-runner";
-import {} from "react-runner";
 import * as rechartsAll from "recharts";
+import { useResizeDetector } from "react-resize-detector";
 
 export function useVisualizationAPI(actionId: number) {
   const [error, setError] = useState<Error | null>(null);
@@ -130,32 +130,75 @@ const useFile = (actionId: number, fileId: string) => {
 // This component renders the generated code.
 // It gets the generated code via message passing to the host window.
 export function VisualizationWrapper({ actionId }: { actionId: string }) {
-  const [code, setCode] = useState<string | null>(null);
+  type RunnerParams = {
+    code: string;
+    scope: Record<string, unknown>;
+  };
+
+  const [runnerParams, setRunnerParams] = useState<RunnerParams | null>(null);
+
   const [errored, setErrored] = useState<Error | null>(null);
   const actionIdParsed = parseInt(actionId, 10);
 
   const { fetchCode, error, retry } = useVisualizationAPI(actionIdParsed);
-  const useFileWrapped = (fileId: string) => useFile(actionIdParsed, fileId);
 
   useEffect(() => {
     const loadCode = async () => {
       try {
         const fetchedCode = await fetchCode();
-        if (fetchedCode) {
-          setCode(fetchedCode);
-        } else {
+        if (!fetchedCode) {
           setErrored(new Error("No visualization code found"));
+        } else {
+          setRunnerParams({
+            code: "() => {import Comp from '@dust/generated-code'; return (<Comp />);}",
+            scope: {
+              import: {
+                recharts: rechartsAll,
+                react: reactAll,
+                "@dust/generated-code": importCode(fetchedCode, {
+                  import: {
+                    recharts: rechartsAll,
+                    react: reactAll,
+                    papaparse: papaparseAll,
+                    "@dust/react-hooks": {
+                      useFile: (fileId: string) =>
+                        useFile(actionIdParsed, fileId),
+                    },
+                  },
+                }),
+              },
+            },
+          });
         }
       } catch (error) {
-        console.error(error);
         setErrored(new Error("Failed to fetch visualization code"));
       }
     };
 
     loadCode();
-  }, [fetchCode]);
+  }, [fetchCode, actionIdParsed]);
+
+  const sendHeightToParent = useCallback(
+    ({ height }: { height: number | null }) => {
+      if (height === null) {
+        return;
+      }
+      const sendHeight = makeIframeMessagePassingFunction<"setContentHeight">(
+        "setContentHeight",
+        actionIdParsed
+      );
+      sendHeight({ height });
+    },
+    [actionIdParsed]
+  );
+
+  const { ref } = useResizeDetector({
+    handleHeight: true,
+    refreshMode: "throttle",
+    refreshRate: 1000,
+    onResize: sendHeightToParent,
+  });
 
-  // Sync the Visualization API error with the local state.
   useEffect(() => {
     if (error) {
       setErrored(error);
@@ -171,45 +214,22 @@ export function VisualizationWrapper({ actionId }: { actionId: string }) {
     );
   }
 
-  if (!code) {
+  if (!runnerParams) {
     return <Spinner />;
   }
 
-  const generatedCodeScope = {
-    recharts: rechartsAll,
-    react: reactAll,
-    papaparse: papaparseAll,
-    "@dust/react-hooks": { useFile: useFileWrapped },
-  };
-
-  const scope = {
-    import: {
-      recharts: rechartsAll,
-      react: reactAll,
-      // Here we expose the code generated as a module to be imported by the wrapper code below.
-      "@dust/generated-code": importCode(code, { import: generatedCodeScope }),
-    },
-  };
-
-  // This code imports and renders the generated code.
-  const wrapperCode = `
-    () => {
-    import Comp from '@dust/generated-code';
-
-    return (<Comp />);
-    }
-  `;
-
   return (
-    <Runner
-      code={wrapperCode}
-      scope={scope}
-      onRendered={(error) => {
-        if (error) {
-          setErrored(error);
-        }
-      }}
-    />
+    <div ref={ref}>
+      <Runner
+        code={runnerParams.code}
+        scope={runnerParams.scope}
+        onRendered={(error) => {
+          if (error) {
+            setErrored(error);
+          }
+        }}
+      />
+    </div>
   );
 }
 

diff --git a/viz/app/globals.css b/viz/app/globals.css
@@ -10,23 +10,6 @@
   --background-end-rgb: 255, 255, 255;
 }
 
-@media (prefers-color-scheme: dark) {
-  :root {
-    --foreground-rgb: 255, 255, 255;
-    --background-start-rgb: 0, 0, 0;
-    --background-end-rgb: 0, 0, 0;
-  }
-}
-
-body {
-  color: rgb(var(--foreground-rgb));
-  background: linear-gradient(
-      to bottom,
-      transparent,
-      rgb(var(--background-end-rgb))
-    )
-    rgb(var(--background-start-rgb));
-}
 
 @layer utilities {
   .text-balance {

diff --git a/viz/package-lock.json b/viz/package-lock.json
diff --git a/viz/package.json b/viz/package.json
@@ -15,6 +15,7 @@
     "papaparse": "^5.4.1",
     "react": "^18",
     "react-dom": "^18",
+    "react-resize-detector": "^11.0.1",
     "react-runner": "^1.0.5",
     "recharts": "^2.12.7"
   },