Skip to content

Commit 2bc837f

Browse files
TreeTreeDiTreeTreeDi
committed
feat: Refactor code structure for improved maintainability & add useSmoothMarkdown to improve performance
1 parent da5a357 commit 2bc837f

15 files changed

Lines changed: 3036 additions & 39 deletions

File tree

‎packages/streaming-markdown/package.json‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@
1818
"scripts": {
1919
"build": "tsup src/index.ts --format cjs,esm --dts",
2020
"dev": "tsup src/index.ts --format cjs,esm --dts --watch",
21-
"lint": "tsc --noEmit"
21+
"lint": "tsc --noEmit",
22+
"test": "vitest run",
23+
"test:watch": "vitest",
24+
"test:coverage": "vitest run --coverage"
2225
},
2326
"keywords": [
2427
"react",
@@ -32,13 +35,24 @@
3235
"react-dom": "^18.0.0 || ^19.0.0"
3336
},
3437
"dependencies": {
38+
"nanoid": "^5.0.0",
3539
"react-markdown": "^9.0.1",
36-
"remark-gfm": "^4.0.0"
40+
"remark-gfm": "^4.0.0",
41+
"remark-parse": "^11.0.0",
42+
"unified": "^11.0.0",
43+
"unist-util-visit": "^5.0.0"
3744
},
3845
"devDependencies": {
46+
"@testing-library/react": "^14.3.1",
3947
"@types/react": "^18.2.0",
4048
"@types/react-dom": "^18.2.0",
49+
"@vitejs/plugin-react": "^5.1.0",
50+
"@vitest/coverage-v8": "^1.6.1",
51+
"@vitest/ui": "^1.6.1",
52+
"happy-dom": "^20.0.10",
53+
"jsdom": "^24.1.3",
4154
"tsup": "^8.0.0",
42-
"typescript": "^5.3.0"
55+
"typescript": "^5.3.0",
56+
"vitest": "^1.6.1"
4357
}
4458
}

‎packages/streaming-markdown/src/__tests__/useSmoothStream.test.tsx‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
import { renderHook, act } from '@testing-library/react';
2+
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
3+
import { useSmoothStream } from '../hooks/useSmoothStream';
4+
5+
describe('useSmoothStream', () => {
6+
beforeEach(() => {
7+
vi.stubGlobal('requestAnimationFrame', (cb: FrameRequestCallback) => {
8+
setTimeout(() => cb(performance.now()), 0);
9+
return 1;
10+
});
11+
vi.stubGlobal('cancelAnimationFrame', vi.fn());
12+
});
13+
14+
afterEach(() => {
15+
vi.unstubAllGlobals();
16+
vi.restoreAllMocks();
17+
});
18+
19+
it('应该初始化为空文本', () => {
20+
const onUpdate = vi.fn();
21+
const { result } = renderHook(() =>
22+
useSmoothStream({
23+
onUpdate,
24+
streamDone: false,
25+
})
26+
);
27+
28+
expect(result.current).toHaveProperty('addChunk');
29+
expect(result.current).toHaveProperty('reset');
30+
});
31+
32+
it('应该初始化为指定文本', () => {
33+
const onUpdate = vi.fn();
34+
renderHook(() =>
35+
useSmoothStream({
36+
onUpdate,
37+
streamDone: false,
38+
initialText: 'Hello',
39+
})
40+
);
41+
42+
expect(onUpdate).toHaveBeenCalledWith('Hello');
43+
});
44+
45+
it('应该通过 addChunk 添加文本块', async () => {
46+
const onUpdate = vi.fn();
47+
const { result } = renderHook(() =>
48+
useSmoothStream({
49+
onUpdate,
50+
streamDone: false,
51+
})
52+
);
53+
54+
await act(async () => {
55+
result.current.addChunk('Hello');
56+
await new Promise((resolve) => setTimeout(resolve, 50));
57+
});
58+
59+
expect(onUpdate).toHaveBeenCalled();
60+
expect(onUpdate.mock.calls.some((call) => call[0].includes('Hello'))).toBe(true);
61+
});
62+
63+
it('应该通过 reset 立即替换文本', async () => {
64+
const onUpdate = vi.fn();
65+
const { result } = renderHook(() =>
66+
useSmoothStream({
67+
onUpdate,
68+
streamDone: false,
69+
})
70+
);
71+
72+
await act(async () => {
73+
result.current.addChunk('Old text');
74+
await new Promise((resolve) => setTimeout(resolve, 10));
75+
});
76+
77+
onUpdate.mockClear();
78+
79+
act(() => {
80+
result.current.reset('New text');
81+
});
82+
83+
expect(onUpdate).toHaveBeenCalledWith('New text');
84+
});
85+
86+
it('应该忽略空字符串', () => {
87+
const onUpdate = vi.fn();
88+
const { result } = renderHook(() =>
89+
useSmoothStream({
90+
onUpdate,
91+
streamDone: false,
92+
})
93+
);
94+
95+
const callCount = onUpdate.mock.calls.length;
96+
97+
act(() => {
98+
result.current.addChunk('');
99+
});
100+
101+
expect(onUpdate.mock.calls.length).toBe(callCount);
102+
});
103+
104+
it('应该在 streamDone 时调用 onComplete', async () => {
105+
const onUpdate = vi.fn();
106+
const onComplete = vi.fn();
107+
const { rerender } = renderHook(
108+
({ streamDone }) =>
109+
useSmoothStream({
110+
onUpdate,
111+
streamDone,
112+
onComplete,
113+
initialText: '',
114+
}),
115+
{ initialProps: { streamDone: false } }
116+
);
117+
118+
await act(async () => {
119+
await new Promise((resolve) => setTimeout(resolve, 50));
120+
});
121+
122+
await act(async () => {
123+
rerender({ streamDone: true });
124+
await new Promise((resolve) => setTimeout(resolve, 100));
125+
});
126+
127+
expect(onComplete).toHaveBeenCalled();
128+
});
129+
130+
it('应该正确处理多语言字符（emoji、中文）', async () => {
131+
const onUpdate = vi.fn();
132+
const { result } = renderHook(() =>
133+
useSmoothStream({
134+
onUpdate,
135+
streamDone: false,
136+
})
137+
);
138+
139+
await act(async () => {
140+
result.current.addChunk('你好👋世界');
141+
await new Promise((resolve) => setTimeout(resolve, 100));
142+
});
143+
144+
expect(onUpdate).toHaveBeenCalled();
145+
const lastCall = onUpdate.mock.calls[onUpdate.mock.calls.length - 1][0];
146+
expect(lastCall).toContain('你好');
147+
});
148+
149+
it('reset 应该清空队列并取消 RAF', () => {
150+
const onUpdate = vi.fn();
151+
const { result } = renderHook(() =>
152+
useSmoothStream({
153+
onUpdate,
154+
streamDone: false,
155+
})
156+
);
157+
158+
act(() => {
159+
result.current.addChunk('Queue this');
160+
result.current.reset('Reset immediately');
161+
});
162+
163+
expect(onUpdate).toHaveBeenCalledWith('Reset immediately');
164+
});
165+
166+
it('应该在组件卸载时清理 RAF', async () => {
167+
const onUpdate = vi.fn();
168+
const cancelSpy = vi.fn();
169+
vi.stubGlobal('cancelAnimationFrame', cancelSpy);
170+
171+
const { result, unmount } = renderHook(() =>
172+
useSmoothStream({
173+
onUpdate,
174+
streamDone: false,
175+
})
176+
);
177+
178+
await act(async () => {
179+
result.current.addChunk('Test');
180+
await new Promise((resolve) => setTimeout(resolve, 10));
181+
});
182+
183+
unmount();
184+
185+
expect(cancelSpy).toHaveBeenCalled();
186+
});
187+
});

‎packages/streaming-markdown/src/components/Markdown/CodeBlock.tsx‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import type { ReactNode } from 'react';
2+
3+
export interface CodeBlockProps {
4+
code: string;
5+
language?: string;
6+
className?: string;
7+
}
8+
9+
export function CodeBlock({ code, language, className }: CodeBlockProps): ReactNode {
10+
return (
11+
<div className={className}>
12+
<pre>
13+
<code data-language={language}>{code}</code>
14+
</pre>
15+
</div>
16+
);
17+
}

‎packages/streaming-markdown/src/components/Markdown/StreamingMarkdown.tsx‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,65 @@
11
import type { ReactNode } from 'react';
2-
import { useMemo } from 'react';
2+
import { useEffect, useMemo, useRef, useState } from 'react';
33
import ReactMarkdown from 'react-markdown';
44
import remarkGfm from 'remark-gfm';
55
import type { Components } from 'react-markdown';
6+
import { useSmoothStream } from '../../hooks/useSmoothStream';
7+
8+
export type StreamingStatus = 'idle' | 'streaming' | 'success' | 'error';
69

710
export interface StreamingMarkdownProps {
811
children?: ReactNode;
912
className?: string;
1013
components?: Partial<Components>;
14+
status?: StreamingStatus;
1115
onComplete?: () => void;
16+
minDelay?: number;
1217
}
1318

1419
export function StreamingMarkdown({
1520
children,
1621
className,
1722
components,
23+
status = 'idle',
24+
onComplete,
25+
minDelay = 10,
1826
}: StreamingMarkdownProps): ReactNode {
19-
const processedContent = useMemo(() => {
20-
const content = typeof children === 'string' ? children : String(children || '');
27+
const [displayedText, setDisplayedText] = useState('');
28+
const previousChildrenRef = useRef<string>('');
29+
30+
const { addChunk, reset } = useSmoothStream({
31+
onUpdate: setDisplayedText,
32+
streamDone: status !== 'streaming',
33+
minDelay,
34+
initialText: '',
35+
onComplete,
36+
});
37+
38+
useEffect(() => {
39+
const currentContent = typeof children === 'string' ? children : String(children || '');
40+
const previousContent = previousChildrenRef.current;
2141

22-
console.log('Original content:123123', JSON.stringify(content));
42+
if (currentContent !== previousContent) {
43+
if (currentContent.startsWith(previousContent)) {
44+
const delta = currentContent.slice(previousContent.length);
45+
addChunk(delta);
46+
} else {
47+
reset(currentContent);
48+
}
49+
previousChildrenRef.current = currentContent;
50+
}
51+
}, [children, addChunk, reset]);
52+
53+
const processedContent = useMemo(() => {
54+
const trimmed = displayedText.trim();
2355

24-
const trimmed = content.trim();
25-
56+
console.log('Displayed content:456456', JSON.stringify(displayedText));
2657
if (trimmed.endsWith('```') && !trimmed.endsWith('```\n')) {
2758
return `${trimmed}\n`;
2859
}
29-
60+
3061
return trimmed;
31-
}, [children]);
62+
}, [displayedText]);
3263

3364
return (
3465
<ReactMarkdown

‎packages/streaming-markdown/src/components/Markdown/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export { StreamingMarkdown } from './StreamingMarkdown';
2+
export type { StreamingMarkdownProps, StreamingStatus } from './StreamingMarkdown';

‎packages/streaming-markdown/src/components/Message/MessageBlockRenderer.tsx‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import type { ReactNode } from 'react';
2+
import type { MessageBlock } from '../../types/message';
3+
import { MessageBlockType, MessageBlockStatus } from '../../types/message';
4+
import { StreamingMarkdown } from '../Markdown/StreamingMarkdown';
5+
6+
export interface MessageBlockRendererProps {
7+
block: MessageBlock;
8+
className?: string;
9+
}
10+
11+
export function MessageBlockRenderer({ block, className }: MessageBlockRendererProps): ReactNode {
12+
switch (block.type) {
13+
case MessageBlockType.MAIN_TEXT:
14+
case MessageBlockType.TRANSLATION:
15+
case MessageBlockType.THINKING:
16+
case MessageBlockType.ERROR:
17+
case MessageBlockType.UNKNOWN: {
18+
const textBlock = block as Extract<MessageBlock, { content: string }>;
19+
return (
20+
<StreamingMarkdown
21+
className={className}
22+
status={block.status === MessageBlockStatus.STREAMING ? 'streaming' : 'success'}
23+
>
24+
{textBlock.content}
25+
</StreamingMarkdown>
26+
);
27+
}
28+
29+
case MessageBlockType.CODE: {
30+
const codeBlock = block as Extract<MessageBlock, { content: string }>;
31+
return (
32+
<div className={className}>
33+
<pre>
34+
<code>{codeBlock.content}</code>
35+
</pre>
36+
</div>
37+
);
38+
}
39+
40+
case MessageBlockType.IMAGE:
41+
case MessageBlockType.VIDEO:
42+
case MessageBlockType.FILE: {
43+
const mediaBlock = block as Extract<MessageBlock, { url: string }>;
44+
return (
45+
<div className={className}>
46+
<a href={mediaBlock.url} target="_blank" rel="noopener noreferrer">
47+
{mediaBlock.name ?? 'Media File'}
48+
</a>
49+
</div>
50+
);
51+
}
52+
53+
case MessageBlockType.TOOL:
54+
case MessageBlockType.CITATION: {
55+
const toolBlock = block as Extract<MessageBlock, { payload?: Record<string, unknown> }>;
56+
return (
57+
<div className={className}>
58+
<pre>{JSON.stringify(toolBlock.payload ?? {}, null, 2)}</pre>
59+
</div>
60+
);
61+
}
62+
63+
default:
64+
return null;
65+
}
66+
}

0 commit comments

Comments
 (0)