diff --git a/src/hooks/useAIChat.test.ts b/src/hooks/useAIChat.test.ts new file mode 100644 index 00000000..4cbed4d5 --- /dev/null +++ b/src/hooks/useAIChat.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import { renderHook, act } from '@testing-library/react' + +// Capture what streamClaudeChat receives +const streamClaudeChatMock = vi.fn< + Parameters, + ReturnType +>() + +vi.mock('../utils/ai-chat', async () => { + const actual = await vi.importActual('../utils/ai-chat') + return { + ...actual, + streamClaudeChat: (...args: Parameters) => { + streamClaudeChatMock(...args) + // Simulate async: call onDone after a tick + const callbacks = args[3] + setTimeout(() => { + callbacks.onInit?.('test-session') + callbacks.onText('mock response') + callbacks.onDone() + }, 10) + return Promise.resolve('test-session') + }, + } +}) + +import { useAIChat } from './useAIChat' + +beforeEach(() => { + streamClaudeChatMock.mockClear() + vi.useFakeTimers() +}) + +afterEach(() => { + vi.useRealTimers() +}) + +describe('useAIChat', () => { + const emptyContent: Record = {} + + it('sends first message without history', async () => { + const { result } = renderHook(() => useAIChat(emptyContent, [])) + + act(() => { result.current.sendMessage('hello') }) + + expect(streamClaudeChatMock).toHaveBeenCalledTimes(1) + const message = streamClaudeChatMock.mock.calls[0][0] + // First message: no history, so just the plain text + expect(message).toBe('hello') + }) + + it('includes conversation history in second message', async () => { + const { result } = renderHook(() => useAIChat(emptyContent, [])) + + // Send first message + act(() => { result.current.sendMessage('What is Rust?') }) + // Wait for mock response + await act(async () => { vi.advanceTimersByTime(50) }) + + // Now messages state has: [user: What is Rust?, assistant: mock response] + expect(result.current.messages).toHaveLength(2) + + // Send second message + act(() => { result.current.sendMessage('Tell me more') }) + + expect(streamClaudeChatMock).toHaveBeenCalledTimes(2) + const secondMessage = streamClaudeChatMock.mock.calls[1][0] + // Should contain conversation history + expect(secondMessage).toContain('What is Rust?') + expect(secondMessage).toContain('mock response') + expect(secondMessage).toContain('Tell me more') + }) + + it('resets history on clearConversation', async () => { + const { result } = renderHook(() => useAIChat(emptyContent, [])) + + // Send a message and get response + act(() => { result.current.sendMessage('hello') }) + await act(async () => { vi.advanceTimersByTime(50) }) + + // Clear conversation + act(() => { result.current.clearConversation() }) + expect(result.current.messages).toHaveLength(0) + + // Send new message — should have no history + act(() => { result.current.sendMessage('fresh start') }) + + const lastCall = streamClaudeChatMock.mock.calls[streamClaudeChatMock.mock.calls.length - 1] + expect(lastCall[0]).toBe('fresh start') + }) + + it('accumulates multiple exchanges in history', async () => { + const { result } = renderHook(() => useAIChat(emptyContent, [])) + + // Exchange 1 + act(() => { result.current.sendMessage('Q1') }) + await act(async () => { vi.advanceTimersByTime(50) }) + + // Exchange 2 + act(() => { result.current.sendMessage('Q2') }) + await act(async () => { vi.advanceTimersByTime(50) }) + + // Exchange 3 + act(() => { result.current.sendMessage('Q3') }) + + expect(streamClaudeChatMock).toHaveBeenCalledTimes(3) + const thirdMessage = streamClaudeChatMock.mock.calls[2][0] + // Should contain all prior exchanges + expect(thirdMessage).toContain('Q1') + expect(thirdMessage).toContain('Q2') + expect(thirdMessage).toContain('Q3') + }) +}) diff --git a/src/hooks/useAIChat.ts b/src/hooks/useAIChat.ts index 14d7fb1a..b0ba1f87 100644 --- a/src/hooks/useAIChat.ts +++ b/src/hooks/useAIChat.ts @@ -7,6 +7,7 @@ import type { VaultEntry } from '../types' import { type ChatMessage, nextMessageId, buildSystemPrompt, streamClaudeChat, + trimHistory, formatMessageWithHistory, MAX_HISTORY_TOKENS, } from '../utils/ai-chat' export function useAIChat( @@ -29,9 +30,11 @@ export function useAIChat( abortRef.current = false const { prompt: systemPrompt } = buildSystemPrompt(contextNotes, allContent) + const history = trimHistory(messages, MAX_HISTORY_TOKENS) + const messageWithHistory = formatMessageWithHistory(history, text.trim()) let accumulated = '' - streamClaudeChat(text.trim(), systemPrompt || undefined, sessionIdRef.current, { + streamClaudeChat(messageWithHistory, systemPrompt || undefined, sessionIdRef.current, { onInit: (sid) => { sessionIdRef.current = sid }, onText: (chunk) => { @@ -58,7 +61,7 @@ export function useAIChat( }).then(sid => { if (sid) sessionIdRef.current = sid }) - }, [isStreaming, allContent, contextNotes]) + }, [isStreaming, allContent, contextNotes, messages]) const clearConversation = useCallback(() => { abortRef.current = true diff --git a/src/utils/ai-chat.test.ts b/src/utils/ai-chat.test.ts index c50a301a..23b0fd69 100644 --- a/src/utils/ai-chat.test.ts +++ b/src/utils/ai-chat.test.ts @@ -8,6 +8,8 @@ vi.mock('../mock-tauri', () => ({ import { estimateTokens, buildSystemPrompt, nextMessageId, checkClaudeCli, streamClaudeChat, + trimHistory, formatMessageWithHistory, + type ChatMessage, MAX_HISTORY_TOKENS, } from './ai-chat' import type { VaultEntry } from '../types' @@ -73,6 +75,107 @@ describe('checkClaudeCli', () => { }) }) +// --- trimHistory --- + +describe('trimHistory', () => { + const msg = (role: 'user' | 'assistant', content: string): ChatMessage => ({ + role, content, id: `msg-${content}`, + }) + + it('returns empty array for empty history', () => { + expect(trimHistory([], 1000)).toEqual([]) + }) + + it('returns all messages when under token limit', () => { + const history = [msg('user', 'hi'), msg('assistant', 'hello')] + expect(trimHistory(history, 1000)).toEqual(history) + }) + + it('drops oldest messages when over token limit', () => { + const history = [ + msg('user', 'a'.repeat(400)), // 100 tokens + msg('assistant', 'b'.repeat(400)), // 100 tokens + msg('user', 'c'.repeat(400)), // 100 tokens + ] + const result = trimHistory(history, 200) + // Should keep the two most recent messages (200 tokens) + expect(result).toHaveLength(2) + expect(result[0].content).toBe('b'.repeat(400)) + expect(result[1].content).toBe('c'.repeat(400)) + }) + + it('keeps at least one message if it fits', () => { + const history = [ + msg('user', 'a'.repeat(2000)), // 500 tokens + msg('assistant', 'b'.repeat(80)), // 20 tokens + ] + const result = trimHistory(history, 30) + expect(result).toHaveLength(1) + expect(result[0].content).toBe('b'.repeat(80)) + }) + + it('returns empty when single message exceeds limit', () => { + const history = [msg('user', 'a'.repeat(4000))] // 1000 tokens + expect(trimHistory(history, 10)).toEqual([]) + }) +}) + +// --- formatMessageWithHistory --- + +describe('formatMessageWithHistory', () => { + const msg = (role: 'user' | 'assistant', content: string): ChatMessage => ({ + role, content, id: `msg-${content}`, + }) + + it('returns bare message when no history', () => { + expect(formatMessageWithHistory([], 'hello')).toBe('hello') + }) + + it('includes conversation history before the new message', () => { + const history = [msg('user', 'What is Rust?'), msg('assistant', 'A systems language.')] + const result = formatMessageWithHistory(history, 'How does it compare to Go?') + expect(result).toContain('What is Rust?') + expect(result).toContain('A systems language.') + expect(result).toContain('How does it compare to Go?') + }) + + it('labels user and assistant messages correctly', () => { + const history = [msg('user', 'Q1'), msg('assistant', 'A1')] + const result = formatMessageWithHistory(history, 'Q2') + expect(result).toContain('[user]: Q1') + expect(result).toContain('[assistant]: A1') + expect(result).toContain('[user]: Q2') + }) + + it('preserves message order', () => { + const history = [ + msg('user', 'first'), + msg('assistant', 'second'), + msg('user', 'third'), + msg('assistant', 'fourth'), + ] + const result = formatMessageWithHistory(history, 'fifth') + const firstIdx = result.indexOf('first') + const secondIdx = result.indexOf('second') + const thirdIdx = result.indexOf('third') + const fourthIdx = result.indexOf('fourth') + const fifthIdx = result.indexOf('fifth') + expect(firstIdx).toBeLessThan(secondIdx) + expect(secondIdx).toBeLessThan(thirdIdx) + expect(thirdIdx).toBeLessThan(fourthIdx) + expect(fourthIdx).toBeLessThan(fifthIdx) + }) +}) + +// --- MAX_HISTORY_TOKENS --- + +describe('MAX_HISTORY_TOKENS', () => { + it('is a reasonable token limit', () => { + expect(MAX_HISTORY_TOKENS).toBeGreaterThan(10_000) + expect(MAX_HISTORY_TOKENS).toBeLessThan(200_000) + }) +}) + // --- streamClaudeChat --- describe('streamClaudeChat', () => { diff --git a/src/utils/ai-chat.ts b/src/utils/ai-chat.ts index fd40e13b..d30b9dee 100644 --- a/src/utils/ai-chat.ts +++ b/src/utils/ai-chat.ts @@ -76,6 +76,34 @@ export function nextMessageId(): string { return `msg-${++msgIdCounter}-${Date.now()}` } +// --- Conversation history --- + +/** Max tokens of history to include in each request. */ +export const MAX_HISTORY_TOKENS = 100_000 + +/** Keep the most recent messages that fit within `maxTokens`. Drops oldest first. */ +export function trimHistory(history: ChatMessage[], maxTokens: number): ChatMessage[] { + let tokenCount = 0 + const result: ChatMessage[] = [] + for (let i = history.length - 1; i >= 0; i--) { + const tokens = estimateTokens(history[i].content) + if (tokenCount + tokens > maxTokens) break + result.unshift(history[i]) + tokenCount += tokens + } + return result +} + +/** Format conversation history + new message into a single prompt for the CLI. */ +export function formatMessageWithHistory(history: ChatMessage[], newMessage: string): string { + if (history.length === 0) return newMessage + + const lines = history.map(m => `[${m.role}]: ${m.content}`) + lines.push(`[user]: ${newMessage}`) + + return `\n${lines.join('\n\n')}\n\n\nContinue the conversation. Respond only to the latest [user] message.` +} + // --- Claude CLI status --- export interface ClaudeCliStatus {