diff --git a/src/App.tsx b/src/App.tsx index ff3ac732..be0114af 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -113,15 +113,8 @@ function App() { [persistTurn], ); - const { - messages, - streamingContent, - ask, - cancel, - isGenerating, - reset, - loadMessages, - } = useOllama(handleTurnComplete); + const { messages, ask, cancel, isGenerating, reset, loadMessages } = + useOllama(handleTurnComplete); const inputRef = useRef(null); @@ -196,7 +189,7 @@ function App() { * complete response. We check for an assistant message rather than any message * so the button never appears during the very first user-only half-turn. */ - const canSave = messages.some((m) => m.role === 'assistant'); + const canSave = !isGenerating && messages.some((m) => m.role === 'assistant'); const shouldRenderOverlay = overlayState === 'visible'; /** @@ -1214,7 +1207,6 @@ function App() { ? [...messages, pendingUserMessage] : messages } - streamingContent={streamingContent} isGenerating={isGenerating || isSubmitPending} onClose={handleCloseOverlay} onSave={handleSave} diff --git a/src/components/ChatBubble.tsx b/src/components/ChatBubble.tsx index fea8b8eb..56757577 100644 --- a/src/components/ChatBubble.tsx +++ b/src/components/ChatBubble.tsx @@ -132,7 +132,7 @@ export function ChatBubble({ )} - {!errorKind && ( + {!errorKind && !isStreaming && (
diff --git a/src/hooks/__tests__/useOllama.test.tsx b/src/hooks/__tests__/useOllama.test.tsx index e2c0b49e..853a4ba1 100644 --- a/src/hooks/__tests__/useOllama.test.tsx +++ b/src/hooks/__tests__/useOllama.test.tsx @@ -71,13 +71,14 @@ describe('useOllama', () => { }); }); - it('adds user message immediately on ask', async () => { + it('adds user message and empty assistant placeholder immediately on ask', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { await result.current.ask('my question'); }); + expect(result.current.messages).toHaveLength(2); expect(result.current.messages[0]).toEqual( expect.objectContaining({ role: 'user', @@ -85,6 +86,12 @@ describe('useOllama', () => { }), ); expect(result.current.messages[0].id).toEqual(expect.any(String)); + expect(result.current.messages[1]).toEqual( + expect.objectContaining({ + role: 'assistant', + content: '', + }), + ); }); it('stores quotedText on user message when provided', async () => { @@ -119,7 +126,7 @@ describe('useOllama', () => { ); }); - it('accumulates streaming tokens into streamingContent', async () => { + it('accumulates streaming tokens into the assistant message', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -134,10 +141,13 @@ describe('useOllama', () => { channel!.simulateMessage({ type: 'Token', data: ', world' }); }); - expect(result.current.streamingContent).toBe('Hello, world'); + const assistantMsg = result.current.messages.find( + (m) => m.role === 'assistant', + ); + expect(assistantMsg?.content).toBe('Hello, world'); }); - it('moves content to messages on Done chunk', async () => { + it('keeps assistant message in place on Done chunk', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -152,7 +162,6 @@ describe('useOllama', () => { channel!.simulateMessage({ type: 'Done' }); }); - expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); expect(result.current.messages).toContainEqual( expect.objectContaining({ @@ -231,8 +240,8 @@ describe('useOllama', () => { await result.current.ask('', undefined, ['/tmp/img1.jpg']); }); - // Should have created a user message (not returned early) - expect(result.current.messages).toHaveLength(1); + // Should have created a user message + assistant placeholder + expect(result.current.messages).toHaveLength(2); expect(result.current.messages[0]).toEqual( expect.objectContaining({ role: 'user', @@ -352,7 +361,7 @@ describe('useOllama', () => { expect(result.current.isGenerating).toBe(false); }); - it('Error chunk creates assistant message with errorKind set', async () => { + it('Error chunk updates assistant placeholder with errorKind', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -381,7 +390,7 @@ describe('useOllama', () => { ); }); - it('Error chunk with partial tokens preserves prior content in separate message', async () => { + it('Error chunk with partial tokens replaces content with error', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -399,7 +408,7 @@ describe('useOllama', () => { }); }); - // The error message should appear as its own assistant message with errorKind + // The error replaces the assistant placeholder content const errorMsg = result.current.messages.find((m) => m.errorKind); expect(errorMsg).toBeDefined(); expect(errorMsg?.errorKind).toBe('Other'); @@ -415,11 +424,11 @@ describe('useOllama', () => { await result.current.ask('test'); }); - const assistantMsg = result.current.messages.find( - (m) => m.role === 'assistant', + const errorMsg = result.current.messages.find( + (m) => m.errorKind === 'Other', ); - expect(assistantMsg?.errorKind).toBe('Other'); - expect(assistantMsg?.content).toBeTruthy(); + expect(errorMsg?.errorKind).toBe('Other'); + expect(errorMsg?.content).toBeTruthy(); }); }); @@ -440,8 +449,11 @@ describe('useOllama', () => { channel!.simulateMessage({ type: 'Token', data: '' }); }); - // streamingContent should still be empty (no crash) - expect(result.current.streamingContent).toBe(''); + // Assistant content should still be empty (no crash) + const assistantMsg = result.current.messages.find( + (m) => m.role === 'assistant', + ); + expect(assistantMsg?.content).toBe(''); }); }); @@ -494,7 +506,7 @@ describe('useOllama', () => { // ─── Cancelled chunk handling ─────────────────────────────────────────────── describe('Cancelled chunk', () => { - it('finalizes partial content as assistant message on Cancelled', async () => { + it('keeps partial content as assistant message on Cancelled', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -510,7 +522,6 @@ describe('useOllama', () => { channel!.simulateMessage({ type: 'Cancelled' }); }); - expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); expect(result.current.messages).toContainEqual( expect.objectContaining({ @@ -520,7 +531,7 @@ describe('useOllama', () => { ); }); - it('does not add empty assistant message when cancelled with no tokens', async () => { + it('removes assistant placeholder when cancelled with no tokens', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { @@ -534,9 +545,8 @@ describe('useOllama', () => { channel!.simulateMessage({ type: 'Cancelled' }); }); - expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); - // Only the user message should exist — no empty assistant message + // Only the user message should exist — empty assistant placeholder was removed expect(result.current.messages).toHaveLength(1); expect(result.current.messages[0].role).toBe('user'); }); @@ -566,7 +576,6 @@ describe('useOllama', () => { }); expect(result.current.messages).toEqual([]); - expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); // Should also reset backend conversation history expect(invoke).toHaveBeenCalledWith('reset_conversation'); @@ -663,7 +672,7 @@ describe('useOllama', () => { expect(result.current.messages).toEqual(loaded); }); - it('clears streaming and generating state when loading messages', async () => { + it('clears generating state when loading messages', async () => { invoke.mockRejectedValueOnce(new Error('boom')); const { result } = renderHook(() => useOllama()); @@ -676,7 +685,6 @@ describe('useOllama', () => { result.current.loadMessages([]); }); - expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); }); }); diff --git a/src/hooks/useOllama.ts b/src/hooks/useOllama.ts index d0499a65..a1bb06b8 100644 --- a/src/hooks/useOllama.ts +++ b/src/hooks/useOllama.ts @@ -43,15 +43,14 @@ export function useOllama( onTurnComplete?: (userMsg: Message, assistantMsg: Message) => void, ) { const [messages, setMessages] = useState([]); - const [streamingContent, setStreamingContent] = useState(''); const [isGenerating, setIsGenerating] = useState(false); /** * Submits a message to the Ollama backend and initiates the streaming response. * The backend manages conversation history — only the new user message is sent. * - * Avoids continuous array copy operations during streaming by maintaining the streaming - * chunk state separately from the main messages state until generation finishes. + * Streams tokens directly into the messages array. An empty assistant placeholder + * is added immediately, then updated in-place on each token until generation finishes. * * @param displayContent The user's query as it should appear in the chat bubble. * @param quotedText Optional selected text quoted alongside this message. @@ -78,57 +77,54 @@ export function useOllama( imagePaths && imagePaths.length > 0 ? imagePaths : undefined, }; - setMessages((prev) => [...prev, userMsg]); - setStreamingContent(''); + const assistantId = crypto.randomUUID(); + const assistantMsg: Message = { + id: assistantId, + role: 'assistant', + content: '', + }; + + setMessages((prev) => [...prev, userMsg, assistantMsg]); setIsGenerating(true); const channel = new Channel(); - // Use block-scoped variable to accumulate the stream and occasionally flush to React state, - // mitigating rendering lag from hundreds of fast chunk events. let currentContent = ''; channel.onmessage = (chunk) => { if (chunk.type === 'Token') { currentContent += chunk.data; - setStreamingContent(currentContent); + setMessages((prev) => + prev.map((m) => + m.id === assistantId ? { ...m, content: currentContent } : m, + ), + ); } else if (chunk.type === 'Done') { - const assistantMsg: Message = { - id: crypto.randomUUID(), - role: 'assistant', - content: currentContent, - }; - setMessages((prev) => [...prev, assistantMsg]); - setStreamingContent(''); setIsGenerating(false); // Notify the caller that a complete turn has finished so it can // persist both messages to SQLite if the conversation is saved. - onTurnComplete?.(userMsg, assistantMsg); + onTurnComplete?.(userMsg, { + ...assistantMsg, + content: currentContent, + }); } else if (chunk.type === 'Cancelled') { - // Finalize partial content as a complete message so the user - // retains everything generated before they hit stop. - if (currentContent) { - setMessages((prev) => [ - ...prev, - { - id: crypto.randomUUID(), - role: 'assistant', - content: currentContent, - }, - ]); + // Remove the empty assistant placeholder if nothing was generated. + if (!currentContent) { + setMessages((prev) => prev.filter((m) => m.id !== assistantId)); } - setStreamingContent(''); setIsGenerating(false); } else { - setMessages((prev) => [ - ...prev, - { - id: crypto.randomUUID(), - role: 'assistant', - content: chunk.data.message, - errorKind: chunk.data.kind, - }, - ]); - setStreamingContent(''); + // Replace the streaming placeholder with an error message. + setMessages((prev) => + prev.map((m) => + m.id === assistantId + ? { + ...m, + content: chunk.data.message, + errorKind: chunk.data.kind, + } + : m, + ), + ); setIsGenerating(false); } }; @@ -150,7 +146,6 @@ export function useOllama( errorKind: 'Other' as const, }, ]); - setStreamingContent(''); setIsGenerating(false); } }, @@ -166,7 +161,6 @@ export function useOllama( /** Resets all conversation state to prepare for a fresh session. */ const reset = useCallback(() => { setMessages([]); - setStreamingContent(''); setIsGenerating(false); void invoke('reset_conversation'); }, []); @@ -182,13 +176,11 @@ export function useOllama( */ const loadMessages = useCallback((msgs: Message[]) => { setMessages(msgs); - setStreamingContent(''); setIsGenerating(false); }, []); return { messages, - streamingContent, ask, cancel, isGenerating, diff --git a/src/view/ConversationView.tsx b/src/view/ConversationView.tsx index d2e0a76b..8170c4bf 100644 --- a/src/view/ConversationView.tsx +++ b/src/view/ConversationView.tsx @@ -12,8 +12,6 @@ import type { Message } from '../hooks/useOllama'; interface ConversationViewProps { /** Array of completed messages in the conversation. */ messages: Message[]; - /** The actively streaming content for the current assistant response. */ - streamingContent: string; /** Whether the underlying LLM engine is currently generating a response. */ isGenerating: boolean; /** Callback fired when the user requests to close the overlay. */ @@ -59,7 +57,6 @@ interface ConversationViewProps { */ export function ConversationView({ messages, - streamingContent, isGenerating, onClose, onSave, @@ -150,7 +147,7 @@ export function ConversationView({ }); return () => cancelAnimationFrame(raf); - }, [messages, streamingContent]); + }, [messages, isGenerating]); return ( - {messages.map((msg, i) => ( - - ))} - - {/* Live-updating streaming bubble */} - {streamingContent ? ( - - ) : null} + {messages.map((msg, i) => { + const isLastAssistant = + isGenerating && + i === messages.length - 1 && + msg.role === 'assistant'; + + // Hide the empty assistant placeholder; the TypingIndicator + // already covers this visual state. + if (isLastAssistant && !msg.content) return null; + + return ( + + ); + })} {/* Typing indicator (pulsing dots) shown before first token arrives */} - {isGenerating && !streamingContent ? : null} + {isGenerating && + messages[messages.length - 1]?.role === 'assistant' && + !messages[messages.length - 1]?.content ? ( + + ) : null} { render( , @@ -20,11 +19,16 @@ describe('ConversationView', () => { expect(screen.getByText('Hi!')).toBeInTheDocument(); }); - it('renders streaming bubble when streamingContent is non-empty', () => { + it('renders streaming assistant message when isGenerating', () => { const { container } = render( , @@ -35,11 +39,10 @@ describe('ConversationView', () => { expect(container.textContent).toContain('response...'); }); - it('shows TypingIndicator when isGenerating with no streaming content', () => { + it('shows TypingIndicator when isGenerating with empty assistant content', () => { const { container } = render( , @@ -49,11 +52,12 @@ describe('ConversationView', () => { expect(dots.length).toBeGreaterThanOrEqual(9); }); - it('hides TypingIndicator when streaming content arrives', () => { + it('hides TypingIndicator when assistant content arrives', () => { const { container } = render( , @@ -65,12 +69,7 @@ describe('ConversationView', () => { it('renders WindowControls with onClose', () => { const onClose = vi.fn(); render( - , + , ); expect( screen.getByRole('button', { name: 'Close window' }), @@ -79,12 +78,7 @@ describe('ConversationView', () => { it('renders empty state with no messages (no .chat-bubble elements)', () => { const { container } = render( - , + , ); expect(container.querySelectorAll('.chat-bubble')).toHaveLength(0); }); @@ -93,7 +87,6 @@ describe('ConversationView', () => { const { container, rerender } = render( , @@ -122,8 +115,10 @@ describe('ConversationView', () => { act(() => { rerender( , @@ -138,7 +133,6 @@ describe('ConversationView', () => { const { container, rerender } = render( , @@ -169,7 +163,6 @@ describe('ConversationView', () => { { id: '1', role: 'user' as const, content: 'first' }, { id: '2', role: 'user' as const, content: 'second question' }, ]} - streamingContent="" isGenerating={true} onClose={vi.fn()} />, @@ -184,8 +177,10 @@ describe('ConversationView', () => { it('auto-scroll stays disabled when assistant message is finalized', () => { const { container, rerender } = render( , @@ -219,7 +214,6 @@ describe('ConversationView', () => { content: 'streaming reply', }, ]} - streamingContent="" isGenerating={false} onClose={vi.fn()} />, @@ -235,7 +229,6 @@ describe('ConversationView', () => { const { container, rerender } = render( , @@ -282,8 +275,10 @@ describe('ConversationView', () => { act(() => { rerender( , @@ -297,7 +292,6 @@ describe('ConversationView', () => { const { container, rerender } = render( , @@ -343,8 +337,10 @@ describe('ConversationView', () => { act(() => { rerender( , @@ -359,7 +355,6 @@ describe('ConversationView', () => { const { container, rerender } = render( , @@ -386,8 +381,10 @@ describe('ConversationView', () => { act(() => { rerender( , @@ -400,7 +397,6 @@ describe('ConversationView', () => { render( { render( , @@ -428,7 +423,6 @@ describe('ConversationView', () => { render( { render( { render( { render( { render( { render( ,