From 26aae5e4c7a8a63cdc66106cb3c30f55ca678d79 Mon Sep 17 00:00:00 2001
From: Parker Stafford <52351508+Parker-Stafford@users.noreply.github.com>
Date: Thu, 24 Oct 2024 12:57:46 -0700
Subject: [PATCH] fix: update to properly initialize when brought into view
 (#5172)

* fix: update  to properly initialize when brought into view

* add link to issue in comment

* update gap fix console warnings

* abstract out into wrapper

* remove log
---
 app/src/components/code/JSONEditor.tsx        | 10 ++-
 app/src/components/code/LazyEditorWrapper.tsx | 75 +++++++++++++++++++
 .../generative/ToolChoiceSelector.tsx         | 10 ++-
 app/src/pages/playground/PlaygroundTool.tsx   | 22 ++++--
 4 files changed, 103 insertions(+), 14 deletions(-)
 create mode 100644 app/src/components/code/LazyEditorWrapper.tsx

diff --git a/app/src/components/code/JSONEditor.tsx b/app/src/components/code/JSONEditor.tsx
index a7459c8afa..8da0793d17 100644
--- a/app/src/components/code/JSONEditor.tsx
+++ b/app/src/components/code/JSONEditor.tsx
@@ -28,6 +28,7 @@ export type JSONEditorProps = Omit<
 
 export function JSONEditor(props: JSONEditorProps) {
   const { theme } = useTheme();
+  const { jsonSchema, ...restProps } = props;
   const codeMirrorTheme = theme === "light" ? githubLight : nord;
   const extensions = useMemo(() => {
     const baseExtensions = [
@@ -35,25 +36,26 @@ export function JSONEditor(props: JSONEditorProps) {
       EditorView.lineWrapping,
       linter(jsonParseLinter()),
     ];
-    if (props.jsonSchema) {
+    if (jsonSchema) {
       baseExtensions.push(
         linter(jsonSchemaLinter(), { needsRefresh: handleRefresh }),
         jsonLanguage.data.of({
           autocomplete: jsonCompletion(),
         }),
         hoverTooltip(jsonSchemaHover()),
-        stateExtensions(props.jsonSchema)
+        stateExtensions(jsonSchema)
       );
     }
     return baseExtensions;
-  }, [props.jsonSchema]);
+  }, [jsonSchema]);
+
   return (
     <CodeMirror
       value={props.value}
       extensions={extensions}
       editable
       theme={codeMirrorTheme}
-      {...props}
+      {...restProps}
     />
   );
 }
diff --git a/app/src/components/code/LazyEditorWrapper.tsx b/app/src/components/code/LazyEditorWrapper.tsx
new file mode 100644
index 0000000000..d8306bbb7d
--- /dev/null
+++ b/app/src/components/code/LazyEditorWrapper.tsx
@@ -0,0 +1,75 @@
+import React, {
+  ReactNode,
+  useEffect,
+  useLayoutEffect,
+  useRef,
+  useState,
+} from "react";
+import { css } from "@emotion/react";
+
+/**
+ * A wrapper for code mirror editors that lazily initializes the editor when it is scrolled into view.
+ * This is necessary in some cases where a code mirror editor is rendered outside of the viewport.
+ * In those cases, the editor may not be initialized properly and may be invisible or cut off when it is scrolled into view.
+ * @param preInitializationMinHeight The minimum height of the container for the JSON editor prior to initialization.
+ */
+export function LazyEditorWrapper({
+  preInitializationMinHeight,
+  children,
+}: {
+  /**
+   * The minimum height of the container for the JSON editor prior to initialization.
+   * After initialization, the height will be set to auto and grow to fit the editor.
+   * This allows for the editor to properly get its dimensions when it is rendered outside of the viewport.
+   */
+  preInitializationMinHeight: number;
+  children: ReactNode;
+}) {
+  const [isVisible, setIsVisible] = useState(false);
+  const [isInitialized, setIsInitialized] = useState(false);
+  const wrapperRef = useRef<HTMLDivElement>(null);
+
+  /**
+   * The two useEffect hooks below are used to initialize the JSON editor.
+   * This is necessary because code mirror needs to calculate its dimensions to initialize properly.
+   * When it is rendered outside of the viewport, the dimensions may not always be calculated correctly,
+   * resulting in the editor being invisible or cut off when it is scrolled into view.
+   * Below we use a combination of an intersection observer and a delay to ensure that the editor is initialized correctly.
+   * For a related issue @see https://github.com/codemirror/dev/issues/1076
+   */
+  useEffect(() => {
+    const observer = new IntersectionObserver(([entry]) => {
+      setIsVisible(entry.isIntersecting);
+    });
+
+    if (wrapperRef.current) {
+      observer.observe(wrapperRef.current);
+    }
+    const current = wrapperRef.current;
+
+    return () => {
+      if (current) {
+        observer.unobserve(current);
+      }
+    };
+  }, []);
+
+  useLayoutEffect(() => {
+    if (isVisible && !isInitialized) {
+      setIsInitialized(true);
+    }
+  }, [isInitialized, isVisible]);
+
+  return (
+    <div
+      ref={wrapperRef}
+      css={css`
+        min-height: ${!isInitialized
+          ? `${preInitializationMinHeight}px`
+          : "auto"};
+      `}
+    >
+      {children}
+    </div>
+  );
+}
diff --git a/app/src/components/generative/ToolChoiceSelector.tsx b/app/src/components/generative/ToolChoiceSelector.tsx
index 32a0404fc3..e4a0091f5f 100644
--- a/app/src/components/generative/ToolChoiceSelector.tsx
+++ b/app/src/components/generative/ToolChoiceSelector.tsx
@@ -77,24 +77,26 @@ export function ToolChoicePicker({
       }}
     >
       {[
-        <Item key="auto">
+        <Item key="auto" textValue="auto">
           <Flex gap={"size-100"}>
             Tools auto-selected by LLM <Label>auto</Label>
           </Flex>
         </Item>,
-        <Item key="required">
+        <Item key="required" textValue="required">
           <Flex gap={"size-100"}>
             Use at least one tool <Label>required</Label>
           </Flex>
         </Item>,
-        <Item key="none">
+        <Item key="none" textValue="none">
           <Flex gap={"size-100"}>
             Don&apos;t use any tools <Label>none</Label>
           </Flex>
         </Item>,
         // Add "TOOL_NAME_PREFIX" prefix to user defined tool names to avoid conflicts with default keys
         ...toolNames.map((toolName) => (
-          <Item key={addToolNamePrefix(toolName)}>{toolName}</Item>
+          <Item key={addToolNamePrefix(toolName)} textValue={toolName}>
+            {toolName}
+          </Item>
         )),
       ]}
     </Picker>
diff --git a/app/src/pages/playground/PlaygroundTool.tsx b/app/src/pages/playground/PlaygroundTool.tsx
index ece8f9a065..3db1853939 100644
--- a/app/src/pages/playground/PlaygroundTool.tsx
+++ b/app/src/pages/playground/PlaygroundTool.tsx
@@ -5,6 +5,7 @@ import { Button, Card, Flex, Icon, Icons, Text } from "@arizeai/components";
 
 import { CopyToClipboardButton } from "@phoenix/components";
 import { JSONEditor } from "@phoenix/components/code";
+import { LazyEditorWrapper } from "@phoenix/components/code/LazyEditorWrapper";
 import { SpanKindIcon } from "@phoenix/components/trace";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { openAIToolJSONSchema, openAIToolSchema } from "@phoenix/schemas";
@@ -13,6 +14,12 @@ import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
 import { PlaygroundInstanceProps } from "./types";
 
+/**
+ * The minimum height for the editor before it is initialized.
+ * This is to ensure that the editor is properly initialized when it is rendered outside of the viewport.
+ */
+const TOOL_EDITOR_PRE_INIT_HEIGHT = 400;
+
 export function PlaygroundTool({
   playgroundInstanceId,
   tool,
@@ -65,7 +72,6 @@ export function PlaygroundTool({
       backgroundColor={"yellow-100"}
       borderColor={"yellow-700"}
       variant="compact"
-      key={tool.id}
       title={
         <Flex direction="row" gap="size-100">
           <SpanKindIcon spanKind="tool" />
@@ -94,11 +100,15 @@ export function PlaygroundTool({
         </Flex>
       }
     >
-      <JSONEditor
-        value={toolDefinition}
-        onChange={onChange}
-        jsonSchema={openAIToolJSONSchema as JSONSchema7}
-      />
+      <LazyEditorWrapper
+        preInitializationMinHeight={TOOL_EDITOR_PRE_INIT_HEIGHT}
+      >
+        <JSONEditor
+          value={toolDefinition}
+          onChange={onChange}
+          jsonSchema={openAIToolJSONSchema as JSONSchema7}
+        />
+      </LazyEditorWrapper>
     </Card>
   );
 }