@vertesia/ui 1.0.0-dev.20260203.130115Z → 1.0.0-dev.20260227.112605Z

package/src/features/agent/chat/ModernAgentOutput/MessageItem.tsx

@@ -1,10 +1,10 @@
-import { AgentMessage, AgentMessageType, AskUserMessageDetails } from "@vertesia/common";
-import { Badge, Button, useToast } from "@vertesia/ui/core";
+import { AgentMessage, AgentMessageType, AskUserMessageDetails, MarkdownRenditionFormat } from "@vertesia/common";
+import { Badge, Button, cn, Dropdown, MenuItem, useToast } from "@vertesia/ui/core";
 import { NavLink } from "@vertesia/ui/router";
 import { useUserSession } from "@vertesia/ui/session";
 import { MarkdownRenderer } from "@vertesia/ui/widgets";
 import dayjs from "dayjs";
-import { AlertCircle, Bot, CheckCircle, Clock, CopyIcon, Info, Layers, MessageSquare, User } from "lucide-react";
+import { AlertCircle, Bot, CheckCircle, Clock, CopyIcon, Download, Info, Layers, type LucideIcon, MessageSquare, User } from "lucide-react";
 import React, { useEffect, useState, useMemo, memo, useRef } from "react";
 import { PulsatingCircle } from "../AnimatedThinkingDots";
 import { AskUserWidget } from "../AskUserWidget";
@@ -12,6 +12,7 @@ import { useImageLightbox } from "../ImageLightbox";
 import { ThinkingMessages } from "../WaitingMessages";
 import { getWorkstreamId } from "./utils";
 import { useArtifactUrlCache, getArtifactCacheKey } from "../useArtifactUrlCache.js";
+import { useDownloadFile } from "../../../store/index.js";
 
 // PERFORMANCE: Move pure function outside component to avoid recreation on every render
 // Process content to enhance markdown detection for lists and thinking messages
@@ -65,11 +66,8 @@ function processContentForMarkdown(content: string | object, messageType: AgentM
     return content;
 }
 
-export interface MessageItemProps {
-    message: AgentMessage;
-    showPulsatingCircle?: boolean;
-    /** Callback when user sends a message (e.g., from proposal selection) */
-    onSendMessage?: (message: string) => void;
+/** className overrides for MessageItem — single source of truth for all className overrides. */
+export interface MessageItemClassNames {
     /** Additional className for the outer container */
     className?: string;
     /** Additional className for the card wrapper */
@@ -88,31 +86,68 @@ export interface MessageItemProps {
     detailsClassName?: string;
     /** Additional className for the artifacts section */
     artifactsClassName?: string;
+    /** Additional className for the prose/markdown wrapper */
+    proseClassName?: string;
+}
+
+/** Keys of {@link MessageItemClassNames} — drives className merging and memo comparison. */
+const MESSAGE_ITEM_CLASS_NAME_KEYS: (keyof MessageItemClassNames)[] = [
+    'className', 'cardClassName', 'headerClassName', 'contentClassName',
+    'timestampClassName', 'senderClassName', 'iconClassName',
+    'detailsClassName', 'artifactsClassName', 'proseClassName',
+];
+
+/** Merge className slots across base, prop, and override layers with consistent priority. */
+function mergeClassNames(
+    base: MessageItemClassNames,
+    props: MessageItemClassNames,
+    ...overrides: (Partial<MessageItemClassNames> | undefined)[]
+): MessageItemClassNames {
+    const result: Record<string, string | undefined> = {};
+    for (const key of MESSAGE_ITEM_CLASS_NAME_KEYS) {
+        result[key] = cn(base[key], props[key], ...overrides.map(o => o?.[key]));
+    }
+    return result as MessageItemClassNames;
 }
 
-// Consolidated message styling - single source of truth
-const MESSAGE_STYLES: Record<AgentMessageType | 'default', {
+/** Per-message-type visual config (border, bg, icon color, sender label, icon component, optional className overrides). */
+export interface MessageStyleConfig extends MessageItemClassNames {
     borderColor: string;
-    bgColor: string;
     iconColor: string;
     sender: string;
-    Icon: typeof Bot;
-}> = {
-    [AgentMessageType.ANSWER]: { borderColor: 'border-l-info', bgColor: 'bg-info', iconColor: 'text-info', sender: 'Agent', Icon: Bot },
-    [AgentMessageType.COMPLETE]: { borderColor: 'border-l-success', bgColor: 'bg-success', iconColor: 'text-success', sender: 'Completed', Icon: CheckCircle },
-    [AgentMessageType.IDLE]: { borderColor: 'border-l-info', bgColor: 'bg-info/50', iconColor: 'text-info', sender: 'Ready', Icon: Clock },
-    [AgentMessageType.REQUEST_INPUT]: { borderColor: 'border-l-attention', bgColor: 'bg-attention', iconColor: 'text-attention', sender: 'Input', Icon: User },
-    [AgentMessageType.QUESTION]: { borderColor: 'border-l-muted', bgColor: 'bg-muted', iconColor: 'text-muted', sender: 'User', Icon: User },
-    [AgentMessageType.THOUGHT]: { borderColor: 'border-l-purple-500', bgColor: 'bg-purple-50/50 dark:bg-purple-900/10', iconColor: 'text-purple-600 dark:text-purple-400', sender: 'Agent', Icon: Bot },
-    [AgentMessageType.ERROR]: { borderColor: 'border-l-destructive', bgColor: 'bg-destructive', iconColor: 'text-destructive', sender: 'Error', Icon: AlertCircle },
-    [AgentMessageType.UPDATE]: { borderColor: 'border-l-success', bgColor: 'bg-success', iconColor: 'text-success', sender: 'Update', Icon: Info },
-    [AgentMessageType.PLAN]: { borderColor: 'border-l-attention', bgColor: 'bg-attention', iconColor: 'text-attention', sender: 'Plan', Icon: MessageSquare },
-    [AgentMessageType.TERMINATED]: { borderColor: 'border-l-muted', bgColor: 'bg-muted', iconColor: 'text-muted', sender: 'Terminated', Icon: CheckCircle },
-    [AgentMessageType.WARNING]: { borderColor: 'border-l-attention', bgColor: 'bg-attention', iconColor: 'text-attention', sender: 'Warning', Icon: AlertCircle },
-    [AgentMessageType.SYSTEM]: { borderColor: 'border-l-muted', bgColor: 'bg-muted', iconColor: 'text-muted', sender: 'System', Icon: Info },
-    [AgentMessageType.STREAMING_CHUNK]: { borderColor: 'border-l-info', bgColor: 'bg-info', iconColor: 'text-info', sender: 'Agent', Icon: Bot },
-    [AgentMessageType.BATCH_PROGRESS]: { borderColor: 'border-l-blue-500', bgColor: 'bg-blue-50/50 dark:bg-blue-900/10', iconColor: 'text-blue-600 dark:text-blue-400', sender: 'Batch', Icon: Layers },
-    default: { borderColor: 'border-l-muted', bgColor: 'bg-muted', iconColor: 'text-muted', sender: 'Agent', Icon: Bot },
+    Icon: LucideIcon;
+}
+
+export interface MessageItemProps extends MessageItemClassNames {
+    message: AgentMessage;
+    showPulsatingCircle?: boolean;
+    /** Callback when user sends a message (e.g., from proposal selection) */
+    onSendMessage?: (message: string) => void;
+    /** Sparse per-type overrides for MESSAGE_STYLES (deep-merged with defaults) */
+    messageStyleOverrides?: Partial<Record<AgentMessageType | 'default', Partial<MessageStyleConfig>>>;
+    /** Custom component to render store/document links instead of default NavLink navigation */
+    StoreLinkComponent?: React.ComponentType<{ href: string; documentId: string; children: React.ReactNode }>;
+    /** Custom component to render store/collection links instead of default NavLink navigation */
+    CollectionLinkComponent?: React.ComponentType<{ href: string; collectionId: string; children: React.ReactNode }>;
+}
+
+// Consolidated Studio/default message styling - single source of truth
+export const MESSAGE_STYLES: Record<AgentMessageType | 'default', MessageStyleConfig> = {
+    [AgentMessageType.ANSWER]: { borderColor: 'border-l-info', iconColor: 'text-info', sender: 'Agent', Icon: Bot },
+    [AgentMessageType.COMPLETE]: { borderColor: 'border-l-success', iconColor: 'text-success', sender: 'Completed', Icon: CheckCircle },
+    [AgentMessageType.IDLE]: { borderColor: 'border-l-info', iconColor: 'text-info', sender: 'Ready', Icon: Clock },
+    [AgentMessageType.REQUEST_INPUT]: { borderColor: 'border-l-attention', iconColor: 'text-attention', sender: 'Input', Icon: User },
+    [AgentMessageType.QUESTION]: { borderColor: 'border-l-muted', iconColor: 'text-muted', sender: 'User', Icon: User },
+    [AgentMessageType.THOUGHT]: { borderColor: 'border-l-purple-500', iconColor: 'text-purple-600 dark:text-purple-400', sender: 'Agent', Icon: Bot },
+    [AgentMessageType.ERROR]: { borderColor: 'border-l-destructive', iconColor: 'text-destructive', sender: 'Error', Icon: AlertCircle },
+    [AgentMessageType.UPDATE]: { borderColor: 'border-l-success', iconColor: 'text-success', sender: 'Update', Icon: Info },
+    [AgentMessageType.PLAN]: { borderColor: 'border-l-attention', iconColor: 'text-attention', sender: 'Plan', Icon: MessageSquare },
+    [AgentMessageType.TERMINATED]: { borderColor: 'border-l-muted', iconColor: 'text-muted', sender: 'Terminated', Icon: CheckCircle },
+    [AgentMessageType.WARNING]: { borderColor: 'border-l-attention', iconColor: 'text-attention', sender: 'Warning', Icon: AlertCircle },
+    [AgentMessageType.SYSTEM]: { borderColor: 'border-l-muted', iconColor: 'text-muted', sender: 'System', Icon: Info },
+    [AgentMessageType.STREAMING_CHUNK]: { borderColor: 'border-l-info', iconColor: 'text-info', sender: 'Agent', Icon: Bot },
+    [AgentMessageType.BATCH_PROGRESS]: { borderColor: 'border-l-blue-500', iconColor: 'text-blue-600 dark:text-blue-400', sender: 'Batch', Icon: Layers },
+    default: { borderColor: 'border-l-muted', iconColor: 'text-muted', sender: 'Agent', Icon: Bot },
 };
 
 function MessageItemComponent({
@@ -128,6 +163,10 @@ function MessageItemComponent({
     iconClassName,
     detailsClassName,
     artifactsClassName,
+    proseClassName,
+    messageStyleOverrides,
+    StoreLinkComponent,
+    CollectionLinkComponent,
 }: MessageItemProps) {
     const [showDetails, setShowDetails] = useState(false);
     const { client } = useUserSession();
@@ -139,9 +178,27 @@ function MessageItemComponent({
     clientRef.current = client;
     const urlCacheRef = useRef(urlCache);
     urlCacheRef.current = urlCache;
-
-    // Get styles from consolidated config
-    const styles = MESSAGE_STYLES[message.type] || MESSAGE_STYLES.default;
+    const { renderContent: exportContent, isDownloading: isExportingFile } = useDownloadFile({ client, toast });
+
+    // Unified style resolution: merge MESSAGE_STYLES base, flat className props, and per-type overrides.
+    // Priority (lowest → highest): base MESSAGE_STYLES → flat props → overrides.default → overrides[type]
+    const resolvedStyle = useMemo(() => {
+        const base = MESSAGE_STYLES[message.type] || MESSAGE_STYLES.default;
+        const defaultOverrides = messageStyleOverrides?.default;
+        const typeOverrides = messageStyleOverrides?.[message.type];
+
+        return {
+            ...base, ...defaultOverrides, ...typeOverrides,
+            ...mergeClassNames(base, {
+                className, cardClassName, headerClassName, contentClassName,
+                timestampClassName, senderClassName, iconClassName,
+                detailsClassName, artifactsClassName, proseClassName,
+            }, defaultOverrides, typeOverrides),
+        };
+    }, [message.type, messageStyleOverrides,
+        className, cardClassName, headerClassName, contentClassName,
+        timestampClassName, senderClassName, iconClassName,
+        detailsClassName, artifactsClassName, proseClassName]);
 
     // PERFORMANCE: Memoize message content extraction - only recalculates when message changes
     const messageContent = useMemo(() => {
@@ -201,12 +258,88 @@ function MessageItemComponent({
         });
     };
 
+    // Export message content to PDF or DOCX
+    const exportToFormat = async (format: MarkdownRenditionFormat) => {
+        const content = typeof messageContent === 'string' ? messageContent : '';
+
+        if (!content.trim()) {
+            toast({
+                status: "error",
+                title: "No content to export",
+                duration: 2000,
+            });
+            return;
+        }
+
+        const title = `Message ${dayjs(message.timestamp).format("YYYY-MM-DD HH-mm-ss")}`;
+        await exportContent(content, {
+            format,
+            title,
+            artifactRunId: runId,
+        });
+    };
+
+    // Check if message has exportable content (markdown text)
+    const hasExportableContent = typeof messageContent === 'string' && messageContent.trim().length > 0;
+
+    // PERFORMANCE: Memoize markdown components to prevent MarkdownRenderer remounts
+    const markdownComponents = useMemo(() => ({
+        a: ({ node, ref, ...props }: { node?: any; ref?: any; href?: string; children?: React.ReactNode }) => {
+            const href = props.href || "";
+            if (href.includes("/store/objects")) {
+                if (StoreLinkComponent) {
+                    const documentId = href.split("/store/objects/")[1] || "";
+                    return <StoreLinkComponent href={href} documentId={documentId}>{props.children}</StoreLinkComponent>;
+                }
+                return (
+                    <NavLink
+                        href={href}
+                        topLevelNav
+                    >
+                        {props.children}
+                    </NavLink>
+                );
+            }
+            if (href.includes("/store/collections")) {
+                if (CollectionLinkComponent) {
+                    const collectionId = href.split("/store/collections/")[1] || "";
+                    return <CollectionLinkComponent href={href} collectionId={collectionId}>{props.children}</CollectionLinkComponent>;
+                }
+                return (
+                    <NavLink
+                        href={href}
+                        topLevelNav
+                    >
+                        {props.children}
+                    </NavLink>
+                );
+            }
+            return (
+                <a
+                    {...props}
+                    target="_blank"
+                    rel="noopener noreferrer"
+                />
+            );
+        },
+        img: ({ node, ref, ...props }: { node?: any; ref?: any; src?: string; alt?: string }) => {
+            return (
+                <img
+                    {...props}
+                    className="max-w-full h-auto rounded-lg shadow-md my-3 cursor-pointer hover:shadow-lg transition-shadow"
+                    loading="lazy"
+                    onClick={() => props.src && openImage(props.src, props.alt)}
+                />
+            );
+        },
+    }), [openImage, StoreLinkComponent, CollectionLinkComponent]);
+
     // Render content with markdown support - all messages now rendered as markdown
     const renderContent = (content: string | object) => {
         // Handle object content (JSON)
         if (typeof content === "object") {
             return (
-                <pre className="text-xs font-mono whitespace-pre-wrap overflow-x-auto bg-gray-100 dark:bg-gray-800 p-2 rounded text-gray-700 ">
+                <pre className="text-xs font-mono whitespace-pre-wrap overflow-x-auto bg-gray-100 dark:bg-gray-800 p-2 rounded text-gray-700">
                     {JSON.stringify(content, null, 2)}
                 </pre>
             );
@@ -216,43 +349,12 @@ function MessageItemComponent({
         const runId = (message as any).workflow_run_id as string | undefined;
 
         return (
-            <div className="vprose prose prose-slate dark:prose-invert prose-p:leading-relaxed prose-p:my-3 prose-headings:font-semibold prose-headings:tracking-normal prose-headings:mt-6 prose-headings:mb-3 prose-h1:text-xl prose-h2:text-lg prose-h3:text-base prose-li:my-1 prose-ul:my-3 prose-ol:my-3 prose-table:my-5 prose-pre:my-4 prose-hr:my-6 max-w-none text-[15px] break-words" style={{ overflowWrap: 'anywhere' }}>
+            <div className={cn("vprose prose prose-slate dark:prose-invert prose-p:leading-relaxed prose-p:my-3 prose-headings:font-semibold prose-headings:tracking-normal prose-headings:mt-6 prose-headings:mb-3 prose-h1:text-xl prose-h2:text-lg prose-h3:text-base prose-li:my-1 prose-ul:my-3 prose-ol:my-3 prose-table:my-5 prose-pre:my-4 prose-hr:my-6 max-w-none text-[15px] break-words", resolvedStyle.proseClassName)} style={{ overflowWrap: 'anywhere' }}>
                 <MarkdownRenderer
                     artifactRunId={runId}
                     onProposalSelect={(optionId) => onSendMessage?.(optionId)}
                     onProposalSubmit={(text) => onSendMessage?.(text)}
-                    components={{
-                        a: ({ node, ref, ...props }: { node?: any; ref?: any; href?: string; children?: React.ReactNode }) => {
-                            const href = props.href || "";
-                            if (href.includes("/store/objects")) {
-                                return (
-                                    <NavLink
-                                        href={href}
-                                        topLevelNav
-                                    >
-                                        {props.children}
-                                    </NavLink>
-                                );
-                            }
-                            return (
-                                <a
-                                    {...props}
-                                    target="_blank"
-                                    rel="noopener noreferrer"
-                                />
-                            );
-                        },
-                        img: ({ node, ref, ...props }: { node?: any; ref?: any; src?: string; alt?: string }) => {
-                            return (
-                                <img
-                                    {...props}
-                                    className="max-w-full h-auto rounded-lg shadow-md my-3 cursor-pointer hover:shadow-lg transition-shadow"
-                                    loading="lazy"
-                                    onClick={() => props.src && openImage(props.src, props.alt)}
-                                />
-                            );
-                        },
-                    }}
+                    components={markdownComponents}
                 >
                     {content as string}
                 </MarkdownRenderer>
@@ -287,12 +389,8 @@ function MessageItemComponent({
                     outputFiles.map(async (name: unknown) => {
                         if (typeof name !== "string" || !name.trim()) return null;
                         const trimmed = name.trim();
-                        // execute_shell returns relative paths like "result.csv" or "plots/chart.png"
-                        // canonical artifact name is under out/
-                        const artifactPath =
-                            trimmed.startsWith("out/") || trimmed.startsWith("files/") || trimmed.startsWith("scripts/")
-                                ? trimmed
-                                : `out/${trimmed}`;
+                        // Strip artifact: protocol prefix to get the artifact-relative path
+                        const artifactPath = trimmed.startsWith("artifact:") ? trimmed.slice(9) : trimmed;
 
                         const ext = artifactPath.split(".").pop()?.toLowerCase() || "";
                         const imageExtensions = new Set(["png", "jpg", "jpeg", "gif", "webp", "bmp", "svg"]);
@@ -342,29 +440,29 @@ function MessageItemComponent({
     }, [runId, outputFilesKey]);
 
     const workstreamId = getWorkstreamId(message);
-    const { Icon } = styles;
+    const { Icon } = resolvedStyle;
 
     // Render icon - show pulsating animation for active messages
     const renderIcon = () => {
         if (showPulsatingCircle) {
             return <PulsatingCircle size="sm" color="blue" />;
         }
-        return <Icon className={`size-4 ${styles.iconColor}`} />;
+        return <Icon className={`size-4 ${resolvedStyle.iconColor}`} />;
     };
 
     return (
-        <div className={`w-full max-w-full ${className || ""}`}>
+        <div className={cn("w-full max-w-full", resolvedStyle.className)}>
             <div
-                className={`border-l-4 bg-white dark:bg-gray-900 mb-4 w-full max-w-full overflow-hidden ${styles.borderColor} ${cardClassName || ""}`}
+                className={cn("border-l-4 bg-white dark:bg-gray-900 mb-4 w-full max-w-full overflow-hidden", resolvedStyle.borderColor, resolvedStyle.cardClassName)}
                 data-workstream-id={workstreamId}
             >
                 {/* Compact header */}
-                <div className={`flex items-center justify-between px-4 py-1.5 ${headerClassName || ""}`}>
+                <div className={cn("flex items-center justify-between px-4 py-1.5", resolvedStyle.headerClassName)}>
                     <div className="flex items-center gap-1.5">
-                        <div className={`${showPulsatingCircle ? "animate-fadeIn" : ""} ${iconClassName || ""}`}>
+                        <div className={cn(showPulsatingCircle ? "animate-fadeIn" : "", resolvedStyle.iconClassName)}>
                             {renderIcon()}
                         </div>
-                        <span className={`text-xs font-medium text-muted ${senderClassName || ""}`}>{styles.sender}</span>
+                        <span className={cn("text-xs font-medium text-muted", resolvedStyle.senderClassName)}>{resolvedStyle.sender}</span>
                         {workstreamId !== "main" && workstreamId !== "all" && (
                             <Badge variant="default" className="text-xs text-muted ml-1">
                                 {workstreamId}
@@ -372,7 +470,7 @@ function MessageItemComponent({
                         )}
                     </div>
                     <div className="flex items-center gap-1.5 print:hidden">
-                        <span className={`text-[11px] text-muted/70 ${timestampClassName || ""}`}>
+                        <span className={cn("text-[11px] text-muted/70", resolvedStyle.timestampClassName)}>
                             {dayjs(message.timestamp).format("HH:mm:ss")}
                         </span>
                         <Button
@@ -383,11 +481,32 @@ function MessageItemComponent({
                         >
                             <CopyIcon className="size-3" />
                         </Button>
+                        {hasExportableContent && (
+                            <Dropdown
+                                trigger={
+                                    <Button
+                                        variant="ghost" size="xs"
+                                        className="text-muted/50 hover:text-muted h-5 w-5 p-0"
+                                        title="Export message"
+                                        disabled={isExportingFile}
+                                    >
+                                        <Download className={`size-3 ${isExportingFile ? 'animate-pulse' : ''}`} />
+                                    </Button>
+                                }
+                            >
+                                <MenuItem onClick={() => exportToFormat(MarkdownRenditionFormat.pdf)}>
+                                    Export as PDF
+                                </MenuItem>
+                                <MenuItem onClick={() => exportToFormat(MarkdownRenditionFormat.docx)}>
+                                    Export as Word
+                                </MenuItem>
+                            </Dropdown>
+                        )}
                     </div>
                 </div>
 
                 {/* Message content */}
-                <div className={`px-4 pb-3 bg-white dark:bg-gray-900 overflow-hidden ${contentClassName || ""}`}>
+                <div className={cn("px-4 pb-3 bg-white dark:bg-gray-900 overflow-hidden", resolvedStyle.contentClassName)}>
                 {/* Check for REQUEST_INPUT with UX config - render AskUserWidget instead of plain text */}
                 {message.type === AgentMessageType.REQUEST_INPUT && (message.details as AskUserMessageDetails)?.ux ? (
                     (() => {
@@ -415,7 +534,7 @@ function MessageItemComponent({
 
                 {/* Auto-surfaced artifacts from tool details (e.g. execute_shell.outputFiles) */}
                 {artifactLinks.length > 0 && (
-                    <div className={`mt-3 text-xs ${artifactsClassName || ""}`}>
+                    <div className={cn("mt-3 text-xs", resolvedStyle.artifactsClassName)}>
                         <div className="font-medium text-muted mb-1">Artifacts</div>
 
                         {/* Inline previews for image artifacts */}
@@ -462,10 +581,10 @@ function MessageItemComponent({
 
                 {/* Optional details section */}
                 {message.details && (
-                    <div className={`mt-2 print:hidden ${detailsClassName || ""}`}>
+                    <div className={cn("mt-2 print:hidden", resolvedStyle.detailsClassName)}>
                         <button
                             onClick={() => setShowDetails(!showDetails)}
-                            className="text-xs text-muted flex items-center"
+                            className="text-[11px] text-muted flex items-center"
                         >
                             {showDetails ? "Hide" : "Show"} details
                             <svg
@@ -480,11 +599,11 @@ function MessageItemComponent({
                         </button>
 
                         {showDetails && (
-                            <div className="mt-2 p-2 bg-muted border border-mixer-muted/40 rounded text-sm">
+                            <div className="mt-1 p-1.5 bg-muted border border-mixer-muted/40 rounded text-sm">
                                 {typeof message.details === "string" ? (
                                     renderContent(message.details)
                                 ) : (
-                                    <pre className="text-xs font-mono whitespace-pre-wrap overflow-x-auto bg-muted p-2 rounded text-muted ">
+                                    <pre className="text-xs font-mono whitespace-pre-wrap overflow-x-auto bg-muted p-2 rounded text-muted">
                                         {JSON.stringify(message.details, null, 2)}
                                     </pre>
                                 )}
@@ -505,15 +624,8 @@ const MessageItem = memo(MessageItemComponent, (prevProps, nextProps) => {
         prevProps.message.timestamp === nextProps.message.timestamp &&
         prevProps.showPulsatingCircle === nextProps.showPulsatingCircle &&
         prevProps.onSendMessage === nextProps.onSendMessage &&
-        prevProps.className === nextProps.className &&
-        prevProps.cardClassName === nextProps.cardClassName &&
-        prevProps.headerClassName === nextProps.headerClassName &&
-        prevProps.contentClassName === nextProps.contentClassName &&
-        prevProps.timestampClassName === nextProps.timestampClassName &&
-        prevProps.senderClassName === nextProps.senderClassName &&
-        prevProps.iconClassName === nextProps.iconClassName &&
-        prevProps.detailsClassName === nextProps.detailsClassName &&
-        prevProps.artifactsClassName === nextProps.artifactsClassName
+        prevProps.messageStyleOverrides === nextProps.messageStyleOverrides &&
+        MESSAGE_ITEM_CLASS_NAME_KEYS.every(key => prevProps[key] === nextProps[key])
     );
 });
 

package/src/features/agent/chat/ModernAgentOutput/StreamingMessage.tsx

@@ -20,12 +20,25 @@ export interface StreamingMessageProps {
     timestamp?: number | string;
     /** Additional className for the outer container */
     className?: string;
+    /** Additional className for the card wrapper */
+    cardClassName?: string;
     /** Additional className for the header section */
     headerClassName?: string;
     /** Additional className for the content section */
     contentClassName?: string;
+    /** Additional className for the prose/markdown container */
+    proseClassName?: string;
+    /** Additional className for the sender label */
+    senderClassName?: string;
+    /** Additional className for the icon wrapper */
+    iconClassName?: string;
 }
 
+/** className overrides for StreamingMessage — subset of StreamingMessageProps containing only className props. */
+export type StreamingMessageClassNames = Partial<Pick<StreamingMessageProps,
+    'className' | 'cardClassName' | 'headerClassName' | 'contentClassName' |
+    'proseClassName' | 'senderClassName' | 'iconClassName'>>;
+
 /**
  * Displays a streaming message with adaptive reveal effect.
  * Automatically adjusts speed to match the model's generation rate,
@@ -38,8 +51,12 @@ function StreamingMessageComponent({
     isComplete = false,
     timestamp,
     className,
+    cardClassName,
     headerClassName,
     contentClassName,
+    proseClassName,
+    senderClassName,
+    iconClassName,
 }: StreamingMessageProps) {
     const [displayedLength, setDisplayedLength] = useState(0);
     const [throttledText, setThrottledText] = useState("");
@@ -247,20 +264,20 @@ function StreamingMessageComponent({
         <div className={cn("w-full max-w-full", className)}>
             {/* Card wrapper matching MessageItem structure */}
             <div
-                className="border-l-4 bg-white dark:bg-gray-900 mb-4 border-l-purple-500 w-full max-w-full overflow-hidden"
+                className={cn("border-l-4 bg-white dark:bg-gray-900 mb-4 border-l-purple-500 w-full max-w-full overflow-hidden", cardClassName)}
                 data-workstream-id={workstreamId}
             >
                 {/* Compact header */}
                 <div className={cn("flex items-center justify-between px-4 py-1.5", headerClassName)}>
                     <div className="flex items-center gap-1.5">
-                        <div className="animate-fadeIn">
+                        <div className={cn("animate-fadeIn", iconClassName)}>
                             {isTyping ? (
                                 <span className="size-2 rounded-full bg-blue-500 animate-pulse inline-block" />
                             ) : (
                                 <Bot className="size-4 text-purple-600 dark:text-purple-400" />
                             )}
                         </div>
-                        <span className="text-xs font-medium text-muted">Agent</span>
+                        <span className={cn("text-xs font-medium text-muted", senderClassName)}>Agent</span>
                         {workstreamId && workstreamId !== "main" && (
                             <span className="text-xs text-muted">• Task {workstreamId}</span>
                         )}
@@ -282,12 +299,12 @@ function StreamingMessageComponent({
                 {/* Content - cursor character is appended directly to text (no DOM manipulation) */}
                 <div
                     className={cn(
-                        "px-4 pb-3 streaming-content",
+                        "px-3 pb-2 streaming-content",
                         isTyping && "streaming-active",
                         contentClassName
                     )}
                 >
-                    <div className="vprose prose prose-slate dark:prose-invert prose-p:leading-relaxed prose-p:my-3 prose-headings:font-semibold prose-headings:tracking-normal prose-headings:mt-6 prose-headings:mb-3 prose-h1:text-xl prose-h2:text-lg prose-h3:text-base prose-li:my-1 prose-ul:my-3 prose-ol:my-3 prose-table:my-5 prose-pre:my-4 prose-hr:my-6 max-w-none text-[15px] break-words" style={{ overflowWrap: 'anywhere' }}>
+                    <div className={cn("vprose prose prose-slate dark:prose-invert prose-p:leading-relaxed prose-p:my-3 prose-headings:font-semibold prose-headings:tracking-normal prose-headings:mt-6 prose-headings:mb-3 prose-h1:text-xl prose-h2:text-lg prose-h3:text-base prose-li:my-1 prose-ul:my-3 prose-ol:my-3 prose-table:my-5 prose-pre:my-4 prose-hr:my-6 max-w-none text-[15px] break-words", proseClassName)} style={{ overflowWrap: 'anywhere' }}>
                         <MarkdownRenderer>
                             {displayTextWithCursor}
                         </MarkdownRenderer>