From 4518a6c857b764ba7e057ddf0a6e3ca519980867 Mon Sep 17 00:00:00 2001
From: Andrey Goder <andy.goder@gmail.com>
Date: Mon, 24 Aug 2026 18:37:36 -0700
Subject: [PATCH] Add configurable chat bubble width

---
 src/components/Chat/ChatMessage.stories.tsx | 38 ++++++++++++++
 src/components/Chat/ChatMessage.test.tsx    | 55 +++++++++++++++++++++
 src/components/Chat/ChatMessageBubble.tsx   | 14 +++++-
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/src/components/Chat/ChatMessage.stories.tsx b/src/components/Chat/ChatMessage.stories.tsx
index 17d55b3d..828cbb56 100644
--- a/src/components/Chat/ChatMessage.stories.tsx
+++ b/src/components/Chat/ChatMessage.stories.tsx
@@ -6,6 +6,8 @@ import {ChatMessageBubble} from 'components/Chat/ChatMessageBubble';
 import {ChatMessageMetadata} from 'components/Chat/ChatMessageMetadata';
 import {ChatSystemMessage} from 'components/Chat/ChatSystemMessage';
 import {Icon} from 'components/Icon';
+import {Table} from 'components/Table/Table';
+import type {TableColumn} from 'components/Table/types';
 import {Timestamp} from 'components/Timestamp';
 import {css} from 'styled-system/css';
 
@@ -37,6 +39,24 @@ const columnStyle = css({
   maxW: '640px',
 });
 
+interface UsageRow extends Record<string, unknown> {
+  model: string;
+  requests: string;
+  tokens: string;
+}
+
+const usageColumns: TableColumn<UsageRow>[] = [
+  {header: 'Model', key: 'model'},
+  {align: 'end', header: 'Requests', key: 'requests'},
+  {align: 'end', header: 'Tokens', key: 'tokens'},
+];
+
+const usageData: UsageRow[] = [
+  {model: 'Atlas', requests: '1,842', tokens: '3.2M'},
+  {model: 'Nova', requests: '967', tokens: '1.1M'},
+  {model: 'Ember', requests: '413', tokens: '620K'},
+];
+
 export const Basic: Story = {
   render: args => (
     <ChatMessage {...args}>
@@ -76,6 +96,24 @@ export const BubbleVariants: Story = {
   ),
 };
 
+export const FullWidthContent: Story = {
+  render: () => (
+    <div className={columnStyle}>
+      <ChatMessage sender="assistant">
+        <ChatMessageBubble width="full">
+          <Table
+            columns={usageColumns}
+            data={usageData}
+            density="compact"
+            idKey="model"
+            label="Model usage"
+          />
+        </ChatMessageBubble>
+      </ChatMessage>
+    </div>
+  ),
+};
+
 export const GroupedBubbles: Story = {
   render: () => (
     <div className={columnStyle}>
diff --git a/src/components/Chat/ChatMessage.test.tsx b/src/components/Chat/ChatMessage.test.tsx
index ae3ee7f9..61fb4301 100644
--- a/src/components/Chat/ChatMessage.test.tsx
+++ b/src/components/Chat/ChatMessage.test.tsx
@@ -277,6 +277,61 @@ describe('ChatMessageBubble', () => {
     expect(ref).toHaveBeenCalledWith(expect.any(HTMLDivElement));
   });
 
+  it('resolves a numeric width to pixels and replaces the default cap', () => {
+    render(
+      <ChatMessageBubble data-testid="bubble" width={480}>
+        Hi
+      </ChatMessageBubble>,
+    );
+
+    expect(screen.getByTestId('bubble')).toHaveStyle({
+      width: '480px',
+      maxWidth: 'none',
+    });
+  });
+
+  it('uses string widths as-is', () => {
+    render(
+      <ChatMessageBubble data-testid="bubble" width="75%">
+        Hi
+      </ChatMessageBubble>,
+    );
+
+    expect(screen.getByTestId('bubble')).toHaveStyle({
+      width: '75%',
+      maxWidth: 'none',
+    });
+  });
+
+  it("fills the message column for width='full'", () => {
+    render(
+      <ChatMessageBubble data-testid="bubble" width="full">
+        Hi
+      </ChatMessageBubble>,
+    );
+
+    expect(screen.getByTestId('bubble')).toHaveStyle({
+      width: '100%',
+      maxWidth: 'none',
+    });
+  });
+
+  it('lets consumer style override width and max-width', () => {
+    render(
+      <ChatMessageBubble
+        data-testid="bubble"
+        style={{width: '12rem', maxWidth: '10rem'}}
+        width="full">
+        Hi
+      </ChatMessageBubble>,
+    );
+
+    expect(screen.getByTestId('bubble')).toHaveStyle({
+      width: '12rem',
+      maxWidth: '10rem',
+    });
+  });
+
   it('forwards the curated passthrough props', () => {
     render(
       <ChatMessage sender="user">
diff --git a/src/components/Chat/ChatMessageBubble.tsx b/src/components/Chat/ChatMessageBubble.tsx
index 3f171055..d6b75a56 100644
--- a/src/components/Chat/ChatMessageBubble.tsx
+++ b/src/components/Chat/ChatMessageBubble.tsx
@@ -5,6 +5,7 @@ import {useChatMessageContext} from 'components/Chat/ChatContext';
 import {chatMessageBubbleRecipe} from 'components/Chat/ChatMessageBubble.recipe';
 import type {ChatPassthroughProps} from 'components/Chat/ChatPassthroughProps';
 import isNonEmptyReactNode from 'internal/isNonEmptyReactNode';
+import {toWidthSize, type WidthValue} from 'internal/toPixelSize';
 import {cx} from 'utils/cx';
 
 export type ChatMessageBubbleVariant = 'filled' | 'ghost';
@@ -54,6 +55,11 @@ export interface ChatMessageBubbleProps extends ChatPassthroughProps {
    * @default 'filled'
    */
   variant?: ChatMessageBubbleVariant;
+  /**
+   * Bubble width. Numbers are pixels, strings are used as-is, `'full'` fills
+   * the message column. When set, replaces the default maximum-width cap.
+   */
+  width?: WidthValue;
 }
 
 /**
@@ -70,12 +76,18 @@ export function ChatMessageBubble({
   ref,
   style,
   variant = 'filled',
+  width,
   ...passthrough
 }: ChatMessageBubbleProps): React.JSX.Element {
   const messageContext = useChatMessageContext();
   const sender = messageContext?.sender ?? 'assistant';
   const density = messageContext?.density ?? 'balanced';
   const classes = chatMessageBubbleRecipe({density, group, sender, variant});
+  // Consumer `style` still wins, matching the layout primitives.
+  const bubbleStyle: CSSProperties | undefined =
+    width == null
+      ? style
+      : {width: toWidthSize(width), maxWidth: 'none', ...style};
 
   return (
     <>
@@ -89,7 +101,7 @@ export function ChatMessageBubble({
         className={cx(classes.bubble, className)}
         data-testid={dataTestId}
         ref={ref}
-        style={style}>
+        style={bubbleStyle}>
         {children}
       </div>
       {isNonEmptyReactNode(metadata) ? (