diff --git a/src/apps/chat/components/message/ChatMessage.tsx b/src/apps/chat/components/message/ChatMessage.tsx
index a6ba90f8b..d0546a570 100644
--- a/src/apps/chat/components/message/ChatMessage.tsx
+++ b/src/apps/chat/components/message/ChatMessage.tsx
@@ -1,10 +1,8 @@
 import * as React from 'react';
-import TimeAgo from 'react-timeago';
 import { shallow } from 'zustand/shallow';
 import { cleanupEfficiency, Diff as TextDiff, makeDiff } from '@sanity/diff-match-patch';
 
-import { Avatar, Box, Button, CircularProgress, IconButton, ListDivider, ListItem, ListItemDecorator, MenuItem, Switch, Tooltip, Typography } from '@mui/joy';
-import { SxProps } from '@mui/joy/styles/types';
+import { Avatar, Box, CircularProgress, IconButton, ListDivider, ListItem, ListItemDecorator, MenuItem, Switch, Tooltip, Typography } from '@mui/joy';
 import AccountTreeIcon from '@mui/icons-material/AccountTree';
 import ClearIcon from '@mui/icons-material/Clear';
 import ContentCopyIcon from '@mui/icons-material/ContentCopy';
@@ -23,30 +21,18 @@ import VerticalAlignBottomIcon from '@mui/icons-material/VerticalAlignBottom';
 
 import { CloseableMenu } from '~/common/components/CloseableMenu';
 import { DMessage } from '~/common/state/store-chats';
-import { InlineError } from '~/common/components/InlineError';
 import { InlineTextarea } from '~/common/components/InlineTextarea';
 import { KeyStroke } from '~/common/components/KeyStroke';
 import { Link } from '~/common/components/Link';
 import { SystemPurposeId, SystemPurposes } from '../../../../data';
 import { copyToClipboard } from '~/common/util/clipboardUtils';
-import { cssRainbowColorKeyframes, lineHeightChatText } from '~/common/app.theme';
+import { cssRainbowColorKeyframes } from '~/common/app.theme';
 import { prettyBaseModel } from '~/common/util/modelUtils';
 import { useUIPreferencesStore } from '~/common/state/store-ui';
 
+import { BlocksRenderer, editBlocksSx } from './blocks/BlocksRenderer';
 import { useChatShowTextDiff } from '../../store-app-chat';
 
-import { RenderCode } from './RenderCode';
-import { RenderHtml } from './RenderHtml';
-import { RenderImage } from './RenderImage';
-import { RenderLatex } from './RenderLatex';
-import { RenderMarkdown } from './RenderMarkdown';
-import { RenderText } from './RenderText';
-import { RenderTextDiff } from './RenderTextDiff';
-import { parseBlocks } from './blocks';
-
-
-// How long is the user collapsed message
-const USER_COLLAPSED_LINES: number = 8;
 
 // Enable the menu on text selection
 const ENABLE_SELECTION_RIGHT_CLICK_MENU: boolean = true;
@@ -206,13 +192,13 @@ export const ChatMessageMemo = React.memo(ChatMessage);
  * or collapsing long user messages.
  *
  */
-export function ChatMessage(props: {
+function ChatMessage(props: {
   message: DMessage,
-  showDate?: boolean, diffPreviousText?: string,
-  hideAvatars?: boolean, codeBackground?: string,
-  noMarkdown?: boolean, diagramMode?: boolean,
-  isBottom?: boolean, noBottomBorder?: boolean,
-  isImagining?: boolean, isSpeaking?: boolean,
+  diffPreviousText?: string,
+  isBottom?: boolean,
+  isImagining?: boolean,
+  isSpeaking?: boolean,
+  blocksShowDate?: boolean,
   onConversationBranch?: (messageId: string) => void,
   onConversationRestartFrom?: (messageId: string, offset: number) => Promise<void>,
   onConversationTruncate?: (messageId: string) => void,
@@ -221,11 +207,9 @@ export function ChatMessage(props: {
   onTextDiagram?: (messageId: string, text: string) => Promise<void>
   onTextImagine?: (text: string) => Promise<void>
   onTextSpeak?: (text: string) => Promise<void>
-  sx?: SxProps,
 }) {
 
   // state
-  const [forceUserExpanded, setForceUserExpanded] = React.useState(false);
   const [isHovering, setIsHovering] = React.useState(false);
   const [opsMenuAnchor, setOpsMenuAnchor] = React.useState<HTMLElement | null>(null);
   const [selMenuAnchor, setSelMenuAnchor] = React.useState<HTMLElement | null>(null);
@@ -257,10 +241,9 @@ export function ChatMessage(props: {
 
   const fromAssistant = messageRole === 'assistant';
   const fromSystem = messageRole === 'system';
-  const fromUser = messageRole === 'user';
   const wasEdited = !!messageUpdated;
 
-  const showAvatars = props.hideAvatars !== true && !cleanerLooks;
+  const showAvatars = !cleanerLooks;
 
   const textSel = selMenuText ? selMenuText : messageText;
   const isSpecialT2I = textSel.startsWith('https://images.prodia.xyz/') || textSel.startsWith('/draw ') || textSel.startsWith('/imagine ') || textSel.startsWith('/img ');
@@ -275,8 +258,6 @@ export function ChatMessage(props: {
       props.onMessageEdit(messageId, editedText);
   };
 
-  const handleUncollapse = () => setForceUserExpanded(true);
-
 
   // Operations Menu
 
@@ -289,12 +270,12 @@ export function ChatMessage(props: {
     closeSelectionMenu();
   };
 
-  const handleOpsEdit = (e: React.MouseEvent) => {
+  const handleOpsEdit = React.useCallback((e: React.MouseEvent) => {
     if (messageTyping && !isEditing) return; // don't allow editing while typing
     setIsEditing(!isEditing);
     e.preventDefault();
     closeOpsMenu();
-  };
+  }, [isEditing, messageTyping]);
 
   const handleOpsConversationBranch = (e: React.MouseEvent) => {
     e.preventDefault();
@@ -396,6 +377,17 @@ export function ChatMessage(props: {
   }, [openSelectionMenu]);
 
 
+  // Blocks renderer
+
+  const handleBlocksContextMenu = React.useCallback((event: React.MouseEvent) => {
+    handleMouseUp(event.nativeEvent);
+  }, [handleMouseUp]);
+
+  const handleBlocksDoubleClick = React.useCallback((event: React.MouseEvent) => {
+    doubleClickToEdit && props.onMessageEdit && handleOpsEdit(event);
+  }, [doubleClickToEdit, handleOpsEdit, props.onMessageEdit]);
+
+
   // prettier upstream errors
   const { isAssistantError, errorMessage } = React.useMemo(
     () => explainErrorInMessage(messageText, fromAssistant, messageOriginLLM),
@@ -411,50 +403,19 @@ export function ChatMessage(props: {
     [messageAvatar, messageOriginLLM, messagePurposeId, messageRole, messageSender, messageTyping, showAvatars],
   );
 
-  // per-blocks css
-  const blockSx: SxProps = {
-    my: 'auto',
-    lineHeight: lineHeightChatText,
-  };
-  const typographySx: SxProps = {
-    lineHeight: lineHeightChatText,
-  };
-  const codeSx: SxProps = {
-    // backgroundColor: fromAssistant ? 'background.level1' : 'background.level1',
-    backgroundColor: props.codeBackground ? props.codeBackground : fromAssistant ? 'neutral.plainHoverBg' : 'primary.plainActiveBg',
-    boxShadow: 'xs',
-    fontFamily: 'code',
-    fontSize: '0.875rem',
-    fontVariantLigatures: 'none',
-    lineHeight: lineHeightChatText,
-    borderRadius: 'var(--joy-radius-sm)',
-  };
-
-  // user message truncation
-  let collapsedText = messageText;
-  let isCollapsed = false;
-  if (fromUser && !forceUserExpanded) {
-    const lines = messageText.split('\n');
-    if (lines.length > USER_COLLAPSED_LINES) {
-      collapsedText = lines.slice(0, USER_COLLAPSED_LINES).join('\n');
-      isCollapsed = true;
-    }
-  }
-
 
   return (
     <ListItem
       sx={{
         display: 'flex', flexDirection: !fromAssistant ? 'row-reverse' : 'row', alignItems: 'flex-start',
-        gap: { xs: 0, md: 1 }, px: { xs: 1, md: 2 }, py: 2,
+        gap: { xs: 0, md: 1 },
+        px: { xs: 1, md: 2 },
+        py: 2,
         backgroundColor,
-        ...(props.noBottomBorder !== true && {
-          borderBottom: '1px solid',
-          borderBottomColor: 'divider',
-        }),
+        borderBottom: '1px solid',
+        borderBottomColor: 'divider',
         ...(ENABLE_COPY_MESSAGE_OVERLAY && { position: 'relative' }),
         '&:hover > button': { opacity: 1 },
-        ...props.sx,
       }}
     >
 
@@ -500,73 +461,24 @@ export function ChatMessage(props: {
 
         <InlineTextarea
           initialText={messageText} onEdit={handleTextEdited}
-          sx={{
-            ...blockSx,
-            flexGrow: 1,
-          }} />
+          sx={editBlocksSx}
+        />
 
       ) : (
 
-        <Box
-          onContextMenu={(ENABLE_SELECTION_RIGHT_CLICK_MENU && props.onMessageEdit) ? event => handleMouseUp(event.nativeEvent) : undefined}
-          onDoubleClick={event => (doubleClickToEdit && props.onMessageEdit) ? handleOpsEdit(event) : null}
-          sx={{
-            ...blockSx,
-            flexGrow: 0,
-            overflowX: 'auto',
-            ...(!!props.diagramMode && {
-              // width: '100%',
-              boxShadow: 'md',
-            }),
-          }}>
-
-          {props.showDate === true && (
-            <Typography level='body-sm' sx={{ mx: 1.5, textAlign: fromAssistant ? 'left' : 'right' }}>
-              <TimeAgo date={messageUpdated || messageCreated} />
-            </Typography>
-          )}
-
-          {/* Warn about user-edited system message */}
-          {fromSystem && wasEdited && (
-            <Typography level='body-sm' color='warning' sx={{ mt: 1, mx: 1.5 }}>modified by user - auto-update disabled</Typography>
-          )}
-
-          {errorMessage && (
-            <Tooltip title={<Typography sx={{ maxWidth: 800 }}>{collapsedText}</Typography>} variant='soft'>
-              <InlineError error={errorMessage} />
-            </Tooltip>
-          )}
+        <BlocksRenderer
+          text={messageText}
+          fromRole={messageRole}
+          renderTextAsMarkdown={renderMarkdown}
+          errorMessage={errorMessage}
+          isBottom={props.isBottom}
+          showDate={props.blocksShowDate === true ? messageUpdated || messageCreated || undefined : undefined}
+          textDiffs={textDiffs || undefined}
+          wasUserEdited={wasEdited}
+          onContextMenu={(props.onMessageEdit && ENABLE_SELECTION_RIGHT_CLICK_MENU) ? handleBlocksContextMenu : undefined}
+          onDoubleClick={(props.onMessageEdit && doubleClickToEdit) ? handleBlocksDoubleClick : undefined}
+        />
 
-          {/* sequence of render components, for each Block */}
-          {!errorMessage && parseBlocks(collapsedText, fromSystem, textDiffs)
-            .filter((block, _, blocks) => !props.diagramMode || block.type === 'code' || blocks.length === 1)
-            .map(
-              (block, index) =>
-                block.type === 'html'
-                  ? <RenderHtml key={'html-' + index} htmlBlock={block} sx={codeSx} />
-                  : block.type === 'code'
-                    ? <RenderCode key={'code-' + index} codeBlock={block} sx={codeSx} noCopyButton={props.diagramMode} />
-                    : block.type === 'image'
-                      ? <RenderImage key={'image-' + index} imageBlock={block} isFirst={!index} allowRunAgain={props.isBottom === true} onRunAgain={handleOpsConversationRestartFrom} />
-                      : block.type === 'latex'
-                        ? <RenderLatex key={'latex-' + index} latexBlock={block} sx={typographySx} />
-                        : block.type === 'diff'
-                          ? <RenderTextDiff key={'latex-' + index} diffBlock={block} sx={typographySx} />
-                          : (renderMarkdown && props.noMarkdown !== true && !fromSystem && !(fromUser && block.content.startsWith('/')))
-                            ? <RenderMarkdown key={'text-md-' + index} textBlock={block} />
-                            : <RenderText key={'text-' + index} textBlock={block} sx={typographySx} />)}
-
-          {isCollapsed && (
-            <Button variant='plain' color='neutral' onClick={handleUncollapse}>... expand ...</Button>
-          )}
-
-          {/* import VisibilityIcon from '@mui/icons-material/Visibility'; */}
-          {/*<br />*/}
-          {/*<Chip variant='outlined' color='warning' sx={{ mt: 1, fontSize: '0.75em' }} startDecorator={<VisibilityIcon />}>*/}
-          {/*  BlockAction*/}
-          {/*</Chip>*/}
-
-        </Box>
       )}
 
 
diff --git a/src/apps/chat/components/message/blocks/BlocksRenderer.tsx b/src/apps/chat/components/message/blocks/BlocksRenderer.tsx
new file mode 100644
index 000000000..14c6377cc
--- /dev/null
+++ b/src/apps/chat/components/message/blocks/BlocksRenderer.tsx
@@ -0,0 +1,167 @@
+import * as React from 'react';
+import TimeAgo from 'react-timeago';
+import type { Diff as TextDiff } from '@sanity/diff-match-patch';
+
+import type { SxProps } from '@mui/joy/styles/types';
+import { Box, Button, Tooltip, Typography } from '@mui/joy';
+
+import type { DMessage } from '~/common/state/store-chats';
+import { InlineError } from '~/common/components/InlineError';
+import { lineHeightChatText } from '~/common/app.theme';
+
+import { RenderCode } from './code/RenderCode';
+import { RenderHtml } from './RenderHtml';
+import { RenderImage } from './RenderImage';
+import { RenderLatex } from './RenderLatex';
+import { RenderMarkdown } from './RenderMarkdown';
+import { RenderText } from './RenderText';
+import { RenderTextDiff } from './RenderTextDiff';
+import { parseMessageBlocks } from './blocks';
+
+
+// How long is the user collapsed message
+const USER_COLLAPSED_LINES: number = 8;
+
+
+const blocksSx: SxProps = {
+  my: 'auto',
+  lineHeight: lineHeightChatText,
+} as const;
+
+export const editBlocksSx: SxProps = {
+  ...blocksSx,
+  flexGrow: 1,
+} as const;
+
+const renderBlocksSx: SxProps = {
+  ...blocksSx,
+  flexGrow: 0,
+  overflowX: 'auto',
+} as const;
+
+const typographySx: SxProps = {
+  lineHeight: lineHeightChatText,
+} as const;
+
+
+export function BlocksRenderer(props: {
+
+  // required
+  text: string;
+  fromRole: DMessage['role'];
+  renderTextAsMarkdown: boolean;
+
+  errorMessage?: React.ReactNode;
+  isBottom?: boolean;
+  showDate?: number;
+  textDiffs?: TextDiff[];
+  wasUserEdited?: boolean;
+
+  specialDiagramMode?: boolean;
+
+  onContextMenu?: (event: React.MouseEvent) => void;
+  onDoubleClick?: (event: React.MouseEvent) => void;
+  onImageRegenerate?: () => void;
+
+}) {
+
+  // state
+  const [forceUserExpanded, setForceUserExpanded] = React.useState(false);
+
+  // derived state
+  const { text: _text, textDiffs, errorMessage, wasUserEdited = false } = props;
+  const fromAssistant = props.fromRole === 'assistant';
+  const fromSystem = props.fromRole === 'system';
+  const fromUser = props.fromRole === 'user';
+
+
+  // Memo text, blocks and styles
+
+  const { text, isTextCollapsed } = React.useMemo(() => {
+    if (fromUser && !forceUserExpanded) {
+      const textLines = _text.split('\n');
+      if (textLines.length > USER_COLLAPSED_LINES)
+        return { text: textLines.slice(0, USER_COLLAPSED_LINES).join('\n'), isTextCollapsed: true };
+    }
+    return { text: _text, isTextCollapsed: false };
+  }, [forceUserExpanded, fromUser, _text]);
+
+  const blocks = React.useMemo(() => {
+    const blocks = errorMessage ? [] : parseMessageBlocks(text, fromSystem, textDiffs || null);
+    return props.specialDiagramMode ? blocks.filter(block => block.type === 'code' || blocks.length === 1) : blocks;
+  }, [errorMessage, fromSystem, props.specialDiagramMode, text, textDiffs]);
+
+  const codeSx: SxProps = React.useMemo(() => (
+    {
+      backgroundColor: props.specialDiagramMode ? 'background.surface' : fromAssistant ? 'neutral.plainHoverBg' : 'primary.plainActiveBg',
+      boxShadow: 'xs',
+      fontFamily: 'code',
+      fontSize: '0.875rem',
+      fontVariantLigatures: 'none',
+      lineHeight: lineHeightChatText,
+      borderRadius: 'var(--joy-radius-sm)',
+    }
+  ), [fromAssistant, props.specialDiagramMode]);
+
+
+  const handleTextUncollapse = React.useCallback(() => {
+    setForceUserExpanded(true);
+  }, []);
+
+
+  return (
+    <Box
+      onContextMenu={props.onContextMenu}
+      onDoubleClick={props.onDoubleClick}
+      sx={renderBlocksSx}
+    >
+
+      {!!props.showDate && (
+        <Typography level='body-sm' sx={{ mx: 1.5, textAlign: fromAssistant ? 'left' : 'right' }}>
+          <TimeAgo date={props.showDate} />
+        </Typography>
+      )}
+
+      {/* Warn about user-edited system message */}
+      {fromSystem && wasUserEdited && (
+        <Typography level='body-sm' color='warning' sx={{ mt: 1, mx: 1.5 }}>modified by user - auto-update disabled</Typography>
+      )}
+
+      {errorMessage ? (
+
+        <Tooltip title={<Typography sx={{ maxWidth: 800 }}>{text}</Typography>} variant='soft'>
+          <InlineError error={errorMessage} />
+        </Tooltip>
+
+      ) : (
+
+        // sequence of render components, for each Block
+        blocks.map(
+          (block, index) =>
+            block.type === 'html'
+              ? <RenderHtml key={'html-' + index} htmlBlock={block} sx={codeSx} />
+              : block.type === 'code'
+                ? <RenderCode key={'code-' + index} codeBlock={block} sx={codeSx} noCopyButton={props.specialDiagramMode} />
+                : block.type === 'image'
+                  ? <RenderImage key={'image-' + index} imageBlock={block} isFirst={!index} allowRunAgain={props.isBottom === true} onRunAgain={props.onImageRegenerate} />
+                  : block.type === 'latex'
+                    ? <RenderLatex key={'latex-' + index} latexBlock={block} sx={typographySx} />
+                    : block.type === 'diff'
+                      ? <RenderTextDiff key={'latex-' + index} diffBlock={block} sx={typographySx} />
+                      : (props.renderTextAsMarkdown && !fromSystem && !(fromUser && block.content.startsWith('/')))
+                        ? <RenderMarkdown key={'text-md-' + index} textBlock={block} />
+                        : <RenderText key={'text-' + index} textBlock={block} sx={typographySx} />)
+
+      )}
+
+      {isTextCollapsed && <Button variant='plain' color='neutral' onClick={handleTextUncollapse}>... expand ...</Button>}
+
+      {/* import VisibilityIcon from '@mui/icons-material/Visibility'; */}
+      {/*<br />*/}
+      {/*<Chip variant='outlined' color='warning' sx={{ mt: 1, fontSize: '0.75em' }} startDecorator={<VisibilityIcon />}>*/}
+      {/*  BlockAction*/}
+      {/*</Chip>*/}
+
+    </Box>
+  );
+}
\ No newline at end of file
diff --git a/src/apps/chat/components/message/RenderHtml.tsx b/src/apps/chat/components/message/blocks/RenderHtml.tsx
similarity index 96%
rename from src/apps/chat/components/message/RenderHtml.tsx
rename to src/apps/chat/components/message/blocks/RenderHtml.tsx
index 2cd4d02c3..ce32a428b 100644
--- a/src/apps/chat/components/message/RenderHtml.tsx
+++ b/src/apps/chat/components/message/blocks/RenderHtml.tsx
@@ -1,14 +1,14 @@
 import * as React from 'react';
 
+import type { SxProps } from '@mui/joy/styles/types';
 import { Box, Button, IconButton, Tooltip, Typography } from '@mui/joy';
-import { SxProps } from '@mui/joy/styles/types';
 import ContentCopyIcon from '@mui/icons-material/ContentCopy';
 import WebIcon from '@mui/icons-material/Web';
 
 import { copyToClipboard } from '~/common/util/clipboardUtils';
 
-import { HtmlBlock } from './blocks';
-import { overlayButtonsSx } from './RenderCode';
+import type { HtmlBlock } from './blocks';
+import { overlayButtonsSx } from './code/RenderCode';
 
 
 // this is used by the blocks parser (for full text detection) and by the Code component (for inline rendering)
diff --git a/src/apps/chat/components/message/RenderImage.tsx b/src/apps/chat/components/message/blocks/RenderImage.tsx
similarity index 98%
rename from src/apps/chat/components/message/RenderImage.tsx
rename to src/apps/chat/components/message/blocks/RenderImage.tsx
index ac85ed6d1..d0c406265 100644
--- a/src/apps/chat/components/message/RenderImage.tsx
+++ b/src/apps/chat/components/message/blocks/RenderImage.tsx
@@ -6,8 +6,8 @@ import ReplayIcon from '@mui/icons-material/Replay';
 
 import { Link } from '~/common/components/Link';
 
-import { ImageBlock } from './blocks';
-import { overlayButtonsSx } from './RenderCode';
+import type { ImageBlock } from './blocks';
+import { overlayButtonsSx } from './code/RenderCode';
 
 
 const mdImageReferenceRegex = /^!\[([^\]]*)]\(([^)]+)\)$/;
diff --git a/src/apps/chat/components/message/RenderLatex.tsx b/src/apps/chat/components/message/blocks/RenderLatex.tsx
similarity index 87%
rename from src/apps/chat/components/message/RenderLatex.tsx
rename to src/apps/chat/components/message/blocks/RenderLatex.tsx
index fb3814c98..5c51a462f 100644
--- a/src/apps/chat/components/message/RenderLatex.tsx
+++ b/src/apps/chat/components/message/blocks/RenderLatex.tsx
@@ -3,7 +3,7 @@ import * as React from 'react';
 import { Box } from '@mui/joy';
 import { SxProps } from '@mui/joy/styles/types';
 
-import { LatexBlock } from './blocks';
+import type { LatexBlock } from './blocks';
 
 
 // Dynamically import the Katex functions
@@ -20,7 +20,7 @@ export const RenderLatex = ({ latexBlock, sx }: { latexBlock: LatexBlock; sx?: S
       mx: 1.5,
       ...(sx || {}),
     }}>
-    <React.Suspense fallback={<div/>}>
+    <React.Suspense fallback={<div />}>
       <RenderLatexDynamic latex={latexBlock.latex} />
     </React.Suspense>
   </Box>;
\ No newline at end of file
diff --git a/src/apps/chat/components/message/RenderMarkdown.tsx b/src/apps/chat/components/message/blocks/RenderMarkdown.tsx
similarity index 99%
rename from src/apps/chat/components/message/RenderMarkdown.tsx
rename to src/apps/chat/components/message/blocks/RenderMarkdown.tsx
index bafeff0fb..f76f7ce2f 100644
--- a/src/apps/chat/components/message/RenderMarkdown.tsx
+++ b/src/apps/chat/components/message/blocks/RenderMarkdown.tsx
@@ -1,14 +1,13 @@
 import * as React from 'react';
-
 import { CSVLink } from 'react-csv';
 
 import { Box, Button, styled } from '@mui/joy';
+import DownloadIcon from '@mui/icons-material/Download';
 
 import { lineHeightChatText } from '~/common/app.theme';
 
 import type { TextBlock } from './blocks';
 
-import DownloadIcon from '@mui/icons-material/Download';
 
 /*
  * For performance reasons, we style this component here and copy the equivalent of 'props.sx' (the lineHeight) locally.
diff --git a/src/apps/chat/components/message/RenderText.tsx b/src/apps/chat/components/message/blocks/RenderText.tsx
similarity index 93%
rename from src/apps/chat/components/message/RenderText.tsx
rename to src/apps/chat/components/message/blocks/RenderText.tsx
index dcb6b221a..b819588e6 100644
--- a/src/apps/chat/components/message/RenderText.tsx
+++ b/src/apps/chat/components/message/blocks/RenderText.tsx
@@ -3,7 +3,7 @@ import * as React from 'react';
 import { Chip, Typography } from '@mui/joy';
 import { SxProps } from '@mui/joy/styles/types';
 
-import { extractChatCommand } from '../../commands/commands.registry';
+import { extractChatCommand } from '../../../commands/commands.registry';
 
 import type { TextBlock } from './blocks';
 
diff --git a/src/apps/chat/components/message/RenderTextDiff.tsx b/src/apps/chat/components/message/blocks/RenderTextDiff.tsx
similarity index 97%
rename from src/apps/chat/components/message/RenderTextDiff.tsx
rename to src/apps/chat/components/message/blocks/RenderTextDiff.tsx
index 1960465c3..35e63a299 100644
--- a/src/apps/chat/components/message/RenderTextDiff.tsx
+++ b/src/apps/chat/components/message/blocks/RenderTextDiff.tsx
@@ -4,7 +4,7 @@ import { Diff as TextDiff, DIFF_DELETE, DIFF_INSERT } from '@sanity/diff-match-p
 import { Box, Typography, useTheme } from '@mui/joy';
 import { SxProps } from '@mui/joy/styles/types';
 
-import { DiffBlock } from './blocks';
+import type { DiffBlock } from './blocks';
 
 
 export const RenderTextDiff = ({ diffBlock, sx }: { diffBlock: DiffBlock; sx?: SxProps; }) => {
diff --git a/src/apps/chat/components/message/blocks.ts b/src/apps/chat/components/message/blocks/blocks.ts
similarity index 90%
rename from src/apps/chat/components/message/blocks.ts
rename to src/apps/chat/components/message/blocks/blocks.ts
index d0cdd2a4a..be26ef8c1 100644
--- a/src/apps/chat/components/message/blocks.ts
+++ b/src/apps/chat/components/message/blocks/blocks.ts
@@ -1,9 +1,10 @@
 import type { Diff as TextDiff } from '@sanity/diff-match-patch';
 
 import { heuristicIsHtml } from './RenderHtml';
-import { heuristicMarkdownImageReferenceBlocks, heuristicLegacyImageBlocks } from './RenderImage';
+import { heuristicLegacyImageBlocks, heuristicMarkdownImageReferenceBlocks } from './RenderImage';
 
-type Block = CodeBlock | DiffBlock | HtmlBlock | ImageBlock | LatexBlock | TextBlock;
+// Block types
+export type Block = CodeBlock | DiffBlock | HtmlBlock | ImageBlock | LatexBlock | TextBlock;
 export type CodeBlock = { type: 'code'; blockTitle: string; blockCode: string; complete: boolean; };
 export type DiffBlock = { type: 'diff'; textDiffs: TextDiff[] };
 export type HtmlBlock = { type: 'html'; html: string; };
@@ -12,9 +13,10 @@ export type LatexBlock = { type: 'latex'; latex: string; };
 export type TextBlock = { type: 'text'; content: string; }; // for Text or Markdown
 
 
-export function parseBlocks(text: string, forceText: boolean, textDiffs: TextDiff[] | null): Block[] {
+export function parseMessageBlocks(text: string, forceText: boolean, textDiffs: TextDiff[] | null): Block[] {
   if (forceText)
     return [{ type: 'text', content: text }];
+
   if (textDiffs && textDiffs.length > 0)
     return [{ type: 'diff', textDiffs }];
 
diff --git a/src/apps/chat/components/message/OpenInCodepen.tsx b/src/apps/chat/components/message/blocks/code/ButtonCodepen.tsx
similarity index 95%
rename from src/apps/chat/components/message/OpenInCodepen.tsx
rename to src/apps/chat/components/message/blocks/code/ButtonCodepen.tsx
index 6f2f5676d..8cc934f8d 100644
--- a/src/apps/chat/components/message/OpenInCodepen.tsx
+++ b/src/apps/chat/components/message/blocks/code/ButtonCodepen.tsx
@@ -9,7 +9,7 @@ interface CodeBlockProps {
   };
 }
 
-export function OpenInCodepen({ codeBlock }: CodeBlockProps): React.JSX.Element {
+export function ButtonCodepen({ codeBlock }: CodeBlockProps): React.JSX.Element {
   const { code, language } = codeBlock;
   const hasCSS = language === 'css';
   const hasJS = ['javascript', 'json', 'typescript'].includes(language || '');
diff --git a/src/apps/chat/components/message/OpenInReplit.tsx b/src/apps/chat/components/message/blocks/code/ButtonReplit.tsx
similarity index 92%
rename from src/apps/chat/components/message/OpenInReplit.tsx
rename to src/apps/chat/components/message/blocks/code/ButtonReplit.tsx
index c15448a4f..fd5b3cf5b 100644
--- a/src/apps/chat/components/message/OpenInReplit.tsx
+++ b/src/apps/chat/components/message/blocks/code/ButtonReplit.tsx
@@ -9,7 +9,7 @@ interface CodeBlockProps {
   };
 }
 
-export function OpenInReplit({ codeBlock }: CodeBlockProps): React.JSX.Element {
+export function ButtonReplit({ codeBlock }: CodeBlockProps): React.JSX.Element {
   const { language } = codeBlock;
 
   const replitLanguageMap: Record<string, string> = {
diff --git a/src/apps/chat/components/message/RenderCode.tsx b/src/apps/chat/components/message/blocks/code/RenderCode.tsx
similarity index 80%
rename from src/apps/chat/components/message/RenderCode.tsx
rename to src/apps/chat/components/message/blocks/code/RenderCode.tsx
index 0a30beb5b..f060751ac 100644
--- a/src/apps/chat/components/message/RenderCode.tsx
+++ b/src/apps/chat/components/message/blocks/code/RenderCode.tsx
@@ -1,8 +1,8 @@
 import * as React from 'react';
 import { useQuery } from '@tanstack/react-query';
 
+import type { SxProps } from '@mui/joy/styles/types';
 import { Box, IconButton, Sheet, Tooltip, Typography } from '@mui/joy';
-import { SxProps } from '@mui/joy/styles/types';
 import ContentCopyIcon from '@mui/icons-material/ContentCopy';
 import HtmlIcon from '@mui/icons-material/Html';
 import SchemaIcon from '@mui/icons-material/Schema';
@@ -10,11 +10,42 @@ import ShapeLineOutlinedIcon from '@mui/icons-material/ShapeLineOutlined';
 
 import { copyToClipboard } from '~/common/util/clipboardUtils';
 
-import { CodeBlock } from './blocks';
-import { OpenInCodepen } from './OpenInCodepen';
-import { OpenInReplit } from './OpenInReplit';
+import type { CodeBlock } from '../blocks';
+import { ButtonCodepen } from './ButtonCodepen';
+import { ButtonReplit } from './ButtonReplit';
 import { RenderCodeMermaid } from './RenderCodeMermaid';
-import { heuristicIsHtml, IFrameComponent } from './RenderHtml';
+import { heuristicIsHtml, IFrameComponent } from '../RenderHtml';
+
+
+async function fetchPlantUmlSvg(plantUmlCode: string): Promise<string | null> {
+  // fetch the PlantUML SVG
+  let text: string = '';
+  try {
+    // Dynamically import the PlantUML encoder - it's a large library that slows down app loading
+    const { encode: plantUmlEncode } = await import('plantuml-encoder');
+
+    // retrieve and manually adapt the SVG, to remove the background
+    const encodedPlantUML: string = plantUmlEncode(plantUmlCode);
+    const response = await fetch(`https://www.plantuml.com/plantuml/svg/${encodedPlantUML}`);
+    text = await response.text();
+  } catch (e) {
+    return null;
+  }
+  // validate/extract the SVG
+  const start = text.indexOf('<svg ');
+  const end = text.indexOf('</svg>');
+  if (start < 0 || end <= start)
+    throw new Error('Could not render PlantUML');
+  const svg = text
+    .slice(start, end + 6) // <svg ... </svg>
+    .replace('background:#FFFFFF;', ''); // transparent background
+
+  // check for syntax errors
+  if (svg.includes('>Syntax Error?</text>'))
+    throw new Error('syntax issue (it happens!). Please regenerate or change generator model.');
+
+  return svg;
+}
 
 
 export const overlayButtonsSx: SxProps = {
@@ -43,12 +74,11 @@ function RenderCodeImpl(props: {
   } = props;
 
   // heuristic for language, and syntax highlight
-  const { highlightedCode, inferredCodeLanguage } = React.useMemo(
-    () => {
-      const inferredCodeLanguage = inferCodeLanguage(blockTitle, blockCode);
-      const highlightedCode = highlightCode(inferredCodeLanguage, blockCode);
-      return { highlightedCode, inferredCodeLanguage };
-    }, [inferCodeLanguage, blockTitle, blockCode, highlightCode]);
+  const { highlightedCode, inferredCodeLanguage } = React.useMemo(() => {
+    const inferredCodeLanguage = inferCodeLanguage(blockTitle, blockCode);
+    const highlightedCode = highlightCode(inferredCodeLanguage, blockCode);
+    return { highlightedCode, inferredCodeLanguage };
+  }, [inferCodeLanguage, blockTitle, blockCode, highlightCode]);
 
 
   // heuristics for specialized rendering
@@ -70,35 +100,7 @@ function RenderCodeImpl(props: {
   const { data: plantUmlHtmlData, error: plantUmlError } = useQuery({
     enabled: renderPlantUML,
     queryKey: ['plantuml', blockCode],
-    queryFn: async () => {
-      // fetch the PlantUML SVG
-      let text: string = '';
-      try {
-        // Dynamically import the PlantUML encoder - it's a large library that slows down app loading
-        const { encode: plantUmlEncode } = await import('plantuml-encoder');
-
-        // retrieve and manually adapt the SVG, to remove the background
-        const encodedPlantUML: string = plantUmlEncode(blockCode);
-        const response = await fetch(`https://www.plantuml.com/plantuml/svg/${encodedPlantUML}`);
-        text = await response.text();
-      } catch (e) {
-        return null;
-      }
-      // validate/extract the SVG
-      const start = text.indexOf('<svg ');
-      const end = text.indexOf('</svg>');
-      if (start < 0 || end <= start)
-        throw new Error('Could not render PlantUML');
-      const svg = text
-        .slice(start, end + 6) // <svg ... </svg>
-        .replace('background:#FFFFFF;', ''); // transparent background
-
-      // check for syntax errors
-      if (svg.includes('>Syntax Error?</text>'))
-        throw new Error('syntax issue (it happens!). Please regenerate or change generator model.');
-
-      return svg;
-    },
+    queryFn: () => fetchPlantUmlSvg(blockCode),
     staleTime: 24 * 60 * 60 * 1000, // 1 day
   });
   renderPlantUML = renderPlantUML && (!!plantUmlHtmlData || !!plantUmlError);
@@ -192,8 +194,8 @@ function RenderCodeImpl(props: {
               </IconButton>
             </Tooltip>
           )}
-          {canCodepen && <OpenInCodepen codeBlock={{ code: blockCode, language: inferredCodeLanguage || undefined }} />}
-          {canReplit && <OpenInReplit codeBlock={{ code: blockCode, language: inferredCodeLanguage || undefined }} />}
+          {canCodepen && <ButtonCodepen codeBlock={{ code: blockCode, language: inferredCodeLanguage || undefined }} />}
+          {canReplit && <ButtonReplit codeBlock={{ code: blockCode, language: inferredCodeLanguage || undefined }} />}
           {props.noCopyButton !== true && <Tooltip title='Copy Code' variant='solid'>
             <IconButton variant='outlined' onClick={handleCopyToClipboard}>
               <ContentCopyIcon />
diff --git a/src/apps/chat/components/message/RenderCodeMermaid.tsx b/src/apps/chat/components/message/blocks/code/RenderCodeMermaid.tsx
similarity index 100%
rename from src/apps/chat/components/message/RenderCodeMermaid.tsx
rename to src/apps/chat/components/message/blocks/code/RenderCodeMermaid.tsx
diff --git a/src/apps/chat/components/message/codePrism.ts b/src/apps/chat/components/message/blocks/code/codePrism.ts
similarity index 100%
rename from src/apps/chat/components/message/codePrism.ts
rename to src/apps/chat/components/message/blocks/code/codePrism.ts
diff --git a/src/apps/link/LinkChat.tsx b/src/apps/link/LinkChat.tsx
index 8e240a794..e9dd79eb0 100644
--- a/src/apps/link/LinkChat.tsx
+++ b/src/apps/link/LinkChat.tsx
@@ -4,7 +4,7 @@ import TimeAgo from 'react-timeago';
 import { Box, Button, Card, CardContent, List, ListItem, Tooltip, Typography } from '@mui/joy';
 import TelegramIcon from '@mui/icons-material/Telegram';
 
-import { ChatMessage } from '../chat/components/message/ChatMessage';
+import { ChatMessageMemo } from '../chat/components/message/ChatMessage';
 import { ScrollToBottom } from '../chat/components/scroll-to-bottom/ScrollToBottom';
 import { useChatShowSystemMessages } from '../chat/store-app-chat';
 
@@ -136,10 +136,11 @@ export function LinkChat(props: { conversation: DConversation, storedAt: Date, e
             </ListItem>
 
             {filteredMessages.map((message, idx) =>
-              <ChatMessage
-                key={'msg-' + message.id} message={message}
-                showDate={idx === 0 || idx === filteredMessages.length - 1}
-                onMessageEdit={text => message.text = text}
+              <ChatMessageMemo
+                key={'msg-' + message.id}
+                message={message}
+                blocksShowDate={idx === 0 || idx === filteredMessages.length - 1 /* first and last message */}
+                onMessageEdit={(_messageId, text: string) => message.text = text}
               />,
             )}
 
diff --git a/src/apps/personas/creator/Creator.tsx b/src/apps/personas/creator/Creator.tsx
index f7189ea4c..646c33bc9 100644
--- a/src/apps/personas/creator/Creator.tsx
+++ b/src/apps/personas/creator/Creator.tsx
@@ -5,7 +5,7 @@ import AddIcon from '@mui/icons-material/Add';
 import ContentCopyIcon from '@mui/icons-material/ContentCopy';
 import SettingsAccessibilityIcon from '@mui/icons-material/SettingsAccessibility';
 
-import { RenderMarkdown } from '../../chat/components/message/RenderMarkdown';
+import { RenderMarkdown } from '../../chat/components/message/blocks/RenderMarkdown';
 
 import { LLMChainStep, useLLMChain } from '~/modules/aifn/useLLMChain';
 
diff --git a/src/apps/settings-modal/ShortcutsModal.tsx b/src/apps/settings-modal/ShortcutsModal.tsx
index 4c78cff8c..b4050153b 100644
--- a/src/apps/settings-modal/ShortcutsModal.tsx
+++ b/src/apps/settings-modal/ShortcutsModal.tsx
@@ -1,13 +1,12 @@
 import * as React from 'react';
 
-import { ChatMessage } from '../chat/components/message/ChatMessage';
+import { BlocksRenderer } from '../chat/components/message/blocks/BlocksRenderer';
 
 import { GoodModal } from '~/common/components/GoodModal';
-import { createDMessage } from '~/common/state/store-chats';
 import { platformAwareKeystrokes } from '~/common/components/KeyStroke';
 
 
-const shortcutsMd = `
+const shortcutsMd = platformAwareKeystrokes(`
 
 | Shortcut            | Description                                     |
 |---------------------|-------------------------------------------------|
@@ -30,15 +29,13 @@ const shortcutsMd = `
 | Ctrl + Shift + O    | Options (current Chat Model)                    |
 | Ctrl + Shift + ?    | Shortcuts                                       |
 
-`.trim();
-
-const shortcutsMessage = createDMessage('assistant', platformAwareKeystrokes(shortcutsMd));
+`).trim();
 
 
 export function ShortcutsModal(props: { onClose: () => void }) {
   return (
     <GoodModal open title='Desktop Shortcuts' onClose={props.onClose}>
-      <ChatMessage message={shortcutsMessage} hideAvatars noBottomBorder sx={{ p: 0, m: 0 }} />
+      <BlocksRenderer text={shortcutsMd} fromRole='assistant' renderTextAsMarkdown />
     </GoodModal>
   );
 }
\ No newline at end of file
diff --git a/src/modules/aifn/digrams/DiagramsModal.tsx b/src/modules/aifn/digrams/DiagramsModal.tsx
index 973287c74..2686d718a 100644
--- a/src/modules/aifn/digrams/DiagramsModal.tsx
+++ b/src/modules/aifn/digrams/DiagramsModal.tsx
@@ -8,9 +8,9 @@ import ReplayIcon from '@mui/icons-material/Replay';
 import StopOutlinedIcon from '@mui/icons-material/StopOutlined';
 import TelegramIcon from '@mui/icons-material/Telegram';
 
-import { llmStreamingChatGenerate } from '~/modules/llms/llm.client';
+import { BlocksRenderer } from '../../../apps/chat/components/message/blocks/BlocksRenderer';
 
-import { ChatMessage } from '../../../apps/chat/components/message/ChatMessage';
+import { llmStreamingChatGenerate } from '~/modules/llms/llm.client';
 
 import { GoodModal } from '~/common/components/GoodModal';
 import { InlineError } from '~/common/components/InlineError';
@@ -183,16 +183,23 @@ export function DiagramsModal(props: { config: DiagramConfig, onClose: () => voi
     </Box>}
 
     {!!message && (!abortController || showOptions) && (
-      <ChatMessage
-        message={message} hideAvatars noBottomBorder noMarkdown diagramMode
-        codeBackground='background.surface'
-        onMessageEdit={(text) => setMessage({ ...message, text })}
-        sx={{
-          backgroundColor: 'background.level2',
-          marginX: 'calc(-1 * var(--Card-padding))',
-          minHeight: 96,
-        }}
-      />
+      <Box sx={{
+        backgroundColor: 'background.level2',
+        marginX: 'calc(-1 * var(--Card-padding))',
+        minHeight: 96,
+        p: { xs: 1, md: 2 },
+        '& > div > div > code': {
+          boxShadow: 'md',
+        },
+      }}>
+        <BlocksRenderer
+          text={message.text}
+          fromRole='assistant'
+          renderTextAsMarkdown={false}
+          specialDiagramMode
+          // onMessageEdit={(text) => setMessage({ ...message, text })}
+        />
+      </Box>
     )}
 
     {!message && <Divider />}