From 21aa62a69a92d1f3cf9b4aee72b4d506b28a4333 Mon Sep 17 00:00:00 2001
From: ephraimduncan <ephraimduncan68@gmail.com>
Date: Fri, 31 Jul 2026 15:14:19 +0000
Subject: [PATCH] fix(elements): prevent consumer className from disabling
 branch visibility

---
 packages/elements/__tests__/message.test.tsx | 29 ++++++++++++++++++++
 packages/elements/src/message.tsx            |  6 ++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/packages/elements/__tests__/message.test.tsx b/packages/elements/__tests__/message.test.tsx
index a6d8a89a..91457c8f 100644
--- a/packages/elements/__tests__/message.test.tsx
+++ b/packages/elements/__tests__/message.test.tsx
@@ -201,6 +201,35 @@ describe("messageBranchContent", () => {
 
     expect(screen.getByText("Branch 1")).toBeInTheDocument();
   });
+
+  it("keeps exactly one branch visible when a consumer className is passed", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <MessageBranch>
+        <MessageBranchContent className="rounded-md">
+          <div key="1">Branch 1</div>
+          <div key="2">Branch 2</div>
+        </MessageBranchContent>
+        <MessageBranchNext />
+      </MessageBranch>
+    );
+
+    const firstWrapper = screen.getByText("Branch 1").parentElement;
+    const secondWrapper = screen.getByText("Branch 2").parentElement;
+
+    expect(firstWrapper).toHaveClass("rounded-md", "grid");
+    expect(firstWrapper).not.toHaveClass("hidden");
+    expect(secondWrapper).toHaveClass("rounded-md", "hidden");
+    expect(secondWrapper).not.toHaveClass("grid");
+
+    await user.click(screen.getByRole("button", { name: NEXT_REGEX }));
+
+    expect(firstWrapper).toHaveClass("hidden");
+    expect(firstWrapper).not.toHaveClass("grid");
+    expect(secondWrapper).toHaveClass("rounded-md", "grid");
+    expect(secondWrapper).not.toHaveClass("hidden");
+  });
 });
 
 describe("messageBranchSelector", () => {
diff --git a/packages/elements/src/message.tsx b/packages/elements/src/message.tsx
index c04a3993..befbf8f8 100644
--- a/packages/elements/src/message.tsx
+++ b/packages/elements/src/message.tsx
@@ -198,6 +198,7 @@ export type MessageBranchContentProps = HTMLAttributes<HTMLDivElement>;
 
 export const MessageBranchContent = ({
   children,
+  className,
   ...props
 }: MessageBranchContentProps) => {
   const { currentBranch, setBranches, branches } = useMessageBranch();
@@ -216,8 +217,9 @@ export const MessageBranchContent = ({
   return childrenArray.map((branch, index) => (
     <div
       className={cn(
-        "grid gap-2 overflow-hidden [&>div]:pb-0",
-        index === currentBranch ? "block" : "hidden"
+        "gap-2 overflow-hidden [&>div]:pb-0",
+        className,
+        index === currentBranch ? "grid" : "hidden"
       )}
       key={branch.key}
       {...props}