@gram-ai/elements 1.23.0 → 1.24.0

package/src/components/assistant-ui/thread.tsx

@@ -4,7 +4,9 @@ import {
   CheckIcon,
   ChevronLeftIcon,
   ChevronRightIcon,
+  CircleIcon,
   CopyIcon,
+  DownloadIcon,
   PencilIcon,
   Settings2,
   Square,
@@ -18,11 +20,20 @@ import {
   ImageMessagePartProps,
   MessagePrimitive,
   ThreadPrimitive,
+  useAssistantState,
 } from '@assistant-ui/react'
 
 import { LazyMotion, MotionConfig, domAnimation } from 'motion/react'
 import * as m from 'motion/react-m'
-import { useEffect, useRef, useState, type FC } from 'react'
+import {
+  createContext,
+  useContext,
+  useEffect,
+  useMemo,
+  useRef,
+  useState,
+  type FC,
+} from 'react'
 import { AnimatePresence } from 'motion/react'
 
 import {
@@ -30,7 +41,9 @@ import {
   ComposerAttachments,
   UserMessageAttachments,
 } from '@/components/assistant-ui/attachment'
+import { FollowOnSuggestions } from '@/components/assistant-ui/follow-on-suggestions'
 import { MarkdownText } from '@/components/assistant-ui/markdown-text'
+import { MessageFeedback } from '@/components/assistant-ui/message-feedback'
 import { Reasoning, ReasoningGroup } from '@/components/assistant-ui/reasoning'
 import { ToolFallback } from '@/components/assistant-ui/tool-fallback'
 import { ToolMentionAutocomplete } from '@/components/assistant-ui/tool-mention-autocomplete'
@@ -46,6 +59,8 @@ import { useThemeProps } from '@/hooks/useThemeProps'
 import { EASE_OUT_QUINT } from '@/lib/easing'
 import { MODELS } from '@/lib/models'
 import { cn } from '@/lib/utils'
+import { useRecordCassette } from '@/hooks/useRecordCassette'
+import { useReplayContext } from '@/contexts/ReplayContext'
 import { ConnectionStatusIndicatorSafe } from './connection-status-indicator'
 import { Popover, PopoverContent, PopoverTrigger } from '../ui/popover'
 import {
@@ -56,6 +71,23 @@ import {
 } from '../ui/tooltip'
 import { ToolGroup } from './tool-group'
 
+// Context for chat resolution state
+const ChatResolutionContext = createContext<{
+  isResolved: boolean
+  feedbackHidden: boolean
+  setResolved: () => void
+  setUnresolved: () => void
+  resetFeedbackHidden: () => void
+}>({
+  isResolved: false,
+  feedbackHidden: false,
+  setResolved: () => {},
+  setUnresolved: () => {},
+  resetFeedbackHidden: () => {},
+})
+
+const useChatResolution = () => useContext(ChatResolutionContext)
+
 const StaticSessionWarning = () => (
   <div className="m-2 rounded-md border border-amber-500 bg-amber-100 px-4 py-3 text-sm text-amber-800 dark:border-amber-600 dark:bg-amber-900/30 dark:text-amber-200">
     <strong>Warning:</strong> You are using a static session token in the
@@ -81,52 +113,89 @@ export const Thread: FC<ThreadProps> = ({ className }) => {
   const { config } = useElements()
   const components = config.components ?? {}
   const showStaticSessionWarning = config.api && 'sessionToken' in config.api
+  const showFeedback = config.thread?.experimental_showFeedback ?? false
+  const [isResolved, setIsResolved] = useState(false)
+  const [feedbackHidden, setFeedbackHidden] = useState(false)
+
+  const setResolved = () => setIsResolved(true)
+  const setUnresolved = () => {
+    setIsResolved(false)
+    setFeedbackHidden(true)
+  }
+  const resetFeedbackHidden = () => setFeedbackHidden(false)
 
   return (
-    <LazyMotion features={domAnimation}>
-      <MotionConfig reducedMotion="user">
-        <ThreadPrimitive.Root
-          className={cn(
-            'aui-root aui-thread-root bg-background @container relative flex h-full flex-col',
-            themeProps.className,
-            className
-          )}
-        >
-          <ConnectionStatusIndicatorSafe />
-          <ThreadPrimitive.Viewport
+    <ChatResolutionContext.Provider
+      value={{
+        isResolved: showFeedback && isResolved,
+        feedbackHidden,
+        setResolved,
+        setUnresolved,
+        resetFeedbackHidden,
+      }}
+    >
+      <LazyMotion features={domAnimation}>
+        <MotionConfig reducedMotion="user">
+          <ThreadPrimitive.Root
             className={cn(
-              'aui-thread-viewport relative mx-auto flex w-full flex-1 flex-col overflow-x-auto overflow-y-scroll pb-0!',
-              d('p-lg')
+              'aui-root aui-thread-root bg-background @container relative flex h-full flex-col',
+              themeProps.className,
+              className
             )}
           >
-            <ThreadPrimitive.If empty>
-              {components.ThreadWelcome ? (
-                <components.ThreadWelcome />
-              ) : (
-                <ThreadWelcome />
+            <ConnectionStatusIndicatorSafe />
+            <ThreadPrimitive.Viewport
+              className={cn(
+                'aui-thread-viewport relative mx-auto flex w-full flex-1 flex-col overflow-x-auto overflow-y-scroll pb-0!',
+                d('p-lg')
+              )}
+            >
+              <ThreadPrimitive.If empty>
+                {components.ThreadWelcome ? (
+                  <components.ThreadWelcome />
+                ) : (
+                  <ThreadWelcome />
+                )}
+              </ThreadPrimitive.If>
+
+              {showStaticSessionWarning && <StaticSessionWarning />}
+
+              <ThreadPrimitive.Messages
+                components={{
+                  UserMessage: components.UserMessage ?? UserMessage,
+                  EditComposer: components.EditComposer ?? EditComposer,
+                  AssistantMessage:
+                    components.AssistantMessage ?? AssistantMessage,
+                }}
+              />
+
+              <ThreadPrimitive.If empty={false} running={false}>
+                <FollowOnSuggestions />
+              </ThreadPrimitive.If>
+
+              <ThreadPrimitive.If empty={false}>
+                <div className="aui-thread-viewport-spacer min-h-8 grow" />
+              </ThreadPrimitive.If>
+
+              <Composer showFeedback={showFeedback} />
+            </ThreadPrimitive.Viewport>
+
+            {/* Resolution overlay - subtle readonly effect */}
+            <AnimatePresence>
+              {showFeedback && isResolved && (
+                <m.div
+                  className="bg-background/40 pointer-events-none absolute inset-0 z-50"
+                  initial={{ opacity: 0 }}
+                  animate={{ opacity: 1 }}
+                  exit={{ opacity: 0 }}
+                  transition={{ duration: 0.3, ease: EASE_OUT_QUINT }}
+                />
               )}
-            </ThreadPrimitive.If>
-
-            {showStaticSessionWarning && <StaticSessionWarning />}
-
-            <ThreadPrimitive.Messages
-              components={{
-                UserMessage: components.UserMessage ?? UserMessage,
-                EditComposer: components.EditComposer ?? EditComposer,
-                AssistantMessage:
-                  components.AssistantMessage ?? AssistantMessage,
-              }}
-            />
-
-            <ThreadPrimitive.If empty={false}>
-              <div className="aui-thread-viewport-spacer min-h-8 grow" />
-            </ThreadPrimitive.If>
-
-            <Composer />
-          </ThreadPrimitive.Viewport>
-        </ThreadPrimitive.Root>
-      </MotionConfig>
-    </LazyMotion>
+            </AnimatePresence>
+          </ThreadPrimitive.Root>
+        </MotionConfig>
+      </LazyMotion>
+    </ChatResolutionContext.Provider>
   )
 }
 
@@ -357,10 +426,57 @@ const ComposerToolMentions: FC<{
   )
 }
 
-const Composer: FC = () => {
+// Resets feedbackHidden when a new message starts generating
+const FeedbackHiddenResetter: FC = () => {
+  const { resetFeedbackHidden } = useChatResolution()
+
+  useEffect(() => {
+    resetFeedbackHidden()
+  }, [resetFeedbackHidden])
+
+  return null
+}
+
+const ComposerFeedback: FC = () => {
+  const { isResolved, feedbackHidden, setResolved } = useChatResolution()
+
+  return (
+    <ThreadPrimitive.If empty={false}>
+      {/* Reset feedbackHidden when a new message starts generating */}
+      <ThreadPrimitive.If running>
+        <FeedbackHiddenResetter />
+      </ThreadPrimitive.If>
+      <ThreadPrimitive.If running={false}>
+        <AnimatePresence>
+          {!isResolved && !feedbackHidden && (
+            <m.div
+              initial={{ opacity: 0, y: 10 }}
+              animate={{ opacity: 1, y: 0 }}
+              exit={{ opacity: 0, y: 10 }}
+              transition={{ duration: 0.2, ease: EASE_OUT_QUINT }}
+              className="mb-3"
+            >
+              <MessageFeedback className="mx-auto" onResolved={setResolved} />
+            </m.div>
+          )}
+        </AnimatePresence>
+      </ThreadPrimitive.If>
+    </ThreadPrimitive.If>
+  )
+}
+
+interface ComposerProps {
+  showFeedback?: boolean
+}
+
+const Composer: FC<ComposerProps> = ({ showFeedback = false }) => {
   const { config, mcpTools } = useElements()
+  const { isResolved, setUnresolved } = useChatResolution()
   const r = useRadius()
   const d = useDensity()
+  const replayCtx = useReplayContext()
+
+  const isReplay = replayCtx?.isReplay ?? false
   const composerConfig = config.composer ?? {
     placeholder: 'Send a message...',
     attachments: true,
@@ -383,37 +499,61 @@ const Composer: FC = () => {
   return (
     <div
       className={cn(
-        'aui-composer-wrapper bg-background sticky bottom-0 flex w-full flex-col overflow-visible',
+        'aui-composer-wrapper bg-background sticky bottom-0 z-[60] flex w-full flex-col overflow-visible',
         d('gap-md'),
         d('py-md'),
         r('xl')
       )}
     >
+      {showFeedback && <ComposerFeedback />}
       <ThreadScrollToBottom />
-      <ComposerPrimitive.Root
-        ref={composerRootRef}
-        className={cn(
-          'aui-composer-root group/input-group border-input bg-background has-[textarea:focus-visible]:border-ring has-[textarea:focus-visible]:ring-ring/5 dark:bg-background relative flex w-full flex-col border px-1 pt-2 shadow-xs transition-[color,box-shadow] outline-none has-[textarea:focus-visible]:ring-1',
-          r('xl')
-        )}
-      >
-        {composerConfig.attachments && <ComposerAttachments />}
-
-        {toolMentionsEnabled && <ComposerToolMentions tools={mcpTools} />}
-
-        <ComposerPrimitive.Input
-          placeholder={composerConfig.placeholder}
+      {showFeedback && isResolved ? (
+        <m.div
+          className="aui-composer-resolved border-input flex min-h-[118px] flex-col items-center justify-center gap-2 border-t px-1"
+          initial={{ opacity: 0 }}
+          animate={{ opacity: 1 }}
+          transition={{ duration: 0.2, ease: EASE_OUT_QUINT }}
+        >
+          <span className="text-muted-foreground text-sm">
+            This conversation has been resolved
+          </span>
+          <Button
+            variant="outline"
+            size="sm"
+            className="text-foreground"
+            onClick={setUnresolved}
+          >
+            Reopen conversation
+          </Button>
+        </m.div>
+      ) : (
+        <ComposerPrimitive.Root
+          ref={composerRootRef}
           className={cn(
-            'aui-composer-input placeholder:text-muted-foreground mb-1 max-h-32 w-full resize-none bg-transparent px-3.5 pt-1.5 pb-3 outline-none focus-visible:ring-0',
-            d('h-input'),
-            d('text-base')
+            'aui-composer-root group/input-group border-input bg-background has-[textarea:focus-visible]:border-ring has-[textarea:focus-visible]:ring-ring/5 dark:bg-background relative flex min-h-[118px] w-full flex-col border px-1 pt-2 shadow-xs transition-[color,box-shadow] outline-none has-[textarea:focus-visible]:ring-1',
+            r('xl'),
+            isReplay && 'pointer-events-none opacity-50'
           )}
-          rows={1}
-          autoFocus
-          aria-label="Message input"
-        />
-        <ComposerAction />
-      </ComposerPrimitive.Root>
+        >
+          {composerConfig.attachments && <ComposerAttachments />}
+
+          {toolMentionsEnabled && <ComposerToolMentions tools={mcpTools} />}
+
+          <ComposerPrimitive.Input
+            placeholder={composerConfig.placeholder}
+            className={cn(
+              'aui-composer-input text-foreground placeholder:text-muted-foreground mb-1 max-h-32 w-full resize-none bg-transparent px-3.5 pt-1.5 pb-3 outline-none focus-visible:ring-0',
+              d('h-input'),
+              d('text-base')
+            )}
+            rows={1}
+            autoFocus={!isReplay}
+            disabled={isReplay}
+            aria-label="Message input"
+          />
+          <ComposerAction />
+        </ComposerPrimitive.Root>
+      )}
     </div>
   )
 }
@@ -508,13 +648,90 @@ const ComposerModelPicker: FC = () => {
   )
 }
 
+const CASSETTE_RECORDING_ENABLED =
+  import.meta.env.VITE_ELEMENTS_ENABLE_CASSETTE_RECORDING === 'true'
+
+const ComposerCassetteRecorder: FC = () => {
+  const [popoverOpen, setPopoverOpen] = useState(false)
+  const [tooltipOpen, setTooltipOpen] = useState(false)
+  const { isRecording, startRecording, stopRecording, download } =
+    useRecordCassette()
+
+  useEffect(() => {
+    if (popoverOpen) setTooltipOpen(false)
+  }, [popoverOpen])
+
+  return (
+    <TooltipProvider>
+      <Tooltip open={tooltipOpen && !popoverOpen} onOpenChange={setTooltipOpen}>
+        <Popover open={popoverOpen} onOpenChange={setPopoverOpen}>
+          <TooltipTrigger asChild>
+            <PopoverTrigger asChild>
+              <Button
+                variant="ghost"
+                size="icon"
+                data-state={popoverOpen ? 'open' : 'closed'}
+                className={cn(
+                  'aui-composer-cassette-recorder data-[state=open]:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30 flex w-fit items-center gap-2 rounded-full px-2.5 py-1 text-xs font-semibold',
+                  isRecording && 'text-red-500'
+                )}
+                aria-label="Cassette Recorder"
+              >
+                <CircleIcon
+                  className={cn(
+                    'size-5 stroke-[1.5px]',
+                    isRecording && 'animate-pulse fill-red-500 text-red-500'
+                  )}
+                />
+              </Button>
+            </PopoverTrigger>
+          </TooltipTrigger>
+          <PopoverContent side="top" align="start" className="w-64 p-3">
+            <div className="flex flex-col gap-3">
+              <div className="text-sm font-medium">Cassette Recorder</div>
+              {!isRecording ? (
+                <Button
+                  size="sm"
+                  variant="outline"
+                  className="w-full justify-start gap-2"
+                  onClick={startRecording}
+                >
+                  <CircleIcon className="size-3 fill-red-500 text-red-500" />
+                  Start Recording
+                </Button>
+              ) : (
+                <Button
+                  size="sm"
+                  variant="outline"
+                  className="w-full justify-start gap-2"
+                  onClick={() => {
+                    stopRecording()
+                    download()
+                    setPopoverOpen(false)
+                  }}
+                >
+                  <DownloadIcon className="size-3" />
+                  Stop &amp; Download
+                </Button>
+              )}
+            </div>
+          </PopoverContent>
+        </Popover>
+        <TooltipContent side="bottom" align="start">
+          {isRecording ? 'Recording…' : 'Cassette Recorder'}
+        </TooltipContent>
+      </Tooltip>
+    </TooltipProvider>
+  )
+}
+
 const ComposerAction: FC = () => {
   const { config } = useElements()
   const r = useRadius()
   const composerConfig = config.composer ?? { attachments: true }
   return (
     <div className="aui-composer-action-wrapper relative mx-1 mt-2 mb-2 flex items-center justify-between">
-      <div className="aui-composer-action-wrapper-inner flex items-center">
+      <div className="aui-composer-action-wrapper-inner text-muted-foreground flex items-center">
         {composerConfig.attachments ? (
           <ComposerAddAttachment />
         ) : (
@@ -524,6 +741,8 @@ const ComposerAction: FC = () => {
         {config.model?.showModelPicker && !config.languageModel && (
           <ComposerModelPicker />
         )}
+
+        {CASSETTE_RECORDING_ENABLED && <ComposerCassetteRecorder />}
       </div>
 
       <ThreadPrimitive.If running={false}>
@@ -572,10 +791,41 @@ const MessageError: FC = () => {
   )
 }
 
+/**
+ * Shows the pulsing dot indicator when the message is still running but the
+ * last rendered part is a tool call (not text). Without this, there's no
+ * visual feedback that the model is still working after a tool call.
+ */
+const ToolCallStreamingIndicator: FC = () => {
+  const show = useAssistantState(({ message }) => {
+    if (message.status?.type !== 'running') return false
+    const lastPart = message.parts[message.parts.length - 1]
+    return lastPart?.type === 'tool-call'
+  })
+  if (!show) return null
+  return <div className="aui-md mt-2" data-status="running" />
+}
+
 const AssistantMessage: FC = () => {
   const { config } = useElements()
   const toolsConfig = config.tools ?? {}
   const components = config.components ?? {}
+
+  const partsComponents = useMemo(
+    () => ({
+      Text: components.Text ?? MarkdownText,
+      Image: components.Image ?? Image,
+      tools: {
+        by_name: toolsConfig.components,
+        Fallback: components.ToolFallback ?? ToolFallback,
+      },
+      Reasoning: components.Reasoning ?? Reasoning,
+      ReasoningGroup: components.ReasoningGroup ?? ReasoningGroup,
+      ToolGroup: components.ToolGroup ?? ToolGroup,
+    }),
+    [components, toolsConfig.components]
+  )
+
   return (
     <MessagePrimitive.Root asChild>
       <div
@@ -583,23 +833,12 @@ const AssistantMessage: FC = () => {
         data-role="assistant"
       >
         <div className="aui-assistant-message-content text-foreground mx-2 leading-7 wrap-break-word">
-          <MessagePrimitive.Parts
-            components={{
-              Text: components.Text ?? MarkdownText,
-              Image: components.Image ?? Image,
-              tools: {
-                by_name: toolsConfig.components,
-                Fallback: components.ToolFallback ?? ToolFallback,
-              },
-              Reasoning: components.Reasoning ?? Reasoning,
-              ReasoningGroup: components.ReasoningGroup ?? ReasoningGroup,
-              ToolGroup: components.ToolGroup ?? ToolGroup,
-            }}
-          />
+          <MessagePrimitive.Parts components={partsComponents} />
+          <ToolCallStreamingIndicator />
           <MessageError />
         </div>
 
-        <div className="aui-assistant-message-footer mt-2 ml-2 flex">
+        <div className="aui-assistant-message-footer mt-2 ml-2 flex items-center gap-3">
           {/* <BranchPicker /> */}
           <AssistantActionBar />
         </div>

package/src/components/assistant-ui/tool-group.tsx

@@ -7,34 +7,45 @@ import { ToolUIGroup } from '@/components/ui/tool-ui'
 
 export const ToolGroup: FC<
   PropsWithChildren<{ startIndex: number; endIndex: number }>
-> = ({ children }) => {
-  const parts = useAssistantState(({ message }) => message).parts
-  const toolCallParts = parts.filter((part) => part.type === 'tool-call')
-  const anyMessagePartsAreRunning = toolCallParts.some(
-    (part) => part.status?.type === 'running'
-  )
+> = ({ children, startIndex, endIndex }) => {
+  // startIndex/endIndex are inclusive indices into message.parts.
+  // assistant-ui only groups consecutive tool-call parts, so every part
+  // in the range is a tool-call — the count is simply the range size.
+  const toolCount = endIndex - startIndex + 1
+
+  const firstToolName = useAssistantState(({ message }) => {
+    const part = message.parts[startIndex]
+    return part?.type === 'tool-call' ? part.toolName : undefined
+  })
+  const anyMessagePartsAreRunning = useAssistantState(({ message }) => {
+    for (let i = startIndex; i <= endIndex; i++) {
+      if (message.parts[i]?.status?.type === 'running') return true
+    }
+    return false
+  })
 
   const { config } = useElements()
   const defaultExpanded = config.tools?.expandToolGroupsByDefault ?? false
 
   const groupTitle = useMemo(() => {
-    const toolParts = parts.filter((part) => part.type === 'tool-call')
-
-    if (toolParts.length === 0) return 'No tools called'
-    if (toolParts.length === 1)
-      return `Calling ${humanizeToolName(toolParts[0].toolName)}...`
+    if (toolCount === 0) return 'No tools called'
+    if (toolCount === 1) {
+      return firstToolName
+        ? `Calling ${humanizeToolName(firstToolName)}...`
+        : 'Calling tool...'
+    }
     return anyMessagePartsAreRunning
-      ? `Calling ${toolParts.length} tools...`
-      : `Executed ${toolParts.length} tools`
-  }, [parts, anyMessagePartsAreRunning])
+      ? `Calling ${toolCount} tools...`
+      : `Executed ${toolCount} tools`
+  }, [toolCount, firstToolName, anyMessagePartsAreRunning])
 
   // If there's a custom component for the single tool, render children directly
-  if (config.tools?.components?.[toolCallParts[0]?.toolName]) {
+  if (firstToolName && config.tools?.components?.[firstToolName]) {
     return children
   }
 
   // For single tool calls, render without the group wrapper
-  if (toolCallParts.length === 1) {
+  if (toolCount === 1) {
     return (
       <div className={cn('my-4 w-full max-w-xl')}>
         <div className="border-border bg-card overflow-hidden rounded-lg border">

package/src/components/ui/tool-ui.tsx

@@ -181,6 +181,22 @@ function CopyButton({ content }: { content: string }) {
  * SyntaxHighlightedCode - Code block with shiki syntax highlighting
  * -------------------------------------------------------------------------- */
 
+/** Max characters to send through shiki — above this we skip highlighting. */
+const SHIKI_CHAR_LIMIT = 8_000
+/** Max lines shown in the collapsed preview. */
+const PREVIEW_LINE_LIMIT = 50
+
+function truncateToLines(text: string, maxLines: number) {
+  let pos = 0
+  for (let i = 0; i < maxLines; i++) {
+    const next = text.indexOf('\n', pos)
+    if (next === -1) return { text, truncated: false, totalLines: i + 1 }
+    pos = next + 1
+  }
+  const totalLines = text.split('\n').length
+  return { text: text.slice(0, pos), truncated: true, totalLines }
+}
+
 function SyntaxHighlightedCode({
   text,
   language,
@@ -191,10 +207,20 @@ function SyntaxHighlightedCode({
   className?: string
 }) {
   const [highlightedCode, setHighlightedCode] = useState<string | null>(null)
+  const [expanded, setExpanded] = useState(false)
+
+  const preview = React.useMemo(
+    () => truncateToLines(text, PREVIEW_LINE_LIMIT),
+    [text]
+  )
+  const displayText = expanded ? text : preview.text
+  const canHighlight = displayText.length <= SHIKI_CHAR_LIMIT
 
   useEffect(() => {
-    if (!language) return
-    codeToHtml(text, {
+    setHighlightedCode(null)
+    if (!language || !canHighlight) return
+    let cancelled = false
+    codeToHtml(displayText, {
       lang: language,
       theme: 'github-dark-default',
       rootStyle: 'background-color: transparent;',
@@ -206,27 +232,43 @@ function SyntaxHighlightedCode({
           },
         },
       ],
-    }).then(setHighlightedCode)
-  }, [text, language])
+    }).then((html) => {
+      if (!cancelled) setHighlightedCode(html)
+    })
+    return () => {
+      cancelled = true
+    }
+  }, [displayText, language, canHighlight])
+
+  const showMoreButton = preview.truncated && !expanded && (
+    <button
+      type="button"
+      onClick={() => setExpanded(true)}
+      className="w-full bg-slate-800/90 px-4 py-2 text-left text-xs text-slate-400 transition-colors hover:text-slate-200"
+    >
+      Show all {preview.totalLines} lines…
+    </button>
+  )
 
-  if (!highlightedCode) {
+  if (!canHighlight || !highlightedCode) {
     return (
-      <pre
-        className={cn(
-          'w-full bg-slate-800/90 px-4 py-3 text-sm whitespace-pre-wrap text-slate-100',
-          className
-        )}
-      >
-        {text}
-      </pre>
+      <div className={cn('w-full', className)}>
+        <pre className="max-h-[300px] w-full overflow-y-auto bg-slate-800/90 px-4 py-3 text-sm whitespace-pre-wrap text-slate-100">
+          {displayText}
+        </pre>
+        {showMoreButton}
+      </div>
     )
   }
 
   return (
-    <div
-      className={cn('w-full bg-slate-800/90', className)}
-      dangerouslySetInnerHTML={{ __html: highlightedCode }}
-    />
+    <div className={cn('w-full', className)}>
+      <div
+        className="w-full bg-slate-800/90"
+        dangerouslySetInnerHTML={{ __html: highlightedCode }}
+      />
+      {showMoreButton}
+    </div>
   )
 }