From b6834a9efe98e89d97369420ef1065cf404889cb Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Sat, 4 Jan 2025 15:15:06 -0500
Subject: [PATCH] fix: better component memozation (#3341)

---
 frontend/src/components/editor/Cell.tsx       |  9 +-
 .../editor/output/ConsoleOutput.tsx           | 83 +++++++++++++------
 2 files changed, 60 insertions(+), 32 deletions(-)

diff --git a/frontend/src/components/editor/Cell.tsx b/frontend/src/components/editor/Cell.tsx
index cd499714ad9..61fcac7350c 100644
--- a/frontend/src/components/editor/Cell.tsx
+++ b/frontend/src/components/editor/Cell.tsx
@@ -11,7 +11,6 @@ import {
   useRef,
   useState,
 } from "react";
-
 import { saveCellConfig, sendRun, sendStdin } from "@/core/network/requests";
 import { autocompletionKeymap } from "@/core/codemirror/cm";
 import type { UserConfig } from "../../core/config/config-schema";
@@ -520,6 +519,10 @@ const CellComponent = (
       : {}),
   });
 
+  const handleRefactorWithAI = useEvent((opts: { prompt: string }) => {
+    setAiCompletionCell({ cellId, initialPrompt: opts.prompt });
+  });
+
   if (!editing) {
     const outputIsError = isErrorMime(output?.mimetype);
     const hidden = errored || interrupted || stopped || outputIsError;
@@ -546,10 +549,6 @@ const CellComponent = (
     return undefined;
   };
 
-  const handleRefactorWithAI = (opts: { prompt: string }) => {
-    setAiCompletionCell({ cellId, initialPrompt: opts.prompt });
-  };
-
   const cellStatusComponent = (
     <CellStatusComponent
       status={status}
diff --git a/frontend/src/components/editor/output/ConsoleOutput.tsx b/frontend/src/components/editor/output/ConsoleOutput.tsx
index 93052fbac4a..262f4ecfdfe 100644
--- a/frontend/src/components/editor/output/ConsoleOutput.tsx
+++ b/frontend/src/components/editor/output/ConsoleOutput.tsx
@@ -75,12 +75,6 @@ export const ConsoleOutput = (props: Props): React.ReactNode => {
     return null;
   }
 
-  const renderText = (text: string) => {
-    return (
-      <span dangerouslySetInnerHTML={{ __html: ansiUp.ansi_to_html(text) }} />
-    );
-  };
-
   const reversedOutputs = [...consoleOutputs].reverse();
 
   return (
@@ -99,7 +93,6 @@ export const ConsoleOutput = (props: Props): React.ReactNode => {
         if (output.channel === "pdb") {
           return null;
         }
-        const originalIdx = consoleOutputs.length - idx - 1;
 
         if (output.channel === "stdin") {
           invariant(
@@ -107,31 +100,24 @@ export const ConsoleOutput = (props: Props): React.ReactNode => {
             "Expected data to be a string",
           );
 
+          const originalIdx = consoleOutputs.length - idx - 1;
+
           if (output.response == null) {
             return (
-              <div key={idx} className="flex gap-2 items-center">
-                {renderText(output.data)}
-                <Input
-                  data-testid="console-input"
-                  type="text"
-                  autoComplete="off"
-                  autoFocus={true}
-                  className="m-0"
-                  placeholder="stdin"
-                  onKeyDown={(e) => {
-                    if (e.key === "Enter" && !e.shiftKey) {
-                      onSubmitDebugger(e.currentTarget.value, originalIdx);
-                    }
-                  }}
-                />
-              </div>
+              <StdInput
+                key={idx}
+                output={output.data}
+                onSubmit={(text) => onSubmitDebugger(text, originalIdx)}
+              />
             );
           }
+
           return (
-            <div key={idx} className="flex gap-2 items-center">
-              {renderText(output.data)}
-              <span className="text-[var(--sky-11)]">{output.response}</span>
-            </div>
+            <StdInputWithResponse
+              key={idx}
+              output={output.data}
+              response={output.response}
+            />
           );
         }
 
@@ -152,3 +138,46 @@ export const ConsoleOutput = (props: Props): React.ReactNode => {
     </div>
   );
 };
+
+const StdInput = (props: {
+  onSubmit: (text: string) => void;
+  output: string;
+  response?: string;
+}) => {
+  return (
+    <div className="flex gap-2 items-center">
+      {renderText(props.output)}
+      <Input
+        data-testid="console-input"
+        type="text"
+        autoComplete="off"
+        autoFocus={true}
+        className="m-0"
+        placeholder="stdin"
+        onKeyDown={(e) => {
+          if (e.key === "Enter" && !e.shiftKey) {
+            props.onSubmit(e.currentTarget.value);
+          }
+        }}
+      />
+    </div>
+  );
+};
+
+const StdInputWithResponse = (props: {
+  output: string;
+  response?: string;
+}) => {
+  return (
+    <div className="flex gap-2 items-center">
+      {renderText(props.output)}
+      <span className="text-[var(--sky-11)]">{props.response}</span>
+    </div>
+  );
+};
+
+const renderText = (text: string) => {
+  return (
+    <span dangerouslySetInnerHTML={{ __html: ansiUp.ansi_to_html(text) }} />
+  );
+};