From aa0b7ed1a82537d49e910b822be9ca32d8ed892a Mon Sep 17 00:00:00 2001
From: one <wangan.cs@gmail.com>
Date: Thu, 12 Jun 2025 16:28:28 +0800
Subject: [PATCH] feat(Markdown): customize table to support source copying
 (#7019)

* feat(Markdown): customize table to support source copying

- add a customized table component
- update ChatNavigation excluded selectors

* refactor: remove redundant feedback

* test: add tests for Table
---
 .../src/pages/home/Markdown/Markdown.tsx      |   4 +-
 .../src/pages/home/Markdown/Table.tsx         | 120 +++++++
 .../home/Markdown/__tests__/Markdown.test.tsx |  27 ++
 .../home/Markdown/__tests__/Table.test.tsx    | 316 ++++++++++++++++++
 .../__snapshots__/Markdown.test.tsx.snap      |  18 +
 .../__snapshots__/Table.test.tsx.snap         | 103 ++++++
 .../pages/home/Messages/ChatNavigation.tsx    |   3 +-
 7 files changed, 589 insertions(+), 2 deletions(-)
 create mode 100644 src/renderer/src/pages/home/Markdown/Table.tsx
 create mode 100644 src/renderer/src/pages/home/Markdown/__tests__/Table.test.tsx
 create mode 100644 src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Table.test.tsx.snap

diff --git a/src/renderer/src/pages/home/Markdown/Markdown.tsx b/src/renderer/src/pages/home/Markdown/Markdown.tsx
index d82905cf7f..8ee90a859a 100644
--- a/src/renderer/src/pages/home/Markdown/Markdown.tsx
+++ b/src/renderer/src/pages/home/Markdown/Markdown.tsx
@@ -24,6 +24,7 @@ import remarkMath from 'remark-math'
 
 import CodeBlock from './CodeBlock'
 import Link from './Link'
+import Table from './Table'
 
 const ALLOWED_ELEMENTS =
   /<(style|p|div|span|b|i|strong|em|ul|ol|li|table|tr|td|th|thead|tbody|h[1-6]|blockquote|pre|code|br|hr|svg|path|circle|rect|line|polyline|polygon|text|g|defs|title|desc|tspan|sub|sup)/i
@@ -83,6 +84,7 @@ const Markdown: FC<Props> = ({ block }) => {
       code: (props: any) => (
         <CodeBlock {...props} id={getCodeBlockId(props?.node?.position?.start)} onSave={onSaveCodeBlock} />
       ),
+      table: (props: any) => <Table {...props} blockId={block.id} />,
       img: (props: any) => <ImageViewer style={{ maxWidth: 500, maxHeight: 500 }} {...props} />,
       pre: (props: any) => <pre style={{ overflow: 'visible' }} {...props} />,
       p: (props) => {
@@ -91,7 +93,7 @@ const Markdown: FC<Props> = ({ block }) => {
         return <p {...props} />
       }
     } as Partial<Components>
-  }, [onSaveCodeBlock])
+  }, [onSaveCodeBlock, block.id])
 
   if (messageContent.includes('<style>')) {
     components.style = MarkdownShadowDOMRenderer as any
diff --git a/src/renderer/src/pages/home/Markdown/Table.tsx b/src/renderer/src/pages/home/Markdown/Table.tsx
new file mode 100644
index 0000000000..06074d55ac
--- /dev/null
+++ b/src/renderer/src/pages/home/Markdown/Table.tsx
@@ -0,0 +1,120 @@
+import store from '@renderer/store'
+import { messageBlocksSelectors } from '@renderer/store/messageBlock'
+import { Tooltip } from 'antd'
+import { Check, Copy } from 'lucide-react'
+import React, { memo, useCallback, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import styled from 'styled-components'
+
+interface Props {
+  children: React.ReactNode
+  node?: any
+  blockId?: string
+}
+
+/**
+ * 自定义 Markdown 表格组件，提供 copy 功能。
+ */
+const Table: React.FC<Props> = ({ children, node, blockId }) => {
+  const { t } = useTranslation()
+  const [copied, setCopied] = useState(false)
+
+  const handleCopyTable = useCallback(() => {
+    const tableMarkdown = extractTableMarkdown(blockId ?? '', node?.position)
+    if (!tableMarkdown) return
+
+    navigator.clipboard
+      .writeText(tableMarkdown)
+      .then(() => {
+        setCopied(true)
+        setTimeout(() => setCopied(false), 2000)
+      })
+      .catch((error) => {
+        window.message?.error({ content: `${t('message.copy.failed')}: ${error}`, key: 'copy-table-error' })
+      })
+  }, [node, blockId, t])
+
+  return (
+    <TableWrapper className="table-wrapper">
+      <table>{children}</table>
+      <ToolbarWrapper className="table-toolbar">
+        <Tooltip title={t('common.copy')} mouseEnterDelay={0.8}>
+          <ToolButton role="button" aria-label={t('common.copy')} onClick={handleCopyTable}>
+            {copied ? (
+              <Check size={14} style={{ color: 'var(--color-primary)' }} data-testid="check-icon" />
+            ) : (
+              <Copy size={14} data-testid="copy-icon" />
+            )}
+          </ToolButton>
+        </Tooltip>
+      </ToolbarWrapper>
+    </TableWrapper>
+  )
+}
+
+/**
+ * 从原始 Markdown 内容中提取表格源代码
+ * @param blockId 消息块 ID
+ * @param position 表格节点的位置信息
+ * @returns 源代码
+ */
+export function extractTableMarkdown(blockId: string, position: any): string {
+  if (!position || !blockId) return ''
+
+  const block = messageBlocksSelectors.selectById(store.getState(), blockId)
+
+  if (!block || !('content' in block) || typeof block.content !== 'string') return ''
+
+  const { start, end } = position
+  const lines = block.content.split('\n')
+
+  // 提取表格对应的行（行号从1开始，数组索引从0开始）
+  const tableLines = lines.slice(start.line - 1, end.line)
+  return tableLines.join('\n').trim()
+}
+
+const TableWrapper = styled.div`
+  position: relative;
+
+  .table-toolbar {
+    border-radius: 4px;
+    opacity: 0;
+    transition: opacity 0.2s ease;
+    transform: translateZ(0);
+    will-change: opacity;
+  }
+  &:hover {
+    .table-toolbar {
+      opacity: 1;
+    }
+  }
+`
+
+const ToolbarWrapper = styled.div`
+  position: absolute;
+  top: 8px;
+  right: 8px;
+  z-index: 10;
+`
+
+const ToolButton = styled.div`
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  width: 24px;
+  height: 24px;
+  border-radius: 4px;
+  cursor: pointer;
+  user-select: none;
+  transition: all 0.2s ease;
+  opacity: 1;
+  color: var(--color-text-3);
+  background-color: var(--color-background-mute);
+  will-change: background-color, opacity;
+
+  &:hover {
+    background-color: var(--color-background-soft);
+  }
+`
+
+export default memo(Table)
diff --git a/src/renderer/src/pages/home/Markdown/__tests__/Markdown.test.tsx b/src/renderer/src/pages/home/Markdown/__tests__/Markdown.test.tsx
index f5769eb4f8..c72f30de98 100644
--- a/src/renderer/src/pages/home/Markdown/__tests__/Markdown.test.tsx
+++ b/src/renderer/src/pages/home/Markdown/__tests__/Markdown.test.tsx
@@ -78,6 +78,18 @@ vi.mock('../Link', () => ({
   )
 }))
 
+vi.mock('../Table', () => ({
+  __esModule: true,
+  default: ({ children, blockId }: any) => (
+    <div data-testid="table-component" data-block-id={blockId}>
+      <table>{children}</table>
+      <button type="button" data-testid="copy-table-button">
+        Copy Table
+      </button>
+    </div>
+  )
+}))
+
 vi.mock('@renderer/components/MarkdownShadowDOMRenderer', () => ({
   __esModule: true,
   default: ({ children }: any) => <div data-testid="shadow-dom">{children}</div>
@@ -104,6 +116,11 @@ vi.mock('react-markdown', () => ({
           {components.code({ children: 'test code', node: { position: { start: { line: 1 } } } })}
         </div>
       )}
+      {components?.table && (
+        <div data-testid="has-table-component">
+          {components.table({ children: 'test table', node: { position: { start: { line: 1 } } } })}
+        </div>
+      )}
       {components?.img && <span data-testid="has-img-component">img</span>}
       {components?.style && <span data-testid="has-style-component">style</span>}
     </div>
@@ -300,6 +317,16 @@ describe('Markdown', () => {
       })
     })
 
+    it('should integrate Table component with copy functionality', () => {
+      const block = createMainTextBlock({ id: 'test-block-456' })
+      render(<Markdown block={block} />)
+
+      expect(screen.getByTestId('has-table-component')).toBeInTheDocument()
+
+      const tableComponent = screen.getByTestId('table-component')
+      expect(tableComponent).toHaveAttribute('data-block-id', 'test-block-456')
+    })
+
     it('should integrate ImagePreview component', () => {
       render(<Markdown block={createMainTextBlock()} />)
 
diff --git a/src/renderer/src/pages/home/Markdown/__tests__/Table.test.tsx b/src/renderer/src/pages/home/Markdown/__tests__/Table.test.tsx
new file mode 100644
index 0000000000..5a5bbeb90d
--- /dev/null
+++ b/src/renderer/src/pages/home/Markdown/__tests__/Table.test.tsx
@@ -0,0 +1,316 @@
+import { act, render, screen, waitFor } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import Table, { extractTableMarkdown } from '../Table'
+
+const mocks = vi.hoisted(() => {
+  return {
+    store: {
+      getState: vi.fn()
+    },
+    messageBlocksSelectors: {
+      selectById: vi.fn()
+    },
+    windowMessage: {
+      error: vi.fn()
+    }
+  }
+})
+
+// Mock dependencies
+vi.mock('@renderer/store', () => ({
+  __esModule: true,
+  default: mocks.store
+}))
+
+vi.mock('@renderer/store/messageBlock', () => ({
+  messageBlocksSelectors: mocks.messageBlocksSelectors
+}))
+
+vi.mock('react-i18next', () => ({
+  useTranslation: () => ({
+    t: (key: string) => key
+  })
+}))
+
+vi.mock('antd', () => ({
+  Tooltip: ({ children, title }: any) => (
+    <div data-testid="tooltip" title={title}>
+      {children}
+    </div>
+  )
+}))
+
+Object.assign(window, {
+  message: mocks.windowMessage
+})
+
+describe('Table', () => {
+  beforeAll(() => {
+    vi.stubGlobal('jest', {
+      advanceTimersByTime: vi.advanceTimersByTime.bind(vi)
+    })
+  })
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    vi.useFakeTimers()
+  })
+
+  afterEach(() => {
+    vi.restoreAllMocks()
+    vi.runOnlyPendingTimers()
+    vi.useRealTimers()
+  })
+
+  afterAll(() => {
+    vi.unstubAllGlobals()
+  })
+
+  // https://testing-library.com/docs/user-event/clipboard/
+  const user = userEvent.setup({
+    advanceTimers: vi.advanceTimersByTime.bind(vi),
+    writeToClipboard: true
+  })
+
+  // Test data factories
+  const createMockBlock = (content: string = defaultTableContent) => ({
+    id: 'test-block-1',
+    content
+  })
+
+  const createTablePosition = (startLine = 1, endLine = 3) => ({
+    start: { line: startLine },
+    end: { line: endLine }
+  })
+
+  const defaultTableContent = `| Header 1 | Header 2 |
+|----------|----------|
+| Cell 1   | Cell 2   |`
+
+  const defaultProps = {
+    children: (
+      <tbody>
+        <tr>
+          <td>Cell 1</td>
+          <td>Cell 2</td>
+        </tr>
+      </tbody>
+    ),
+    blockId: 'test-block-1',
+    node: { position: createTablePosition() }
+  }
+
+  const getCopyButton = () => screen.getByRole('button', { name: /common\.copy/i })
+  const getCopyIcon = () => screen.getByTestId('copy-icon')
+  const getCheckIcon = () => screen.getByTestId('check-icon')
+  const queryCheckIcon = () => screen.queryByTestId('check-icon')
+  const queryCopyIcon = () => screen.queryByTestId('copy-icon')
+
+  describe('rendering', () => {
+    it('should render table with children and toolbar', () => {
+      render(<Table {...defaultProps} />)
+
+      expect(screen.getByRole('table')).toBeInTheDocument()
+      expect(screen.getByText('Cell 1')).toBeInTheDocument()
+      expect(screen.getByText('Cell 2')).toBeInTheDocument()
+      expect(screen.getByTestId('tooltip')).toBeInTheDocument()
+    })
+
+    it('should render with table-wrapper and table-toolbar classes', () => {
+      const { container } = render(<Table {...defaultProps} />)
+
+      expect(container.querySelector('.table-wrapper')).toBeInTheDocument()
+      expect(container.querySelector('.table-toolbar')).toBeInTheDocument()
+    })
+
+    it('should render copy button with correct tooltip', () => {
+      render(<Table {...defaultProps} />)
+
+      const tooltip = screen.getByTestId('tooltip')
+      expect(tooltip).toHaveAttribute('title', 'common.copy')
+    })
+
+    it('should match snapshot', () => {
+      const { container } = render(<Table {...defaultProps} />)
+      expect(container.firstChild).toMatchSnapshot()
+    })
+  })
+
+  describe('extractTableMarkdown', () => {
+    beforeEach(() => {
+      mocks.store.getState.mockReturnValue({})
+    })
+
+    it('should extract table content from specified line range', () => {
+      const block = createMockBlock()
+      const position = createTablePosition(1, 3)
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(block)
+
+      const result = extractTableMarkdown('test-block-1', position)
+
+      expect(result).toBe(defaultTableContent)
+      expect(mocks.messageBlocksSelectors.selectById).toHaveBeenCalledWith({}, 'test-block-1')
+    })
+
+    it('should handle line range extraction correctly', () => {
+      const multiLineContent = `Line 0
+| Header 1 | Header 2 |
+|----------|----------|
+| Cell 1   | Cell 2   |
+Line 4`
+      const block = createMockBlock(multiLineContent)
+      const position = createTablePosition(2, 4) // Extract lines 2-4 (table part)
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(block)
+
+      const result = extractTableMarkdown('test-block-1', position)
+
+      expect(result).toBe(`| Header 1 | Header 2 |
+|----------|----------|
+| Cell 1   | Cell 2   |`)
+    })
+
+    it('should return empty string when blockId is empty', () => {
+      const result = extractTableMarkdown('', createTablePosition())
+      expect(result).toBe('')
+      expect(mocks.messageBlocksSelectors.selectById).not.toHaveBeenCalled()
+    })
+
+    it('should return empty string when position is null', () => {
+      const result = extractTableMarkdown('test-block-1', null)
+      expect(result).toBe('')
+      expect(mocks.messageBlocksSelectors.selectById).not.toHaveBeenCalled()
+    })
+
+    it('should return empty string when position is undefined', () => {
+      const result = extractTableMarkdown('test-block-1', undefined)
+      expect(result).toBe('')
+      expect(mocks.messageBlocksSelectors.selectById).not.toHaveBeenCalled()
+    })
+
+    it('should return empty string when block does not exist', () => {
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(null)
+
+      const result = extractTableMarkdown('non-existent-block', createTablePosition())
+
+      expect(result).toBe('')
+    })
+
+    it('should return empty string when block has no content property', () => {
+      const blockWithoutContent = { id: 'test-block-1' }
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(blockWithoutContent)
+
+      const result = extractTableMarkdown('test-block-1', createTablePosition())
+
+      expect(result).toBe('')
+    })
+
+    it('should return empty string when block content is not a string', () => {
+      const blockWithInvalidContent = { id: 'test-block-1', content: 123 }
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(blockWithInvalidContent)
+
+      const result = extractTableMarkdown('test-block-1', createTablePosition())
+
+      expect(result).toBe('')
+    })
+
+    it('should handle boundary line numbers correctly', () => {
+      const block = createMockBlock('Line 1\nLine 2\nLine 3')
+      const position = createTablePosition(1, 3)
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(block)
+
+      const result = extractTableMarkdown('test-block-1', position)
+
+      expect(result).toBe('Line 1\nLine 2\nLine 3')
+    })
+  })
+
+  describe('copy functionality', () => {
+    beforeEach(() => {
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(createMockBlock())
+    })
+
+    it('should copy table content to clipboard on button click', async () => {
+      render(<Table {...defaultProps} />)
+
+      const copyButton = getCopyButton()
+      await user.click(copyButton)
+
+      await waitFor(() => {
+        expect(getCheckIcon()).toBeInTheDocument()
+        expect(queryCopyIcon()).not.toBeInTheDocument()
+      })
+    })
+
+    it('should show check icon after successful copy', async () => {
+      render(<Table {...defaultProps} />)
+
+      // Initially shows copy icon
+      expect(getCopyIcon()).toBeInTheDocument()
+
+      const copyButton = getCopyButton()
+      await user.click(copyButton)
+
+      await waitFor(() => {
+        expect(getCheckIcon()).toBeInTheDocument()
+        expect(queryCopyIcon()).not.toBeInTheDocument()
+      })
+    })
+
+    it('should reset to copy icon after 2 seconds', async () => {
+      render(<Table {...defaultProps} />)
+
+      const copyButton = getCopyButton()
+      await user.click(copyButton)
+
+      await waitFor(() => {
+        expect(getCheckIcon()).toBeInTheDocument()
+      })
+
+      // Fast forward 2 seconds
+      act(() => {
+        vi.advanceTimersByTime(2000)
+      })
+
+      await waitFor(() => {
+        expect(getCopyIcon()).toBeInTheDocument()
+        expect(queryCheckIcon()).not.toBeInTheDocument()
+      })
+    })
+
+    it('should not copy when extractTableMarkdown returns empty string', async () => {
+      mocks.messageBlocksSelectors.selectById.mockReturnValue(null)
+
+      render(<Table {...defaultProps} />)
+
+      const copyButton = getCopyButton()
+      await user.click(copyButton)
+
+      await waitFor(() => {
+        expect(getCopyIcon()).toBeInTheDocument()
+        expect(queryCheckIcon()).not.toBeInTheDocument()
+      })
+    })
+  })
+
+  describe('edge cases', () => {
+    it('should work without blockId', () => {
+      const propsWithoutBlockId = { ...defaultProps, blockId: undefined }
+
+      expect(() => render(<Table {...propsWithoutBlockId} />)).not.toThrow()
+
+      const copyButton = getCopyButton()
+      expect(copyButton).toBeInTheDocument()
+    })
+
+    it('should work without node position', () => {
+      const propsWithoutPosition = { ...defaultProps, node: undefined }
+
+      expect(() => render(<Table {...propsWithoutPosition} />)).not.toThrow()
+
+      const copyButton = getCopyButton()
+      expect(copyButton).toBeInTheDocument()
+    })
+  })
+})
diff --git a/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Markdown.test.tsx.snap b/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Markdown.test.tsx.snap
index e055c83f52..29aae68dc0 100644
--- a/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Markdown.test.tsx.snap
+++ b/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Markdown.test.tsx.snap
@@ -30,6 +30,24 @@ This is **bold** text.
       </button>
     </div>
   </div>
+  <div
+    data-testid="has-table-component"
+  >
+    <div
+      data-block-id="test-block-1"
+      data-testid="table-component"
+    >
+      <table>
+        test table
+      </table>
+      <button
+        data-testid="copy-table-button"
+        type="button"
+      >
+        Copy Table
+      </button>
+    </div>
+  </div>
   <span
     data-testid="has-img-component"
   >
diff --git a/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Table.test.tsx.snap b/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Table.test.tsx.snap
new file mode 100644
index 0000000000..da85d514be
--- /dev/null
+++ b/src/renderer/src/pages/home/Markdown/__tests__/__snapshots__/Table.test.tsx.snap
@@ -0,0 +1,103 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`Table > rendering > should match snapshot 1`] = `
+.c0 {
+  position: relative;
+}
+
+.c0 .table-toolbar {
+  border-radius: 4px;
+  opacity: 0;
+  transition: opacity 0.2s ease;
+  transform: translateZ(0);
+  will-change: opacity;
+}
+
+.c0:hover .table-toolbar {
+  opacity: 1;
+}
+
+.c1 {
+  position: absolute;
+  top: 8px;
+  right: 8px;
+  z-index: 10;
+}
+
+.c2 {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  width: 24px;
+  height: 24px;
+  border-radius: 4px;
+  cursor: pointer;
+  user-select: none;
+  transition: all 0.2s ease;
+  opacity: 1;
+  color: var(--color-text-3);
+  background-color: var(--color-background-mute);
+  will-change: background-color,opacity;
+}
+
+.c2:hover {
+  background-color: var(--color-background-soft);
+}
+
+<div
+  class="c0 table-wrapper"
+>
+  <table>
+    <tbody>
+      <tr>
+        <td>
+          Cell 1
+        </td>
+        <td>
+          Cell 2
+        </td>
+      </tr>
+    </tbody>
+  </table>
+  <div
+    class="c1 table-toolbar"
+  >
+    <div
+      data-testid="tooltip"
+      title="common.copy"
+    >
+      <div
+        aria-label="common.copy"
+        class="c2"
+        role="button"
+      >
+        <svg
+          class="lucide lucide-copy"
+          data-testid="copy-icon"
+          fill="none"
+          height="14"
+          stroke="currentColor"
+          stroke-linecap="round"
+          stroke-linejoin="round"
+          stroke-width="2"
+          viewBox="0 0 24 24"
+          width="14"
+          xmlns="http://www.w3.org/2000/svg"
+        >
+          <rect
+            height="14"
+            rx="2"
+            ry="2"
+            width="14"
+            x="8"
+            y="8"
+          />
+          <path
+            d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"
+          />
+        </svg>
+      </div>
+    </div>
+  </div>
+</div>
+`;
diff --git a/src/renderer/src/pages/home/Messages/ChatNavigation.tsx b/src/renderer/src/pages/home/Messages/ChatNavigation.tsx
index ee3ae1fffa..d38c05618e 100644
--- a/src/renderer/src/pages/home/Messages/ChatNavigation.tsx
+++ b/src/renderer/src/pages/home/Messages/ChatNavigation.tsx
@@ -24,7 +24,8 @@ const EXCLUDED_SELECTORS = [
   '.ant-collapse-header',
   '.group-menu-bar',
   '.code-block',
-  '.message-editor'
+  '.message-editor',
+  '.table-wrapper'
 ]
 
 // Gap between the navigation bar and the right element