feat(markdown): useIsMarkdownCodeBlock hook (#586)

Yonom · Jul 26, 2024 · 5114a8f · 5114a8f
1 parent 53541d8
commit 5114a8f
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 3 deletions.
diff --git a/.changeset/lovely-dancers-collect.md b/.changeset/lovely-dancers-collect.md
@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react-markdown": patch
+---
+
+feat: useIsMarkdownCodeBlock hook
diff --git a/packages/react-markdown/src/index.ts b/packages/react-markdown/src/index.ts
@@ -8,6 +8,8 @@ export type {
   SyntaxHighlighterProps,
 } from "./overrides/types";
 
+export { useIsMarkdownCodeBlock } from "./overrides/PreOverride";
+
 export {
   makeMarkdownText,
   type MakeMarkdownTextProps,

diff --git a/packages/react-markdown/src/overrides/PreOverride.tsx b/packages/react-markdown/src/overrides/PreOverride.tsx
@@ -1,11 +1,15 @@
-import { createContext, ComponentPropsWithoutRef } from "react";
+import { createContext, ComponentPropsWithoutRef, useContext } from "react";
 import { PreComponent } from "./types";
 
 export const PreContext = createContext<Omit<
   ComponentPropsWithoutRef<PreComponent>,
   "children"
 > | null>(null);
 
+export const useIsMarkdownCodeBlock = () => {
+  return useContext(PreContext) !== null;
+};
+
 export const PreOverride: PreComponent = ({ children, ...rest }) => {
   return <PreContext.Provider value={rest}>{children}</PreContext.Provider>;
 };
diff --git a/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx b/packages/shadcn-registry/registry/assistant-ui/markdown-text.tsx
@@ -3,6 +3,7 @@
 import {
   CodeHeaderProps,
   MarkdownTextPrimitive,
+  useIsMarkdownCodeBlock,
 } from "@assistant-ui/react-markdown";
 import remarkGfm from "remark-gfm";
 import rehypeKatex from "rehype-katex";
@@ -155,11 +156,21 @@ const MarkdownTextImpl = () => {
             {...props}
           />
         ),
-        code: ({ node, className, ...props }) => {
+        pre: ({ node, className, ...props }) => (
+          <pre
+            className={cn(
+              "overflow-x-auto rounded-b-lg bg-black p-4 text-white",
+              className,
+            )}
+            {...props}
+          />
+        ),
+        code: function Code({ node, className, ...props }) {
+          const isCodeBlock = useIsMarkdownCodeBlock();
           return (
             <code
               className={cn(
-                "overflow-x-auto rounded-b-lg bg-black p-4 text-white",
+                !isCodeBlock && "bg-aui-muted rounded border font-semibold",
                 className,
               )}
               {...props}