From 42d8c0617d6b3729782aef5053b4c2f67cadfd84 Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Wed, 4 Sep 2024 19:17:05 +0200
Subject: [PATCH] fix: fix theme in ai assit editor (#2212)

---
 frontend/src/components/editor/Output.tsx              | 10 +++++++++-
 frontend/src/components/editor/ai/add-cell-with-ai.tsx |  7 +++++++
 frontend/src/plugins/impl/code/any-language-editor.tsx |  2 ++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/frontend/src/components/editor/Output.tsx b/frontend/src/components/editor/Output.tsx
index e52c944159e..4c9bcf20f02 100644
--- a/frontend/src/components/editor/Output.tsx
+++ b/frontend/src/components/editor/Output.tsx
@@ -28,6 +28,7 @@ import { invariant } from "@/utils/invariant";
 import { CsvViewer } from "./file-tree/renderers";
 import { LazyAnyLanguageCodeMirror } from "@/plugins/impl/code/LazyAnyLanguageCodeMirror";
 import { MarimoTracebackOutput } from "./output/MarimoTracebackOutput";
+import { useTheme } from "@emotion/react";
 
 /**
  * Renders an output based on an OutputMessage.
@@ -37,6 +38,7 @@ export const OutputRenderer: React.FC<{
   onRefactorWithAI?: (opts: { prompt: string }) => void;
 }> = memo((props) => {
   const { message, onRefactorWithAI } = props;
+  const theme = useTheme();
 
   // Memoize parsing the json data
   const parsedJsonData = useMemo(() => {
@@ -125,7 +127,13 @@ export const OutputRenderer: React.FC<{
         typeof data === "string",
         `Expected string data for mime=${mimetype}. Got ${typeof data}`,
       );
-      return <LazyAnyLanguageCodeMirror value={data} language="markdown" />;
+      return (
+        <LazyAnyLanguageCodeMirror
+          theme={theme === "dark" ? "dark" : "light"}
+          value={data}
+          language="markdown"
+        />
+      );
     default:
       logNever(mimetype);
       return null;
diff --git a/frontend/src/components/editor/ai/add-cell-with-ai.tsx b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
index d08e9f2c5a4..a69a42df886 100644
--- a/frontend/src/components/editor/ai/add-cell-with-ai.tsx
+++ b/frontend/src/components/editor/ai/add-cell-with-ai.tsx
@@ -34,6 +34,7 @@ import {
 import { sql } from "@codemirror/lang-sql";
 import { SQLLanguageAdapter } from "@/core/codemirror/language/sql";
 import { atomWithStorage } from "jotai/utils";
+import { type ResolvedTheme, useTheme } from "@/theme/useTheme";
 
 const pythonExtensions = [
   customPythonLanguageSupport(),
@@ -89,6 +90,7 @@ export const AddCellWithAI: React.FC<{
   const { createNewCell } = useCellActions();
   const [completionBody, setCompletionBody] = useState<object>({});
   const [language, setLanguage] = useAtom(languageAtom);
+  const { theme } = useTheme();
 
   const {
     completion,
@@ -139,6 +141,7 @@ export const AddCellWithAI: React.FC<{
         </DropdownMenuContent>
       </DropdownMenu>
       <PromptInput
+        theme={theme}
         onClose={() => {
           setCompletion("");
           onClose();
@@ -207,6 +210,7 @@ export const AddCellWithAI: React.FC<{
           value={completion}
           className="cm border-t"
           onChange={setCompletion}
+          theme={theme === "dark" ? "dark" : "light"}
           extensions={language === "python" ? pythonExtensions : sqlExtensions}
         />
       )}
@@ -219,6 +223,7 @@ interface PromptInputProps {
   onClose: () => void;
   onChange: (value: string) => void;
   onSubmit: () => void;
+  theme: ResolvedTheme;
 }
 
 const PromptInput = ({
@@ -226,6 +231,7 @@ const PromptInput = ({
   onChange,
   onSubmit,
   onClose,
+  theme,
 }: PromptInputProps) => {
   const handleSubmit = onSubmit;
   const handleEscape = onClose;
@@ -327,6 +333,7 @@ const PromptInput = ({
       basicSetup={false}
       extensions={extensions}
       onChange={onChange}
+      theme={theme === "dark" ? "dark" : "light"}
       placeholder={"Generate with AI"}
     />
   );
diff --git a/frontend/src/plugins/impl/code/any-language-editor.tsx b/frontend/src/plugins/impl/code/any-language-editor.tsx
index 3919a95a20d..cba30623600 100644
--- a/frontend/src/plugins/impl/code/any-language-editor.tsx
+++ b/frontend/src/plugins/impl/code/any-language-editor.tsx
@@ -11,6 +11,7 @@ import {
 import React, { useMemo } from "react";
 import { Logger } from "@/utils/Logger";
 import { ErrorBanner } from "../common/error-banner";
+import type { ResolvedTheme } from "@/theme/useTheme";
 
 /**
  * A code editor that supports any language.
@@ -21,6 +22,7 @@ import { ErrorBanner } from "../common/error-banner";
 const AnyLanguageCodeMirror: React.FC<
   ReactCodeMirrorProps & {
     language: string;
+    theme: ResolvedTheme;
   }
 > = ({ language, extensions = [], ...props }) => {
   const isNotSupported = !(language in langs);