Convert TemplateEditor, VariableEditor, and PlaygroundTool (editor) t…

…o uncontrolled components

Workaround for uiwjs/react-codemirror#694

app/src/components/templateEditor/TemplateEditor.tsx

@@ -1,4 +1,4 @@
-import React, { useMemo } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { githubLight } from "@uiw/codemirror-theme-github";
 import { nord } from "@uiw/codemirror-theme-nord";
 import CodeMirror, {
@@ -16,8 +16,9 @@ import { MustacheLikeTemplating } from "./language/mustacheLike";
 import { TemplateLanguages } from "./constants";
 import { TemplateLanguage } from "./types";
 
-type TemplateEditorProps = ReactCodeMirrorProps & {
+type TemplateEditorProps = Omit<ReactCodeMirrorProps, "value"> & {
   templateLanguage: TemplateLanguage;
+  defaultValue: string;
 };
 
 const basicSetupOptions: BasicSetupOptions = {
@@ -28,10 +29,22 @@ const basicSetupOptions: BasicSetupOptions = {
   bracketMatching: false,
 };
 
+/**
+ * A template editor that is used to edit the template of a tool.
+ *
+ * This is an uncontrolled editor.
+ * You can only reset the value of the editor by triggering a re-mount, like with the `key` prop,
+ * or, when the readOnly prop is true, the editor will reset on all value changes.
+ * This is necessary because controlled react-codemirror editors incessantly reset
+ * cursor position when value is updated.
+ */
 export const TemplateEditor = ({
   templateLanguage,
+  defaultValue,
+  readOnly,
   ...props
 }: TemplateEditorProps) => {
+  const [value, setValue] = useState(() => defaultValue);
   const { theme } = useTheme();
   const codeMirrorTheme = theme === "light" ? githubLight : nord;
   const extensions = useMemo(() => {
@@ -51,12 +64,20 @@ export const TemplateEditor = ({
     return ext;
   }, [templateLanguage]);
 
+  useEffect(() => {
+    if (readOnly) {
+      setValue(defaultValue);
+    }
+  }, [readOnly, defaultValue]);
+
   return (
     <CodeMirror
       theme={codeMirrorTheme}
       extensions={extensions}
       basicSetup={basicSetupOptions}
+      readOnly={readOnly}
       {...props}
+      value={value}
     />
   );
 };

app/src/pages/playground/Playground.tsx

@@ -1,4 +1,4 @@
-import React, { Suspense, useCallback, useEffect } from "react";
+import React, { Fragment, Suspense, useCallback, useEffect } from "react";
 import { graphql, useLazyLoadQuery } from "react-relay";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 import { BlockerFunction, useBlocker, useSearchParams } from "react-router-dom";
@@ -205,6 +205,9 @@ const playgroundInputOutputPanelContentCSS = css`
  */
 const PLAYGROUND_PROMPT_PANEL_MIN_WIDTH = 632;
 
+const DEFAULT_EXPANDED_PROMPTS = ["prompts"];
+const DEFAULT_EXPANDED_PARAMS = ["input", "output"];
+
 function PlaygroundContent() {
   const instances = usePlaygroundContext((state) => state.instances);
   const templateLanguage = usePlaygroundContext(
@@ -246,7 +249,7 @@ function PlaygroundContent() {
   }, [isRunning, anyDirtyInstances]);
 
   return (
-    <>
+    <Fragment key="playground-content">
       <PanelGroup
         direction={
           isSingleInstance && !isDatasetMode ? "horizontal" : "vertical"
@@ -257,7 +260,7 @@ function PlaygroundContent() {
       >
         <Panel>
           <div css={playgroundPromptPanelContentCSS}>
-            <DisclosureGroup defaultExpandedKeys={["prompts"]}>
+            <DisclosureGroup defaultExpandedKeys={DEFAULT_EXPANDED_PROMPTS}>
               <Disclosure id="prompts" size="L">
                 <DisclosureTrigger
                   arrowPosition="start"
@@ -277,11 +280,10 @@ function PlaygroundContent() {
                       {instances.map((instance) => (
                         <View
                           flex="1 1 0px"
-                          key={instance.id}
+                          key={`${instance.id}-prompt`}
                           minWidth={PLAYGROUND_PROMPT_PANEL_MIN_WIDTH}
                         >
                           <PlaygroundTemplate
-                            key={instance.id}
                             playgroundInstanceId={instance.id}
                           />
                         </View>
@@ -301,7 +303,7 @@ function PlaygroundContent() {
             </Suspense>
           ) : (
             <div css={playgroundInputOutputPanelContentCSS}>
-              <DisclosureGroup defaultExpandedKeys={["input", "output"]}>
+              <DisclosureGroup defaultExpandedKeys={DEFAULT_EXPANDED_PARAMS}>
                 {templateLanguage !== TemplateLanguages.NONE ? (
                   <Disclosure id="input" size="L">
                     <DisclosureTrigger arrowPosition="start">
@@ -321,10 +323,9 @@ function PlaygroundContent() {
                   <DisclosurePanel>
                     <View padding="size-200" height="100%">
                       <Flex direction="row" gap="size-200">
-                        {instances.map((instance, i) => (
-                          <View key={i} flex="1 1 0px">
+                        {instances.map((instance) => (
+                          <View key={`${instance.id}-output`} flex="1 1 0px">
                             <PlaygroundOutput
-                              key={i}
                               playgroundInstanceId={instance.id}
                             />
                           </View>
@@ -348,6 +349,6 @@ function PlaygroundContent() {
           }
         />
       )}
-    </>
+    </Fragment>
   );
 }

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -199,6 +199,12 @@ function MessageEditor({
   updateMessage: (patch: Partial<ChatMessage>) => void;
   messageMode: MessageMode;
 }) {
+  const onChange = useCallback(
+    (val: string) => {
+      updateMessage({ content: val });
+    },
+    [updateMessage]
+  );
   if (messageMode === "toolCalls") {
     return (
       <View
@@ -262,18 +268,15 @@ function MessageEditor({
       </Form>
     );
   }
+
   return (
     <TemplateEditorWrap>
       <TemplateEditor
         height="100%"
-        value={
-          typeof message.content === "string"
-            ? message.content
-            : JSON.stringify(message.content, null, 2)
-        }
+        defaultValue={message.content || ""}
         aria-label="Message content"
         templateLanguage={templateLanguage}
-        onChange={(val) => updateMessage({ content: val })}
+        onChange={onChange}
         placeholder={
           message.role === "system"
             ? "You are a helpful assistant"

app/src/pages/playground/PlaygroundInput.tsx

@@ -56,7 +56,7 @@ export function PlaygroundInput() {
             // change rapidly for a given variable
             key={i}
             label={variableKey}
-            value={variablesMap[variableKey]}
+            defaultValue={variablesMap[variableKey] ?? ""}
             onChange={(value) => setVariableValue(variableKey, value)}
           />
         );

app/src/pages/playground/PlaygroundTool.tsx

@@ -32,12 +32,20 @@ import { PlaygroundInstanceProps } from "./types";
  */
 const TOOL_EDITOR_PRE_INIT_HEIGHT = 400;
 
+/**
+ * A tool editor that is used to edit the definition of a tool.
+ *
+ * This is a mostly un-controlled editor that re-mounts when the tool definition changes externally.
+ * This is necessary because controlled react-codemirror editors incessantly remount and reset
+ * cursor position when value is updated.
+ */
 export function PlaygroundTool({
   playgroundInstanceId,
   toolId,
 }: PlaygroundInstanceProps & {
   toolId: Tool["id"];
 }) {
+  const [version, setVersion] = useState(0);
   const updateInstance = usePlaygroundContext((state) => state.updateInstance);
 
   const instance = usePlaygroundContext((state) =>
@@ -56,7 +64,7 @@ export function PlaygroundTool({
     throw new Error(`Tool ${toolId} not found`);
   }
 
-  const [editorValue, setEditorValue] = useState(
+  const [initialEditorValue, setEditorValue] = useState(
     JSON.stringify(tool.definition, null, 2)
   );
 
@@ -71,12 +79,15 @@ export function PlaygroundTool({
     ) {
       setLastValidDefinition(tool.definition);
       setEditorValue(JSON.stringify(tool.definition, null, 2));
+      setVersion((prev) => prev + 1);
     }
   }, [tool.definition, lastValidDefinition]);
 
   const onChange = useCallback(
     (value: string) => {
-      setEditorValue(value);
+      // note that we do not update initialEditorValue here, we only want to update it when
+      // we are okay with the editor state resetting, which is basically only when the provider
+      // changes
       const { json: definition } = safelyParseJSON(value);
       if (definition == null) {
         return;
@@ -141,7 +152,7 @@ export function PlaygroundTool({
       bodyStyle={{ padding: 0 }}
       extra={
         <Flex direction="row" gap="size-100">
-          <CopyToClipboardButton text={editorValue} />
+          <CopyToClipboardButton text={initialEditorValue} />
           <Button
             aria-label="Delete tool"
             icon={<Icon svg={<Icons.TrashOutline />} />}
@@ -177,7 +188,10 @@ export function PlaygroundTool({
         preInitializationMinHeight={TOOL_EDITOR_PRE_INIT_HEIGHT}
       >
         <JSONEditor
-          value={editorValue}
+          // force remount of the editor when the tool definition changes externally
+          // usually when switching between providers
+          key={version}
+          value={initialEditorValue}
           onChange={onChange}
           jsonSchema={toolDefinitionJSONSchema}
         />

app/src/pages/playground/VariableEditor.tsx

@@ -1,4 +1,4 @@
-import React from "react";
+import React, { useEffect, useRef, useState } from "react";
 import { githubLight } from "@uiw/codemirror-theme-github";
 import { nord } from "@uiw/codemirror-theme-nord";
 import ReactCodeMirror, {
@@ -13,7 +13,7 @@ import { useTheme } from "@phoenix/contexts";
 
 type VariableEditorProps = {
   label?: string;
-  value?: string;
+  defaultValue: string;
   onChange?: (value: string) => void;
 };
 
@@ -28,20 +28,43 @@ const basicSetupOptions: BasicSetupOptions = {
 
 const extensions = [EditorView.lineWrapping];
 
+/**
+ * A mostly uncontrolled editor that re-mounts when the label changes.
+ *
+ * The re-mount ensures that value is reset to the initial value when the label (variable name) changes.
+ *
+ * This is necessary because controlled react-codemirror editors incessantly remount and reset
+ * cursor position when value is updated.
+ */
 export const VariableEditor = ({
   label,
-  value,
+  defaultValue,
   onChange,
 }: VariableEditorProps) => {
   const { theme } = useTheme();
+  const valueRef = useRef(defaultValue);
+  const [version, setVersion] = useState(0);
+  const [initialValue, setInitialValue] = useState(() => defaultValue);
+  useEffect(() => {
+    if (defaultValue == null) {
+      setInitialValue("");
+      setVersion((prev) => prev + 1);
+    }
+    valueRef.current = defaultValue;
+  }, [defaultValue]);
+  useEffect(() => {
+    setInitialValue(valueRef.current);
+    setVersion((prev) => prev + 1);
+  }, [label]);
   const codeMirrorTheme = theme === "light" ? githubLight : nord;
   return (
     <Field label={label}>
       <CodeWrap width="100%">
         <ReactCodeMirror
+          key={version}
           theme={codeMirrorTheme}
           basicSetup={basicSetupOptions}
-          value={value}
+          value={initialValue}
           extensions={extensions}
           onChange={onChange}
         />

app/src/pages/prompt/ChatTemplateMessageCard.tsx

@@ -69,7 +69,7 @@ export function ChatTemplateMessageToolResultPart({
             <TemplateEditor
               readOnly
               height="100%"
-              value={value}
+              defaultValue={value}
               templateLanguage={TemplateLanguages.NONE}
             />
           </TemplateEditorWrap>
@@ -103,7 +103,7 @@ export function ChatTemplateMessageToolCallPart({
         <TemplateEditor
           readOnly
           height="100%"
-          value={value}
+          defaultValue={value}
           templateLanguage={TemplateLanguages.NONE}
         />
       </TemplateEditorWrap>
@@ -130,7 +130,7 @@ export function ChatTemplateMessageTextPart(
         <TemplateEditor
           readOnly
           height="100%"
-          value={text}
+          defaultValue={text}
           templateLanguage={templateFormat}
         />
       </TemplateEditorWrap>