diff --git a/render-markdown-codehighlight/hooks.ts b/render-markdown-codehighlight/hooks.ts
index 99bf52af..1cec6808 100644
--- a/render-markdown-codehighlight/hooks.ts
+++ b/render-markdown-codehighlight/hooks.ts
@@ -19,10 +19,12 @@
 
 import { useEffect, useState } from 'react';
 import hljs from 'highlight.js';
-import { themeStyles } from './themeStyles'
+import { themeStyles } from './themeStyles';
 
 
-const useHighlightCode = (element: HTMLElement | null) => {
+const useHighlightCode = (props: HTMLElement | null | {
+  current: HTMLElement | null;
+}) => {
   const [selectTheme, setSelectTheme] = useState<string>('default');
 
   // Fetch theme from API
@@ -39,7 +41,14 @@ const useHighlightCode = (element: HTMLElement | null) => {
   }, []);
 
   useEffect(() => {
-    if (!element) return;
+    let element;
+    if (props instanceof HTMLElement) {
+      element = props;
+    } else if (props && props.current instanceof HTMLElement) {
+      element = props.current;
+    } else {
+      return;
+    }
 
     const applyThemeCSS = async (theme: string) => {
       const existingStyleElement = document.querySelector('style[data-theme-style="highlight"]');
@@ -96,10 +105,9 @@ const useHighlightCode = (element: HTMLElement | null) => {
       contentObserver.disconnect();
       themeObserver.disconnect();
     };
-  }, [element, selectTheme]);
+  }, [props, selectTheme]);
 
   return null;
 };
 
 export { useHighlightCode };
-
diff --git a/render-markdown-codehighlight/info.yaml b/render-markdown-codehighlight/info.yaml
index 10cdb8f1..b6fa9a96 100644
--- a/render-markdown-codehighlight/info.yaml
+++ b/render-markdown-codehighlight/info.yaml
@@ -17,6 +17,6 @@
 
 slug_name: render_markdown_codehighlight
 type: render
-version: 0.0.3
+version: 0.0.4
 author: Chen Jiaji, Zhu Xuanlyu
 link: https://github.com/apache/incubator-answer-plugins/tree/main/render-markdown-codehighlight
diff --git a/render-markdown-codehighlight/package.json b/render-markdown-codehighlight/package.json
index 147088ad..1f5d2281 100644
--- a/render-markdown-codehighlight/package.json
+++ b/render-markdown-codehighlight/package.json
@@ -1,6 +1,6 @@
 {
   "name": "render-markdown-codehighlight",
-  "version": "0.0.3",
+  "version": "0.0.4",
   "description": "",
   "type": "module",
   "files": [