From 0e5b38fd15745bafd1a85f2f83fe0662c1626b97 Mon Sep 17 00:00:00 2001 From: Evandro Fonseca Junior Date: Sat, 10 Oct 2026 04:28:38 -0300 Subject: [PATCH] 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) --- __tests__/MessageList.test.tsx | 78 ++++++++++++++++++++++++++++++ __tests__/debouncedStorage.test.ts | 72 +++++++++++++++++++++++++++ __tests__/useAppStore.test.ts | 3 +- components/chat/MessageList.tsx | 54 +++++++++++++++++---- store/debouncedStorage.ts | 54 +++++++++++++++++++++ store/useAppStore.ts | 4 +- 6 files changed, 252 insertions(+), 13 deletions(-) create mode 100644 __tests__/MessageList.test.tsx create mode 100644 __tests__/debouncedStorage.test.ts create mode 100644 store/debouncedStorage.ts diff --git a/__tests__/MessageList.test.tsx b/__tests__/MessageList.test.tsx new file mode 100644 index 0000000..38f77c7 --- /dev/null +++ b/__tests__/MessageList.test.tsx @@ -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 = {}; +jest.mock('../components/chat/MessageBubble', () => { + const { Text } = require('react-native'); + return { + MessageBubble: ({ message }: any) => { + mockRenders[message.id] = (mockRenders[message.id] ?? 0) + 1; + return {`${message.id}:${message.content}`}; + }, + }; +}); +jest.mock('../components/chat/TypingIndicator', () => { + const { Text } = require('react-native'); + return { TypingIndicator: () => typing-indicator }; +}); + +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( + , + ); + 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( + , + ); + expect(mockRenders).toEqual({ u1: 1, a0: 1, a1: 1 }); + + rerender( + , + ); + 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( + , + ); + expect(screen.getByText('typing-indicator')).toBeTruthy(); + }); +}); diff --git a/__tests__/debouncedStorage.test.ts b/__tests__/debouncedStorage.test.ts new file mode 100644 index 0000000..5541d85 --- /dev/null +++ b/__tests__/debouncedStorage.test.ts @@ -0,0 +1,72 @@ +import { createDebouncedStorage } from '../store/debouncedStorage'; + +const makeBase = () => ({ + getItem: jest.fn(async (_name: string): Promise => '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'); + }); +}); diff --git a/__tests__/useAppStore.test.ts b/__tests__/useAppStore.test.ts index f6e43a2..07619f7 100644 --- a/__tests__/useAppStore.test.ts +++ b/__tests__/useAppStore.test.ts @@ -1,4 +1,5 @@ import AsyncStorage from '@react-native-async-storage/async-storage'; +import { appStorage } from '../store/debouncedStorage'; import { useAppStore } from '../store/useAppStore'; const state = () => useAppStore.getState(); @@ -172,7 +173,7 @@ describe('persistence', () => { }); 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 persisted = JSON.parse(raw as string).state; diff --git a/components/chat/MessageList.tsx b/components/chat/MessageList.tsx index bf808b6..f4be8a9 100644 --- a/components/chat/MessageList.tsx +++ b/components/chat/MessageList.tsx @@ -14,6 +14,41 @@ interface MessageListProps { 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 ( + + onEdit(message.id)} + onRegenerate={onRegenerate} + onDelete={() => onDelete(message.id)} + isStreaming={isStreaming} + isLastAssistant={isLastAssistant} + /> + + ); +}); + export function MessageList({ messages, isStreaming, onEdit, onDelete, onRegenerate }: MessageListProps) { const flatListRef = useRef(null); const lastAssistantId = findLastAssistantId(messages); @@ -34,16 +69,15 @@ export function MessageList({ messages, isStreaming, onEdit, onDelete, onRegener data={messages} keyExtractor={(item) => item.id} renderItem={({ item, index }) => ( - messages.length - 4 ? 100 : 0)}> - onEdit(item.id)} - onRegenerate={onRegenerate} - onDelete={() => onDelete(item.id)} - isStreaming={isStreaming} - isLastAssistant={item.id === lastAssistantId} - /> - + messages.length - 4} + isStreaming={isStreaming} + isLastAssistant={item.id === lastAssistantId} + onEdit={onEdit} + onDelete={onDelete} + onRegenerate={onRegenerate} + /> )} contentContainerStyle={styles.listContent} /> diff --git a/store/debouncedStorage.ts b/store/debouncedStorage.ts new file mode 100644 index 0000000..dfb0862 --- /dev/null +++ b/store/debouncedStorage.ts @@ -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; +} + +/** + * 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(); + let timer: ReturnType | 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(); +}); diff --git a/store/useAppStore.ts b/store/useAppStore.ts index d796d1f..22856fe 100644 --- a/store/useAppStore.ts +++ b/store/useAppStore.ts @@ -1,7 +1,7 @@ -import AsyncStorage from '@react-native-async-storage/async-storage'; import { create } from 'zustand'; import { createJSONStorage, persist } from 'zustand/middleware'; import { DEFAULT_SERVER_CONFIG, LMModel, ServerConfig } from '../services/types'; +import { appStorage } from './debouncedStorage'; export interface ChatMessage { id: string; @@ -164,7 +164,7 @@ export const useAppStore = create()( }), { name: 'localaiapp-storage', - storage: createJSONStorage(() => AsyncStorage), + storage: createJSONStorage(() => appStorage), partialize: (state) => ({ theme: state.theme, serverConfig: state.serverConfig,