diff --git a/src/App.tsx b/src/App.tsx index 39064bf..294e795 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -386,6 +386,14 @@ function App() { * deferred until Framer Motion finishes the exit transition. */ const requestHideOverlay = useCallback(() => { + // Cancel any in-flight work: active Ollama streaming, image processing, + // or screen capture. This ensures hiding the overlay (via double-tap + // Control, Escape, Cmd+W, or the X button) behaves like pressing Stop. + cancel(); + pendingSubmitRef.current = null; + setIsSubmitPending(false); + setPendingUserMessage(null); + windowAnchorRef.current = null; isPreExpandedRef.current = false; /* v8 ignore start -- DOM ref null guard: always set when overlay is visible */ @@ -407,7 +415,7 @@ function App() { } return 'hiding'; }); - }, []); + }, [cancel]); /** Ref attached to the chat-mode history dropdown for click-outside detection. */ const historyDropdownRef = useRef(null); diff --git a/src/hooks/__tests__/useOllama.test.tsx b/src/hooks/__tests__/useOllama.test.tsx index e2c0b49..9338469 100644 --- a/src/hooks/__tests__/useOllama.test.tsx +++ b/src/hooks/__tests__/useOllama.test.tsx @@ -573,6 +573,118 @@ describe('useOllama', () => { }); }); + // ─── Stale channel after reset ─────────────────────────────────────────────── + + describe('stale channel after reset', () => { + it('ignores Token chunks arriving after reset()', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.ask('hello'); + }); + + const channel = getChannel(); + + act(() => { + channel!.simulateMessage({ type: 'Token', data: 'Partial' }); + }); + expect(result.current.streamingContent).toBe('Partial'); + + // Reset clears state and bumps the epoch + act(() => { + result.current.reset(); + }); + expect(result.current.messages).toEqual([]); + expect(result.current.streamingContent).toBe(''); + + // Old channel sends more chunks after the reset + act(() => { + channel!.simulateMessage({ type: 'Token', data: ' stale token' }); + }); + + // Should remain empty; the stale token must be discarded + expect(result.current.streamingContent).toBe(''); + expect(result.current.messages).toEqual([]); + }); + + it('ignores Cancelled chunk arriving after reset()', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.ask('hello'); + }); + + const channel = getChannel(); + + act(() => { + channel!.simulateMessage({ type: 'Token', data: 'Partial' }); + }); + + act(() => { + result.current.reset(); + }); + + // Old channel delivers Cancelled after reset + act(() => { + channel!.simulateMessage({ type: 'Cancelled' }); + }); + + // The partial content must NOT reappear as a finalized message + expect(result.current.messages).toEqual([]); + expect(result.current.isGenerating).toBe(false); + }); + + it('ignores Done chunk arriving after reset()', async () => { + const onTurnComplete = vi.fn(); + const { result } = renderHook(() => useOllama(onTurnComplete)); + + await act(async () => { + await result.current.ask('hello'); + }); + + const channel = getChannel(); + + act(() => { + channel!.simulateMessage({ type: 'Token', data: 'Full answer' }); + }); + + act(() => { + result.current.reset(); + }); + + // Old channel delivers Done after reset + act(() => { + channel!.simulateMessage({ type: 'Done' }); + }); + + expect(result.current.messages).toEqual([]); + expect(onTurnComplete).not.toHaveBeenCalled(); + }); + + it('ignores Error chunk arriving after reset()', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.ask('hello'); + }); + + const channel = getChannel(); + + act(() => { + result.current.reset(); + }); + + act(() => { + channel!.simulateMessage({ + type: 'Error', + data: { kind: 'Other', message: 'Something went wrong\nHTTP 500' }, + }); + }); + + expect(result.current.messages).toEqual([]); + }); + }); + // ─── onTurnComplete callback ───────────────────────────────────────────────── describe('onTurnComplete callback', () => { diff --git a/src/hooks/useOllama.ts b/src/hooks/useOllama.ts index d0499a6..304ec2f 100644 --- a/src/hooks/useOllama.ts +++ b/src/hooks/useOllama.ts @@ -1,4 +1,4 @@ -import { useState, useCallback } from 'react'; +import { useState, useCallback, useRef } from 'react'; import { invoke, Channel } from '@tauri-apps/api/core'; /** Mirrors the Rust OllamaErrorKind enum sent over IPC. */ @@ -46,6 +46,10 @@ export function useOllama( const [streamingContent, setStreamingContent] = useState(''); const [isGenerating, setIsGenerating] = useState(false); + // Epoch counter: bumped on every reset so that stale channel callbacks from + // a previous generation can detect they are outdated and bail out. + const epochRef = useRef(0); + /** * Submits a message to the Ollama backend and initiates the streaming response. * The backend manages conversation history — only the new user message is sent. @@ -82,12 +86,19 @@ export function useOllama( setStreamingContent(''); setIsGenerating(true); + // Snapshot the epoch so this channel's callbacks can detect a reset. + const epochAtStart = epochRef.current; + 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) => { + // A reset occurred since this generation started; discard all + // remaining chunks so stale content never re-populates the UI. + if (epochRef.current !== epochAtStart) return; + if (chunk.type === 'Token') { currentContent += chunk.data; setStreamingContent(currentContent); @@ -165,6 +176,7 @@ export function useOllama( /** Resets all conversation state to prepare for a fresh session. */ const reset = useCallback(() => { + epochRef.current += 1; setMessages([]); setStreamingContent(''); setIsGenerating(false);