Merge pull request #71 from ai-action/fix/messages

Author: remarkablemark

src/components/CodeBlock/CodeBlock.tsx

@@ -11,16 +11,30 @@ interface CodeBlockProps {
 
 const highlightCache = new Map<string, string>();
 
-export const CODE_BLOCK_REGEX = /^(`{3,})(\w+)?[ \t]*\n([\s\S]*?)^\1[ \t]*$/gm;
+const CODE_BLOCK_REGEX =
+  /^(?<indent>[ \t]*)(`{3,})(\w+)?[ \t]*\n([\s\S]*?)^\k<indent>\2[ \t]*$/gm;
+
+export function normalizeCodeBlockContent(
+  content: string,
+  indent = '',
+): string {
+  if (!indent) {
+    return content.trim();
+  }
+
+  const indentPattern = new RegExp(`^${indent}`, 'gm');
+  return content.replace(indentPattern, '').trim();
+}
 
 export async function prewarmCodeBlocks(content: string): Promise<void> {
   const promises: Promise<void>[] = [];
   let match;
   CODE_BLOCK_REGEX.lastIndex = 0;
 
   while ((match = CODE_BLOCK_REGEX.exec(content)) !== null) {
-    const language = match[2];
-    const code = match[3].trim();
+    const indent = match[1];
+    const language = match[3];
+    const code = normalizeCodeBlockContent(match[4], indent);
     // v8 ignore next 2
     if (code) {
       promises.push(prewarmHighlight(code, language));

src/components/Markdown/Markdown.test.tsx

@@ -1,9 +1,38 @@
+import { useStdout } from 'ink';
 import { render } from 'ink-testing-library';
 
 import { UI } from '../../constants';
 import { Markdown } from './Markdown';
 
+const { mockColumns } = vi.hoisted(() => ({
+  mockColumns: {
+    value: 100,
+  },
+}));
+
+vi.mock('ink', async () => ({
+  ...(await vi.importActual('ink')),
+  useStdout: vi.fn(() => ({
+    stdout: {
+      columns: mockColumns.value,
+    },
+  })),
+}));
+
+function setTerminalWidth(columns: number) {
+  mockColumns.value = columns;
+}
+
+function stripAnsi(value: string | undefined) {
+  return value?.replaceAll(new RegExp(String.raw`\u001B\[[0-9;]*m`, 'g'), '');
+}
+
 describe('Markdown', () => {
+  beforeEach(() => {
+    setTerminalWidth(100);
+    vi.mocked(useStdout).mockClear();
+  });
+
   it('renders markdown content', () => {
     const { lastFrame } = render(<Markdown content="# Hello" />);
     expect(lastFrame()).toContain('Hello');
@@ -79,4 +108,18 @@ describe('Markdown', () => {
     const { lastFrame } = render(<Markdown content="$dx \\, dt$" />);
     expect(lastFrame()).not.toContain('\\,');
   });
+
+  it('reflows wrapped markdown lists before Ink wraps ANSI output', () => {
+    setTerminalWidth(40);
+
+    const content =
+      '4. **Restructure the "Usage" section** to clearly separate **Interactive TUI** from **CLI Commands**.';
+
+    const { lastFrame } = render(<Markdown content={content} />);
+    const frame = stripAnsi(lastFrame()) ?? '';
+
+    expect(frame).toContain('CLI');
+    expect(frame).toContain('Commands');
+    expect(frame).not.toContain('**');
+  });
 });

src/components/Markdown/Markdown.tsx

@@ -1,5 +1,5 @@
 import { Text, useStdout } from 'ink';
-import { marked } from 'marked';
+import { Marked, type Token } from 'marked';
 import { markedTerminal } from 'marked-terminal';
 import { memo, useMemo } from 'react';
 
@@ -14,24 +14,41 @@ interface MarkdownProps {
 
 const HR_PLACEHOLDER = '__CODE_OLLAMA_HR_PLACEHOLDER__';
 
-marked.use(
-  markedTerminal({
-    theme: 'gitHub',
-  }),
-);
-
-marked.use({
-  extensions: [inlineMathExtension],
-  renderer: {
-    hr: () => `${HR_PLACEHOLDER}\n`,
-  },
-});
-
 function renderMarkdown(content: string, hrWidth: number): string {
   const hr = UI.MARKDOWN_HR_CHARACTER.repeat(Math.max(1, hrWidth));
+  const markdown = new Marked();
+  const rendererExtension = {
+    extensions: [inlineMathExtension],
+    useNewRenderer: true,
+    renderer: {
+      hr: () => `${HR_PLACEHOLDER}\n`,
+      text(token: Token) {
+        const textToken = token as Token & {
+          text?: string;
+          tokens?: Token[];
+        };
+
+        if (typeof token === 'object' && Array.isArray(textToken.tokens)) {
+          return this.parser.parseInline(textToken.tokens);
+        }
+
+        return String(textToken.text);
+      },
+    },
+  } as Parameters<Marked['use']>[0];
+
+  markdown.use(
+    markedTerminal({
+      theme: 'gitHub',
+      reflowText: true,
+      width: Math.max(1, hrWidth),
+    }),
+  );
+
+  markdown.use(rendererExtension);
 
   try {
-    const result = marked.parse(content);
+    const result = markdown.parse(content);
     // v8 ignore start
     const text = typeof result === 'string' ? result.trim() : content;
     return text.replaceAll(HR_PLACEHOLDER, hr);