fix: update codeblock display

components/chat-message.tsx

@@ -29,9 +29,8 @@ 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
-            className="prose leading-6 dark:prose-invert prose-p:leading-[1.8rem] prose-pre:rounded-md prose-pre:bg-[#282c34]"
+          className="prose dark:prose-invert prose-p:leading-relaxed prose-pre:p-0"
           remarkPlugins={[remarkGfm, remarkMath]}
           components={{
             p({ children }) {
@@ -41,9 +40,7 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
               if (children.length) {
                 if (children[0] == '▍') {
                   return (
-                      <span className="mt-1 animate-pulse cursor-default">
+                    <span className="mt-1 animate-pulse cursor-default">▍</span>
-                        ▍
-                      </span>
                   )
                 }
 
@@ -52,24 +49,27 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
 
               const match = /language-(\w+)/.exec(className || '')
 
-                return !inline ? (
+              if (inline) {
+                return (
+                  <code className={className} {...props}>
+                    {children}
+                  </code>
+                )
+              }
+
+              return (
                 <CodeBlock
                   key={Math.random()}
                   language={(match && match[1]) || ''}
                   value={String(children).replace(/\n$/, '')}
                   {...props}
                 />
-                ) : (
-                  <code className={className} {...props}>
-                    {children}
-                  </code>
               )
             }
           }}
         >
           {message.content}
         </MemoizedReactMarkdown>
-        </div>
         <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: {