From 017788dbcc91bf07df03cc966ead275d7393effc Mon Sep 17 00:00:00 2001
From: Andrii Ieroshenko <ieroshenkoa@gmail.com>
Date: Mon, 8 Jan 2024 10:30:00 -0800
Subject: [PATCH] use rendermime directrly in the react component

---
 .../src/components/chat-messages.tsx          |  4 +-
 .../src/components/lumino-component.tsx       | 39 ------------
 .../src/components/markdown-component.tsx     | 59 -------------------
 .../src/components/rendermime-markdown.tsx    | 35 +++++++++++
 4 files changed, 37 insertions(+), 100 deletions(-)
 delete mode 100644 packages/jupyter-ai/src/components/lumino-component.tsx
 delete mode 100644 packages/jupyter-ai/src/components/markdown-component.tsx
 create mode 100644 packages/jupyter-ai/src/components/rendermime-markdown.tsx

diff --git a/packages/jupyter-ai/src/components/chat-messages.tsx b/packages/jupyter-ai/src/components/chat-messages.tsx
index 3a541249d..a2c78f9f8 100644
--- a/packages/jupyter-ai/src/components/chat-messages.tsx
+++ b/packages/jupyter-ai/src/components/chat-messages.tsx
@@ -7,7 +7,7 @@ import 'katex/dist/katex.min.css';
 import { AiService } from '../handler';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { Jupyternaut } from '../icons';
-import { MarkdownComponent } from './markdown-component';
+import { RendermimeMarkdown } from './rendermime-markdown';
 import { useCollaboratorsContext } from '../contexts/collaborators-context';
 
 type ChatMessagesProps = {
@@ -143,7 +143,7 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element {
             timestamp={timestamps[message.id]}
             sx={{ marginBottom: 3 }}
           />
-          <MarkdownComponent
+          <RendermimeMarkdown
             rmRegistry={props.rmRegistry}
             markdownStr={message.body}
           />
diff --git a/packages/jupyter-ai/src/components/lumino-component.tsx b/packages/jupyter-ai/src/components/lumino-component.tsx
deleted file mode 100644
index cb1767675..000000000
--- a/packages/jupyter-ai/src/components/lumino-component.tsx
+++ /dev/null
@@ -1,39 +0,0 @@
-import React, { useRef, useEffect } from 'react';
-import { MessageLoop } from '@lumino/messaging';
-import { Widget } from '@lumino/widgets';
-
-type LuminoComponentProps = {
-  widget: Widget;
-};
-
-export function LuminoComponent({
-  widget
-}: LuminoComponentProps): React.ReactElement {
-  const ref = useRef<HTMLDivElement>(null);
-
-  useEffect(() => {
-    if (!ref.current) {
-      return;
-    }
-
-    try {
-      MessageLoop.sendMessage(widget, Widget.Msg.BeforeAttach);
-      ref.current.appendChild(widget.node);
-      MessageLoop.sendMessage(widget, Widget.Msg.AfterAttach);
-    } catch (e) {
-      console.warn('Exception while attaching Lumino widget:', e);
-    }
-
-    return () => {
-      try {
-        if (widget.isAttached || widget.node.isConnected) {
-          Widget.detach(widget);
-        }
-      } catch (e) {
-        console.warn('Exception while detaching Lumino widget:', e);
-      }
-    };
-  }, [widget]);
-
-  return <div ref={ref} />;
-}
diff --git a/packages/jupyter-ai/src/components/markdown-component.tsx b/packages/jupyter-ai/src/components/markdown-component.tsx
deleted file mode 100644
index 8ece2fcc0..000000000
--- a/packages/jupyter-ai/src/components/markdown-component.tsx
+++ /dev/null
@@ -1,59 +0,0 @@
-import React, { useRef, useEffect } from 'react';
-import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
-import { Widget } from '@lumino/widgets';
-import { LuminoComponent } from './lumino-component';
-
-type MarkdownWidgetProps = {
-  markdownStr: string;
-  rmRegistry: IRenderMimeRegistry;
-};
-
-type MarkdownComponentProps = {
-  markdownStr: string;
-  rmRegistry: IRenderMimeRegistry;
-};
-
-class MarkdownWidget extends Widget {
-  private rmRegistry: IRenderMimeRegistry;
-  private markdownStr: string;
-
-  constructor(props: MarkdownWidgetProps) {
-    super();
-    this.rmRegistry = props.rmRegistry;
-    this.markdownStr = props.markdownStr;
-    this.initializeMarkdownRendering();
-  }
-
-  async initializeMarkdownRendering(): Promise<void> {
-    const mimeType = 'text/markdown';
-    const model = this.rmRegistry.createModel({
-      data: { [mimeType]: this.markdownStr }
-    });
-    const renderer = this.rmRegistry.createRenderer(mimeType);
-    await renderer.renderModel(model);
-    this.node.appendChild(renderer.node);
-  }
-}
-
-export function MarkdownComponent(
-  props: MarkdownComponentProps
-): React.ReactElement | null {
-  const widgetRef = useRef<MarkdownWidget | null>(null);
-
-  useEffect(() => {
-    if (!widgetRef.current) {
-      widgetRef.current = new MarkdownWidget({
-        markdownStr: props.markdownStr,
-        rmRegistry: props.rmRegistry
-      });
-    }
-
-    return () => {
-      widgetRef.current?.dispose();
-    };
-  }, []);
-
-  return widgetRef.current ? (
-    <LuminoComponent widget={widgetRef.current} />
-  ) : null;
-}
diff --git a/packages/jupyter-ai/src/components/rendermime-markdown.tsx b/packages/jupyter-ai/src/components/rendermime-markdown.tsx
new file mode 100644
index 000000000..33b79bf0a
--- /dev/null
+++ b/packages/jupyter-ai/src/components/rendermime-markdown.tsx
@@ -0,0 +1,35 @@
+import React, { useRef, useEffect } from 'react';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+
+type RendermimeMarkdownProps = {
+  markdownStr: string;
+  rmRegistry: IRenderMimeRegistry;
+};
+
+export function RendermimeMarkdown(
+  props: RendermimeMarkdownProps
+): JSX.Element {
+  const { markdownStr, rmRegistry } = props;
+  const ref = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    const renderMarkdown = async () => {
+      if (!ref.current) {
+        return;
+      }
+      ref.current.innerHTML = '';
+
+      const mimeType = 'text/markdown';
+      const model = rmRegistry.createModel({
+        data: { [mimeType]: markdownStr }
+      });
+      const renderer = rmRegistry.createRenderer(mimeType);
+      await renderer.renderModel(model);
+      ref.current.appendChild(renderer.node);
+    };
+
+    renderMarkdown();
+  }, [markdownStr, rmRegistry]);
+
+  return <div ref={ref} />;
+}