diff --git a/src/message.tsx b/src/message.tsx
index 916ac66..7b62d81 100644
--- a/src/message.tsx
+++ b/src/message.tsx
@@ -124,7 +124,6 @@ export default function Message(props: Props) {
         variant={chat.role === "assistant" ? "received" : "sent"}
         className={chat.role !== "assistant" ? "flex-row-reverse" : ""}
       >
-        <ChatBubbleAvatar fallback={chat.role === "assistant" ? "AI" : "US"} />
         <ChatBubbleMessage isLoading={false}>
           {chat.hide ? (
             <MessageHide chat={chat} />
diff --git a/src/messageHide.tsx b/src/messageHide.tsx
index f8bfd69..3f11728 100644
--- a/src/messageHide.tsx
+++ b/src/messageHide.tsx
@@ -1,5 +1,6 @@
 import { ChatStoreMessage } from "@/types/chatstore";
 import { getMessageText } from "@/chatgpt";
+import { Badge } from "./components/ui/badge";
 
 interface Props {
   chat: ChatStoreMessage;
@@ -9,9 +10,7 @@ export function MessageHide({ chat }: Props) {
   return (
     <div className="flex items-center gap-2 text-sm text-muted-foreground">
       <span>{getMessageText(chat).split("\n")[0].slice(0, 18)} ...</span>
-      <span className="rounded-md bg-secondary px-2 py-1 text-xs">
-        Removed from context
-      </span>
+      <Badge variant="secondary">Removed from context</Badge>
     </div>
   );
 }