fix: update codeblock display

components/chat-message.tsx

@@ -29,47 +29,47 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
         {message.role === 'user' ? <IconUser /> : <IconOpenAI />}
       </div>
       <div className="ml-4 flex-1 space-y-2 px-1">
-        <div>
+        <MemoizedReactMarkdown
-          <MemoizedReactMarkdown
+          className="prose dark:prose-invert prose-p:leading-relaxed prose-pre:p-0"
-            className="prose leading-6 dark:prose-invert prose-p:leading-[1.8rem] prose-pre:rounded-md prose-pre:bg-[#282c34]"
+          remarkPlugins={[remarkGfm, remarkMath]}
-            remarkPlugins={[remarkGfm, remarkMath]}
+          components={{
-            components={{
+            p({ children }) {
-              p({ children }) {
+              return <p className="mb-2 last:mb-0">{children}</p>
-                return <p className="mb-2 last:mb-0">{children}</p>
+            },
-              },
+            code({ node, inline, className, children, ...props }) {
-              code({ node, inline, className, children, ...props }) {
+              if (children.length) {
-                if (children.length) {
+                if (children[0] == '▍') {
-                  if (children[0] == '▍') {
+                  return (
-                    return (
+                    <span className="mt-1 animate-pulse cursor-default">▍</span>
-                      <span className="mt-1 animate-pulse cursor-default">
+                  )
-                        ▍
-                      </span>
-                    )
-                  }
-
-                  children[0] = (children[0] as string).replace('`▍`', '▍')
                 }
 
-                const match = /language-(\w+)/.exec(className || '')
+                children[0] = (children[0] as string).replace('`▍`', '▍')
+              }
 
-                return !inline ? (
+              const match = /language-(\w+)/.exec(className || '')
-                  <CodeBlock
+
-                    key={Math.random()}
+              if (inline) {
-                    language={(match && match[1]) || ''}
+                return (
-                    value={String(children).replace(/\n$/, '')}
-                    {...props}
-                  />
-                ) : (
                   <code className={className} {...props}>
                     {children}
                   </code>
                 )
               }
-            }}
+
-          >
+              return (
-            {message.content}
+                <CodeBlock
-          </MemoizedReactMarkdown>
+                  key={Math.random()}
-        </div>
+                  language={(match && match[1]) || ''}
+                  value={String(children).replace(/\n$/, '')}
+                  {...props}
+                />
+              )
+            }
+          }}
+        >
+          {message.content}
+        </MemoizedReactMarkdown>
         <ChatMessageActions message={message} />
       </div>
     </div>

components/ui/codeblock.tsx

@@ -6,6 +6,7 @@ import { coldarkDark } from 'react-syntax-highlighter/dist/cjs/styles/prism'
 
 import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
 import { IconCheck, IconCopy, IconDownload } from '@/components/ui/icons'
+import { Button } from '@/components/ui/button'
 
 interface Props {
   language: string
@@ -51,8 +52,10 @@ export const generateRandomString = (length: number, lowercase = false) => {
   }
   return lowercase ? result.toLowerCase() : result
 }
+
 const CodeBlock: FC<Props> = memo(({ language, value }) => {
   const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 })
+
   const downloadAsFile = () => {
     if (typeof window === 'undefined') {
       return
@@ -65,7 +68,7 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
     const fileName = window.prompt('Enter file name' || '', suggestedFileName)
 
     if (!fileName) {
-      // user pressed cancel on prompt
+      // User pressed cancel on prompt.
       return
     }
 
@@ -80,37 +83,42 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
     document.body.removeChild(link)
     URL.revokeObjectURL(url)
   }
-  return (
-    <div className="codeblock relative w-full font-sans text-[16px]">
-      <div className="flex w-full items-center justify-between px-4 py-1.5">
-        <span className="text-xs lowercase text-white">{language}</span>
 
-        <div className="flex items-center">
+  return (
-          <button
+    <div className="codeblock relative w-full bg-zinc-950 font-sans">
-            type="button"
+      <div className="flex w-full items-center justify-between bg-zinc-900 px-6 py-2 pr-4 text-zinc-100">
-            className="flex items-center gap-1.5 rounded bg-none p-1 text-xs text-white"
+        <span className="text-xs lowercase">{language}</span>
+        <div className="flex items-center space-x-1">
+          <Button
+            variant="ghost"
+            className="hover:bg-zinc-800 focus-visible:ring-1 focus-visible:ring-slate-700 focus-visible:ring-offset-0"
+            onClick={downloadAsFile}
+            size="icon"
+          >
+            <IconDownload />
+            <span className="sr-only">Download</span>
+          </Button>
+          <Button
+            variant="ghost"
+            size="icon"
+            className="text-xs hover:bg-zinc-800 focus-visible:ring-1 focus-visible:ring-slate-700 focus-visible:ring-offset-0"
             onClick={() => copyToClipboard(value)}
           >
             {isCopied ? <IconCheck /> : <IconCopy />}
-            {isCopied ? 'Copied!' : 'Copy code'}
+            <span className="sr-only">Copy code</span>
-          </button>
+          </Button>
-          <button
-            type="button"
-            className="flex items-center rounded bg-none p-1 text-xs text-white"
-            onClick={downloadAsFile}
-          >
-            <IconDownload />
-          </button>
         </div>
       </div>
-
       <SyntaxHighlighter
         language={language}
         style={coldarkDark}
+        PreTag="div"
+        showLineNumbers
         customStyle={{
           margin: 0,
           width: '100%',
-          background: 'transparent'
+          background: 'transparent',
+          padding: '1.5rem 1rem'
         }}
         codeTagProps={{
           style: {