diff --git a/app/packages/rehype-momiji/buildCodeBlock.ts b/app/packages/rehype-momiji/buildCodeBlock.ts
index d602c18..78432aa 100644
--- a/app/packages/rehype-momiji/buildCodeBlock.ts
+++ b/app/packages/rehype-momiji/buildCodeBlock.ts
@@ -1,6 +1,6 @@
 import type { BundledLanguage, HighlighterGeneric, BundledTheme } from "shiki";
-import { toHtml } from "hast-util-to-html";
 import { COLOR_SHIRONERI, COLOR_HAI } from "./colors";
+import type { Root } from "hast";
 
 const buildCodeBlock = (
   highlighter: HighlighterGeneric<BundledLanguage, BundledTheme>,
@@ -10,26 +10,25 @@ const buildCodeBlock = (
   filename: string,
   filenameBGColor?: string,
   filenameTextColor?: string,
-) => {
-  const filenameColorStyle = `background-color: ${filenameBGColor ?? COLOR_HAI}; color: ${filenameTextColor ?? COLOR_SHIRONERI};`;
-
+): Root => {
   const hast = highlighter.codeToHast(rawCode, {
     lang: lang,
     theme: theme,
     transformers: [
       {
         pre(this, node) {
-          if (filename === "") return;
-
+          // Add filename to the code block
           if (
+            filename !== "" &&
             node.properties.class &&
             typeof node.properties.class === "string" &&
             node.properties.class.includes("shiki")
           ) {
+            const filenameColorStyle = `background-color: ${filenameBGColor ?? COLOR_HAI}; color: ${filenameTextColor ?? COLOR_SHIRONERI};`;
+
             node.properties = {
-              class: node.properties.class,
-              style: `${node.properties.style}; padding-top: 48px;`,
-              tabindex: node.properties.tabindex,
+              ...node.properties,
+              style: `${node.properties.style}; padding-top: 48px; position: relative;`,
             };
 
             node.children.unshift({
@@ -51,7 +50,7 @@ const buildCodeBlock = (
     ],
   });
 
-  return toHtml(hast);
+  return hast;
 };
 
 export { buildCodeBlock };
diff --git a/app/packages/rehype-momiji/rehypeMomiji.ts b/app/packages/rehype-momiji/index.ts
similarity index 75%
rename from app/packages/rehype-momiji/rehypeMomiji.ts
rename to app/packages/rehype-momiji/index.ts
index 7e0b747..019def6 100644
--- a/app/packages/rehype-momiji/rehypeMomiji.ts
+++ b/app/packages/rehype-momiji/index.ts
@@ -1,10 +1,8 @@
 import { bundledLanguages, bundledThemes, type BundledLanguage, type BundledTheme } from "shiki";
-import type { Node } from "unist";
 import type { Plugin } from "unified";
-import type { Text, Element, Nodes } from "hast";
+import type { ElementContent, Root } from "hast";
 import { visit } from "unist-util-visit";
 import { getHighlighter } from "./highlighter";
-import { parser } from "./parser";
 import { buildCodeBlock } from "./buildCodeBlock";
 
 interface Option {
@@ -17,7 +15,7 @@ interface Option {
 
 const defaultHighlighter = await getHighlighter({ themes: bundledThemes, langs: bundledLanguages });
 
-const rehypeMomiji: Plugin<Option[]> = (options) => {
+const rehypeMomiji: Plugin<Option[], Root> = (options) => {
   const {
     theme = "github-dark-default",
     fallbackLang = "c",
@@ -38,20 +36,27 @@ const rehypeMomiji: Plugin<Option[]> = (options) => {
   };
 
   const checkSupportedLanguage = (lang: string): string => {
-    if (lang && langs.includes(lang)) return lang;
+    if (langs.includes(lang)) return lang;
     return fallbackLang;
   };
 
-  return (tree: Node) => {
-    visit(tree as Nodes, "element", (node) => {
+  return (tree) => {
+    visit(tree, "element", (node) => {
       // Check if the node is a pre tag with a single child
       if (!(node.tagName === "pre" && Array.isArray(node.children) && node.children.length === 1)) {
         return;
       }
 
       // Check if the child is a code tag
-      const codeElem = node.children[0] as Element;
-      if (!(codeElem !== null && typeof codeElem === "object" && codeElem.tagName === "code")) {
+      const [codeElem] = node.children;
+      if (
+        !(
+          codeElem.type === "element" &&
+          codeElem !== null &&
+          typeof codeElem === "object" &&
+          codeElem.tagName === "code"
+        )
+      ) {
         return;
       }
 
@@ -68,8 +73,8 @@ const rehypeMomiji: Plugin<Option[]> = (options) => {
       }
 
       // Check if the code tag has a text child
-      const textNode = codeElem.children[0] as Text;
-      if (!(typeof textNode.value === "string")) {
+      const [textNode] = codeElem.children;
+      if (!(textNode.type === "text" && typeof textNode.value === "string")) {
         return;
       }
 
@@ -100,15 +105,8 @@ const rehypeMomiji: Plugin<Option[]> = (options) => {
         filenameTextColor,
       );
 
-      const container = `
-        <div style="position: relative; display: flex; flex-direction: column; gap: 2px;">
-          ${codeBlock}
-        </div>
-      `;
-
-      const parsedRoot = parser.parse(container);
       node.tagName = "div";
-      node.children = parsedRoot.children as Element[];
+      node.children = codeBlock.children as ElementContent[];
     });
   };
 };
diff --git a/app/packages/rehype-momiji/parser.ts b/app/packages/rehype-momiji/parser.ts
deleted file mode 100644
index 9d82f1d..0000000
--- a/app/packages/rehype-momiji/parser.ts
+++ /dev/null
@@ -1,6 +0,0 @@
-import { unified } from "unified";
-import rehypeParse from "rehype-parse";
-
-const parser = unified().use(rehypeParse, { fragment: true });
-
-export { parser };
diff --git a/vite.config.ts b/vite.config.ts
index e658304..7f04ef6 100644
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -8,7 +8,7 @@ import remarkBreaks from 'remark-breaks';
 import remarkFrontmatter from 'remark-frontmatter';
 import remarkMdxFrontmatter from 'remark-mdx-frontmatter';
 import remarkGfm from 'remark-gfm'
-import rehypeMomiji from './app/packages/rehype-momiji/rehypeMomiji'
+import rehypeMomiji from './app/packages/rehype-momiji'
 import remarkMomijiCodeFilename from './app/packages/remark-momiji-filename'
 import rehypeMermaid from './app/packages/rehype-mermaid/rehypeMermaid'