feat: re-implement message group

frontend/src/components/ai-elements/chain-of-thought.tsx

@@ -15,7 +15,13 @@ import {
   type LucideIcon,
 } from "lucide-react";
 import type { ComponentProps, ReactNode } from "react";
-import { createContext, memo, useContext, useMemo } from "react";
+import {
+  createContext,
+  isValidElement,
+  memo,
+  useContext,
+  useMemo,
+} from "react";
 
 type ChainOfThoughtContextValue = {
   isOpen: boolean;
@@ -108,7 +114,7 @@ export const ChainOfThoughtHeader = memo(
 );
 
 export type ChainOfThoughtStepProps = ComponentProps<"div"> & {
-  icon?: LucideIcon;
+  icon?: LucideIcon | React.ReactElement;
   label: ReactNode;
   description?: ReactNode;
   status?: "complete" | "active" | "pending";
@@ -141,7 +147,7 @@ export const ChainOfThoughtStep = memo(
         {...props}
       >
         <div className="relative mt-0.5">
-          <Icon className="size-4" />
+          {isValidElement(Icon) ? Icon : <Icon className="size-4" />}
           <div className="bg-border absolute top-7 bottom-0 left-1/2 -mx-px w-px" />
         </div>
         <div className="flex-1 space-y-2 overflow-hidden">

frontend/src/components/ai-elements/code-block.tsx

@@ -114,7 +114,7 @@ export const CodeBlock = ({
             dangerouslySetInnerHTML={{ __html: html }}
           />
           <div
-            className="[&>pre]:bg-background! [&>pre]:text-foreground! hidden overflow-auto dark:block [&_code]:font-mono [&_code]:text-sm [&>pre]:m-0 [&>pre]:p-4 [&>pre]:text-sm"
+            className="[&>pre]:bg-background! [&>pre]:text-foreground! hidden size-full overflow-auto dark:block [&_code]:font-mono [&_code]:text-sm [&>pre]:m-0 [&>pre]:p-4 [&>pre]:text-sm [&>pre]:whitespace-pre-wrap"
             // biome-ignore lint/security/noDangerouslySetInnerHtml: "this is needed."
             dangerouslySetInnerHTML={{ __html: darkHtml }}
           />