diff --git a/.changeset/perf-message-token-cache.md b/.changeset/perf-message-token-cache.md new file mode 100644 index 000000000..69b55034f --- /dev/null +++ b/.changeset/perf-message-token-cache.md @@ -0,0 +1,5 @@ +--- +"@nanocollective/nanocoder": patch +--- + +Update the message token cache in place instead of copying every entry into a new map on each cache miss. diff --git a/source/hooks/useAppState.spec.tsx b/source/hooks/useAppState.spec.tsx index 4797269d2..eb58a02ce 100644 --- a/source/hooks/useAppState.spec.tsx +++ b/source/hooks/useAppState.spec.tsx @@ -22,14 +22,21 @@ console.log('\nuseAppState.spec.tsx'); type AppStateHook = ReturnType; let captured: AppStateHook | null = null; +let renderCount = 0; function Probe({initialMode}: {initialMode?: DevelopmentMode}) { + renderCount++; captured = useAppState(initialMode ?? 'normal'); return null; } +function tokenCacheKey(message: Message, model = '') { + return (message.content || '') + message.role + model; +} + function setup(initialMode: DevelopmentMode = 'normal') { captured = null; + renderCount = 0; const instance = render(); if (!captured) throw new Error('useAppState did not initialize'); return {hook: captured as AppStateHook, instance}; @@ -252,14 +259,106 @@ test('tokenizer is rebuilt when provider or model changes', t => { t.not(captured!.tokenizer, initial); }); -test('getMessageTokens returns a number and caches the result', t => { - const {hook} = setup(); +test('getMessageTokens returns a number and caches it in place', t => { + const {hook, instance} = setup(); + const cache = hook.messageTokenCache; const msg: Message = {role: 'user', content: 'hello world'} as Message; const tokens = hook.getMessageTokens(msg); t.is(typeof tokens, 'number'); t.true(tokens >= 0); + t.is(cache.size, 1); + t.is(cache.get(tokenCacheKey(msg)), tokens); + + t.is(hook.getMessageTokens(msg), tokens); + t.is(cache.size, 1); + + instance.rerender(); + + t.is(captured!.messageTokenCache, cache); + t.is(captured!.messageTokenCache.get(tokenCacheKey(msg)), tokens); +}); + +test('getMessageTokens returns a cached entry instead of recomputing', t => { + const {hook} = setup(); + + const msg: Message = {role: 'user', content: 'seeded'} as Message; + t.true(hook.getMessageTokens(msg) > 0); + + hook.messageTokenCache.set(tokenCacheKey(msg), 0); + + t.is(hook.getMessageTokens(msg), 0); + t.is(hook.messageTokenCache.size, 1); +}); + +test('a cache miss neither re-renders nor invalidates getMessageTokens', async t => { + const {hook, instance} = setup(); + const rendersAfterMount = renderCount; + const {getMessageTokens} = hook; + + getMessageTokens({role: 'user', content: 'uncached'} as Message); + await new Promise(resolve => setTimeout(resolve, 20)); + + t.is(renderCount, rendersAfterMount); + + instance.rerender(); + + t.is(captured!.getMessageTokens, getMessageTokens); + t.is(captured!.messageTokenCache.size, 1); +}); + +test('token cache keys separate content, role and model', t => { + const {hook, instance} = setup(); + const cache = hook.messageTokenCache; + const user: Message = {role: 'user', content: 'same text'} as Message; + + const userTokens = hook.getMessageTokens(user); + hook.getMessageTokens({role: 'assistant', content: 'same text'} as Message); + hook.getMessageTokens({role: 'user', content: 'other text'} as Message); + + t.is(cache.size, 3); + + hook.setCurrentModel('gpt-4o'); + instance.rerender(); + const switchedTokens = captured!.getMessageTokens(user); + + t.is(captured!.messageTokenCache, cache); + t.is(cache.size, 4); + t.is(cache.get(tokenCacheKey(user)), userTokens); + t.is(cache.get(tokenCacheKey(user, 'gpt-4o')), switchedTokens); +}); + +test('token cache stays bounded and evicts the oldest entry', t => { + const {hook} = setup(); + const cache = hook.messageTokenCache; + const oldest: Message = {role: 'user', content: 'message 0'} as Message; + const newest: Message = {role: 'user', content: 'message 1000'} as Message; + + for (let i = 0; i <= 1000; i++) { + hook.getMessageTokens({role: 'user', content: `message ${i}`} as Message); + } + + t.is(cache.size, 1000); + t.is(cache.get(tokenCacheKey(oldest)), undefined); + t.true(cache.get(tokenCacheKey(newest))! > 0); + + const recomputed = hook.getMessageTokens(oldest); + + t.is(cache.size, 1000); + t.is(cache.get(tokenCacheKey(oldest)), recomputed); +}); + +test('getMessageTokens handles messages without content', t => { + const {hook} = setup(); + const empty: Message = {role: 'user', content: ''} as Message; + const missing = {role: 'user'} as unknown as Message; + + const emptyTokens = hook.getMessageTokens(empty); + + t.is(typeof emptyTokens, 'number'); + t.is(hook.getMessageTokens(missing), emptyTokens); + t.is(hook.messageTokenCache.size, 1); }); test('exposes setters for every state slice', t => { diff --git a/source/hooks/useAppState.tsx b/source/hooks/useAppState.tsx index 28030d511..44b6f92ba 100644 --- a/source/hooks/useAppState.tsx +++ b/source/hooks/useAppState.tsx @@ -52,14 +52,14 @@ export function useAppState( const [client, setClient] = useState(null); const [messages, setMessages] = useState([]); - const [messageTokenCache, setMessageTokenCache] = useState< - BoundedMap - >( - new BoundedMap({ + const messageTokenCacheRef = useRef | null>(null); + if (!messageTokenCacheRef.current) { + messageTokenCacheRef.current = new BoundedMap({ maxSize: 1000, // No TTL - cache is session-based and cleared on app restart - }), - ); + }); + } + const messageTokenCache = messageTokenCacheRef.current; const [currentModel, setCurrentModel] = useState(''); const [currentProvider, setCurrentProvider] = useState('openai-compatible'); @@ -261,22 +261,7 @@ export function useAppState( } const tokens = tokenizer.countTokens(message); - // Defer cache update to avoid "Cannot update a component while rendering" error - // This can happen when components call getMessageTokens during their render - queueMicrotask(() => { - setMessageTokenCache(prev => { - const newCache = new BoundedMap({ - maxSize: 1000, - }); - // Copy existing entries - for (const [k, v] of prev.entries()) { - newCache.set(k, v); - } - // Add new entry - newCache.set(cacheKey, tokens); - return newCache; - }); - }); + messageTokenCache.set(cacheKey, tokens); return tokens; }, [messageTokenCache, tokenizer, currentModel], @@ -390,7 +375,6 @@ export function useAppState( // Setters setClient, setMessages, - setMessageTokenCache, setCurrentModel, setCurrentProvider, setCurrentProviderConfig,