From 5114a8f6529b44a885778956882cbc326dbe1443 Mon Sep 17 00:00:00 2001
From: Simon Farshid <simon.farshid@outlook.com>
Date: Fri, 26 Jul 2024 16:52:40 -0700
Subject: [PATCH] feat(markdown): useIsMarkdownCodeBlock hook (#586)

---
 .changeset/lovely-dancers-collect.md              |  5 +++++
 packages/react-markdown/src/index.ts              |  2 ++
 .../react-markdown/src/overrides/PreOverride.tsx  |  6 +++++-
 .../registry/assistant-ui/markdown-text.tsx       | 15 +++++++++++++--
 4 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/lovely-dancers-collect.md

diff --git a/.changeset/lovely-dancers-collect.md b/.changeset/lovely-dancers-collect.md
new file mode 100644
index 000000000..297498df8
--- /dev/null
+++ 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
index b69d03683..aad8c5877 100644
--- 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
index ed0c725c0..20766193b 100644
--- a/packages/react-markdown/src/overrides/PreOverride.tsx
+++ b/packages/react-markdown/src/overrides/PreOverride.tsx
@@ -1,4 +1,4 @@
-import { createContext, ComponentPropsWithoutRef } from "react";
+import { createContext, ComponentPropsWithoutRef, useContext } from "react";
 import { PreComponent } from "./types";
 
 export const PreContext = createContext<Omit<
@@ -6,6 +6,10 @@ export const PreContext = createContext<Omit<
   "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
index d6c4c7752..69f2a7aeb 100644
--- 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}