perf(chat): debounce persisted writes and re-render only the streaming message
- the Zustand persist middleware serialized every saved session and wrote to AsyncStorage on each streamed token (appendToMessage -> set); writes are now debounced (400 ms) and flushed when the app leaves the foreground - every token re-rendered every MessageBubble, re-parsing the Markdown of the whole conversation; rows are memoized so only the changed message renders (covered by MessageList.test.tsx)
This commit is contained in:
1 parent
200fb6afd7
commit
0e5b38fd15
6 files changed
+248
-9
No files matched your search
@@ -0,0 +1,78 @@
|
|||||||
|
import { render, screen } from '@testing-library/react-native';
|
||||||
|
import React from 'react';
|
||||||
|
import { MessageList } from '../components/chat/MessageList';
|
||||||
|
import { ChatMessage } from '../store/useAppStore';
|
||||||
|
|
||||||
|
jest.mock('react-native-reanimated', () => {
|
||||||
|
const { View } = require('react-native');
|
||||||
|
const chain: any = new Proxy(() => chain, { get: () => chain });
|
||||||
|
return { __esModule: true, default: { View }, FadeInDown: chain };
|
||||||
|
});
|
||||||
|
|
||||||
|
const mockRenders: Record<string, number> = {};
|
||||||
|
jest.mock('../components/chat/MessageBubble', () => {
|
||||||
|
const { Text } = require('react-native');
|
||||||
|
return {
|
||||||
|
MessageBubble: ({ message }: any) => {
|
||||||
|
mockRenders[message.id] = (mockRenders[message.id] ?? 0) + 1;
|
||||||
|
return <Text>{`${message.id}:${message.content}`}</Text>;
|
||||||
|
},
|
||||||
|
};
|
||||||
|
});
|
||||||
|
jest.mock('../components/chat/TypingIndicator', () => {
|
||||||
|
const { Text } = require('react-native');
|
||||||
|
return { TypingIndicator: () => <Text>typing-indicator</Text> };
|
||||||
|
});
|
||||||
|
|
||||||
|
const msg = (id: string, role: ChatMessage['role'], content: string): ChatMessage => ({
|
||||||
|
id,
|
||||||
|
role,
|
||||||
|
content,
|
||||||
|
timestamp: 0,
|
||||||
|
});
|
||||||
|
|
||||||
|
const handlers = { onEdit: jest.fn(), onDelete: jest.fn(), onRegenerate: jest.fn() };
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
Object.keys(mockRenders).forEach((k) => delete mockRenders[k]);
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('MessageList', () => {
|
||||||
|
it('renders every message', () => {
|
||||||
|
render(
|
||||||
|
<MessageList
|
||||||
|
messages={[msg('u1', 'user', 'q'), msg('a1', 'assistant', 'a')]}
|
||||||
|
isStreaming={false}
|
||||||
|
{...handlers}
|
||||||
|
/>,
|
||||||
|
);
|
||||||
|
expect(screen.getByText('u1:q')).toBeTruthy();
|
||||||
|
expect(screen.getByText('a1:a')).toBeTruthy();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('re-renders only the message that changed while an answer streams in', () => {
|
||||||
|
const user = msg('u1', 'user', 'question');
|
||||||
|
const earlier = msg('a0', 'assistant', 'earlier answer');
|
||||||
|
const streaming = msg('a1', 'assistant', 'He');
|
||||||
|
|
||||||
|
const { rerender } = render(
|
||||||
|
<MessageList messages={[user, earlier, streaming]} isStreaming {...handlers} />,
|
||||||
|
);
|
||||||
|
expect(mockRenders).toEqual({ u1: 1, a0: 1, a1: 1 });
|
||||||
|
|
||||||
|
rerender(
|
||||||
|
<MessageList messages={[user, earlier, { ...streaming, content: 'Hello' }]} isStreaming {...handlers} />,
|
||||||
|
);
|
||||||
|
expect(screen.getByText('a1:Hello')).toBeTruthy();
|
||||||
|
expect(mockRenders.u1).toBe(1);
|
||||||
|
expect(mockRenders.a0).toBe(1);
|
||||||
|
expect(mockRenders.a1).toBe(2);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('shows the typing indicator while the last assistant message is still empty', () => {
|
||||||
|
render(
|
||||||
|
<MessageList messages={[msg('u1', 'user', 'q'), msg('a1', 'assistant', '')]} isStreaming {...handlers} />,
|
||||||
|
);
|
||||||
|
expect(screen.getByText('typing-indicator')).toBeTruthy();
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -0,0 +1,72 @@
|
|||||||
|
import { createDebouncedStorage } from '../store/debouncedStorage';
|
||||||
|
|
||||||
|
const makeBase = () => ({
|
||||||
|
getItem: jest.fn(async (_name: string): Promise<string | null> => 'from-disk'),
|
||||||
|
setItem: jest.fn(async (_name: string, _value: string) => {}),
|
||||||
|
removeItem: jest.fn(async (_name: string) => {}),
|
||||||
|
});
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
jest.useFakeTimers();
|
||||||
|
});
|
||||||
|
afterEach(() => {
|
||||||
|
jest.useRealTimers();
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('createDebouncedStorage', () => {
|
||||||
|
it('collapses a burst of writes into one trailing write with the latest value', () => {
|
||||||
|
const base = makeBase();
|
||||||
|
const storage = createDebouncedStorage(base, 400);
|
||||||
|
|
||||||
|
storage.setItem('k', '1');
|
||||||
|
jest.advanceTimersByTime(300);
|
||||||
|
storage.setItem('k', '2');
|
||||||
|
jest.advanceTimersByTime(300);
|
||||||
|
storage.setItem('k', '3');
|
||||||
|
expect(base.setItem).not.toHaveBeenCalled();
|
||||||
|
|
||||||
|
jest.advanceTimersByTime(400);
|
||||||
|
expect(base.setItem).toHaveBeenCalledTimes(1);
|
||||||
|
expect(base.setItem).toHaveBeenCalledWith('k', '3');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('reads a pending value before it reaches the disk, and falls back to the base storage', async () => {
|
||||||
|
const base = makeBase();
|
||||||
|
const storage = createDebouncedStorage(base);
|
||||||
|
|
||||||
|
await expect(storage.getItem('k')).resolves.toBe('from-disk');
|
||||||
|
storage.setItem('k', 'fresh');
|
||||||
|
await expect(storage.getItem('k')).resolves.toBe('fresh');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('flush writes immediately and cancels the timer', async () => {
|
||||||
|
const base = makeBase();
|
||||||
|
const storage = createDebouncedStorage(base);
|
||||||
|
|
||||||
|
storage.setItem('a', '1');
|
||||||
|
storage.setItem('b', '2');
|
||||||
|
await storage.flush();
|
||||||
|
expect(base.setItem.mock.calls).toEqual([['a', '1'], ['b', '2']]);
|
||||||
|
|
||||||
|
jest.advanceTimersByTime(1000);
|
||||||
|
expect(base.setItem).toHaveBeenCalledTimes(2);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('flush with nothing pending does not write', async () => {
|
||||||
|
const base = makeBase();
|
||||||
|
await createDebouncedStorage(base).flush();
|
||||||
|
expect(base.setItem).not.toHaveBeenCalled();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('removeItem drops a pending write and removes from the base storage', async () => {
|
||||||
|
const base = makeBase();
|
||||||
|
const storage = createDebouncedStorage(base);
|
||||||
|
|
||||||
|
storage.setItem('k', '1');
|
||||||
|
await storage.removeItem('k');
|
||||||
|
jest.advanceTimersByTime(1000);
|
||||||
|
|
||||||
|
expect(base.setItem).not.toHaveBeenCalled();
|
||||||
|
expect(base.removeItem).toHaveBeenCalledWith('k');
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -1,4 +1,5 @@
|
|||||||
import AsyncStorage from '@react-native-async-storage/async-storage';
|
import AsyncStorage from '@react-native-async-storage/async-storage';
|
||||||
|
import { appStorage } from '../store/debouncedStorage';
|
||||||
import { useAppStore } from '../store/useAppStore';
|
import { useAppStore } from '../store/useAppStore';
|
||||||
|
|
||||||
const state = () => useAppStore.getState();
|
const state = () => useAppStore.getState();
|
||||||
@@ -172,7 +173,7 @@ describe('persistence', () => {
|
|||||||
});
|
});
|
||||||
state().addMessage({ role: 'user', content: 'persist me' });
|
state().addMessage({ role: 'user', content: 'persist me' });
|
||||||
|
|
||||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
await appStorage.flush();
|
||||||
const raw = await AsyncStorage.getItem('localaiapp-storage');
|
const raw = await AsyncStorage.getItem('localaiapp-storage');
|
||||||
const persisted = JSON.parse(raw as string).state;
|
const persisted = JSON.parse(raw as string).state;
|
||||||
|
|
||||||
|
|||||||
@@ -14,6 +14,41 @@ interface MessageListProps {
|
|||||||
onRegenerate: () => void;
|
onRegenerate: () => void;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
interface MessageRowProps {
|
||||||
|
message: ChatMessage;
|
||||||
|
delayed: boolean;
|
||||||
|
isStreaming: boolean;
|
||||||
|
isLastAssistant: boolean;
|
||||||
|
onEdit: (id: string) => void;
|
||||||
|
onDelete: (id: string) => void;
|
||||||
|
onRegenerate: () => void;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Memoized so that a streamed token only re-renders (and re-parses the Markdown of)
|
||||||
|
// the message that changed, not the whole conversation.
|
||||||
|
const MessageRow = React.memo(function MessageRow({
|
||||||
|
message,
|
||||||
|
delayed,
|
||||||
|
isStreaming,
|
||||||
|
isLastAssistant,
|
||||||
|
onEdit,
|
||||||
|
onDelete,
|
||||||
|
onRegenerate,
|
||||||
|
}: MessageRowProps) {
|
||||||
|
return (
|
||||||
|
<Animated.View entering={FadeInDown.duration(300).delay(delayed ? 100 : 0)}>
|
||||||
|
<MessageBubble
|
||||||
|
message={message}
|
||||||
|
onEdit={() => onEdit(message.id)}
|
||||||
|
onRegenerate={onRegenerate}
|
||||||
|
onDelete={() => onDelete(message.id)}
|
||||||
|
isStreaming={isStreaming}
|
||||||
|
isLastAssistant={isLastAssistant}
|
||||||
|
/>
|
||||||
|
</Animated.View>
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
export function MessageList({ messages, isStreaming, onEdit, onDelete, onRegenerate }: MessageListProps) {
|
export function MessageList({ messages, isStreaming, onEdit, onDelete, onRegenerate }: MessageListProps) {
|
||||||
const flatListRef = useRef<FlatList>(null);
|
const flatListRef = useRef<FlatList>(null);
|
||||||
const lastAssistantId = findLastAssistantId(messages);
|
const lastAssistantId = findLastAssistantId(messages);
|
||||||
@@ -34,16 +69,15 @@ export function MessageList({ messages, isStreaming, onEdit, onDelete, onRegener
|
|||||||
data={messages}
|
data={messages}
|
||||||
keyExtractor={(item) => item.id}
|
keyExtractor={(item) => item.id}
|
||||||
renderItem={({ item, index }) => (
|
renderItem={({ item, index }) => (
|
||||||
<Animated.View entering={FadeInDown.duration(300).delay(index > messages.length - 4 ? 100 : 0)}>
|
<MessageRow
|
||||||
<MessageBubble
|
|
||||||
message={item}
|
message={item}
|
||||||
onEdit={() => onEdit(item.id)}
|
delayed={index > messages.length - 4}
|
||||||
onRegenerate={onRegenerate}
|
|
||||||
onDelete={() => onDelete(item.id)}
|
|
||||||
isStreaming={isStreaming}
|
isStreaming={isStreaming}
|
||||||
isLastAssistant={item.id === lastAssistantId}
|
isLastAssistant={item.id === lastAssistantId}
|
||||||
|
onEdit={onEdit}
|
||||||
|
onDelete={onDelete}
|
||||||
|
onRegenerate={onRegenerate}
|
||||||
/>
|
/>
|
||||||
</Animated.View>
|
|
||||||
)}
|
)}
|
||||||
contentContainerStyle={styles.listContent}
|
contentContainerStyle={styles.listContent}
|
||||||
/>
|
/>
|
||||||
|
|||||||
@@ -0,0 +1,54 @@
|
|||||||
|
import AsyncStorage from '@react-native-async-storage/async-storage';
|
||||||
|
import { AppState } from 'react-native';
|
||||||
|
import type { StateStorage } from 'zustand/middleware';
|
||||||
|
|
||||||
|
export interface DebouncedStorage extends StateStorage {
|
||||||
|
/** Writes every pending value right away. */
|
||||||
|
flush: () => Promise<void>;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Wraps a storage so bursts of writes collapse into one trailing write.
|
||||||
|
*
|
||||||
|
* The chat store is persisted on every state change and a streamed answer
|
||||||
|
* changes it once per token, which would serialize every saved session and hit
|
||||||
|
* AsyncStorage dozens of times per second.
|
||||||
|
*/
|
||||||
|
export function createDebouncedStorage(base: StateStorage, delayMs = 400): DebouncedStorage {
|
||||||
|
const pending = new Map<string, string>();
|
||||||
|
let timer: ReturnType<typeof setTimeout> | null = null;
|
||||||
|
|
||||||
|
const flush = async () => {
|
||||||
|
if (timer) {
|
||||||
|
clearTimeout(timer);
|
||||||
|
timer = null;
|
||||||
|
}
|
||||||
|
const entries = [...pending.entries()];
|
||||||
|
pending.clear();
|
||||||
|
await Promise.all(entries.map(([name, value]) => base.setItem(name, value)));
|
||||||
|
};
|
||||||
|
|
||||||
|
return {
|
||||||
|
getItem: (name) => {
|
||||||
|
const value = pending.get(name);
|
||||||
|
return value === undefined ? base.getItem(name) : Promise.resolve(value);
|
||||||
|
},
|
||||||
|
setItem: (name, value) => {
|
||||||
|
pending.set(name, value);
|
||||||
|
if (timer) clearTimeout(timer);
|
||||||
|
timer = setTimeout(flush, delayMs);
|
||||||
|
},
|
||||||
|
removeItem: (name) => {
|
||||||
|
pending.delete(name);
|
||||||
|
return base.removeItem(name);
|
||||||
|
},
|
||||||
|
flush,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
export const appStorage = createDebouncedStorage(AsyncStorage);
|
||||||
|
|
||||||
|
// Do not lose the last tokens when the app is sent to the background or closed.
|
||||||
|
AppState.addEventListener('change', (state) => {
|
||||||
|
if (state !== 'active') void appStorage.flush();
|
||||||
|
});
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
import AsyncStorage from '@react-native-async-storage/async-storage';
|
|
||||||
import { create } from 'zustand';
|
import { create } from 'zustand';
|
||||||
import { createJSONStorage, persist } from 'zustand/middleware';
|
import { createJSONStorage, persist } from 'zustand/middleware';
|
||||||
import { DEFAULT_SERVER_CONFIG, LMModel, ServerConfig } from '../services/types';
|
import { DEFAULT_SERVER_CONFIG, LMModel, ServerConfig } from '../services/types';
|
||||||
|
import { appStorage } from './debouncedStorage';
|
||||||
|
|
||||||
export interface ChatMessage {
|
export interface ChatMessage {
|
||||||
id: string;
|
id: string;
|
||||||
@@ -164,7 +164,7 @@ export const useAppStore = create<AppState>()(
|
|||||||
}),
|
}),
|
||||||
{
|
{
|
||||||
name: 'localaiapp-storage',
|
name: 'localaiapp-storage',
|
||||||
storage: createJSONStorage(() => AsyncStorage),
|
storage: createJSONStorage(() => appStorage),
|
||||||
partialize: (state) => ({
|
partialize: (state) => ({
|
||||||
theme: state.theme,
|
theme: state.theme,
|
||||||
serverConfig: state.serverConfig,
|
serverConfig: state.serverConfig,
|
||||||
|
|||||||
Reference in new issue
Block a user