From 7395c125a27ec57f59a057f287c36a9e76ba959b Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 15:23:03 -0500 Subject: [PATCH 1/8] feat: add stream cancellation with stop generating button Implement CancellationToken-based stream cancellation so users can stop Ollama inference mid-generation. Dropping the HTTP connection signals Ollama to halt via Go's context.Context propagation. - Add tokio-util CancellationToken in GenerationState (Tauri managed state) - Refactor stream_ollama to use tokio::select! for instant cancellation - New cancel_generation Tauri command + StreamChunk::Cancelled variant - Stop button replaces spinner in AskBarView during generation - Partial content preserved as assistant message on cancel - 100% test coverage maintained (144 frontend + 50 backend tests) Co-Authored-By: Claude Opus 4.6 (1M context) --- src-tauri/Cargo.lock | 1 + src-tauri/Cargo.toml | 3 +- src-tauri/src/commands.rs | 318 +++++++++++++++++++++---- src-tauri/src/lib.rs | 5 + src/App.tsx | 12 +- src/hooks/__tests__/useOllama.test.tsx | 97 ++++++++ src/hooks/useOllama.ts | 32 ++- src/view/AskBarView.tsx | 45 ++-- src/view/__tests__/AskBarView.test.tsx | 64 +++++ 9 files changed, 509 insertions(+), 68 deletions(-) diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index eda19a23..dd7eecf9 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -4119,6 +4119,7 @@ dependencies = [ "tauri-build", "tauri-nspanel", "tokio", + "tokio-util", ] [[package]] diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 1e68d1c8..f04c49b2 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -22,8 +22,9 @@ tauri = { version = "2", features = ["macos-private-api", "tray-icon", "image-pn serde = { version = "1", features = ["derive"] } serde_json = "1" reqwest = { version = "0.13.2", features = ["json", "stream"] } -tokio = "1.50.0" +tokio = { version = "1.50.0", features = ["macros"] } futures-util = "0.3.32" +tokio-util = "0.7" [target.'cfg(target_os = "macos")'.dependencies] tauri-nspanel = { git = "https://github.com/ahkohd/tauri-nspanel", branch = "v2.1" } diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 4d92d3c2..28611696 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1,24 +1,28 @@ use futures_util::StreamExt; use serde::{Deserialize, Serialize}; +use std::sync::Mutex; use tauri::{ipc::Channel, State}; +use tokio_util::sync::CancellationToken; /// Default configuration constants as the application currently lacks a Settings UI. const DEFAULT_OLLAMA_URL: &str = "http://127.0.0.1:11434"; const DEFAULT_MODEL_NAME: &str = "llama3.2:3b"; -/// Payload emitted back to the frontend per token chunk +/// Payload emitted back to the frontend per token chunk. #[derive(Clone, Serialize)] #[serde(tag = "type", content = "data")] pub enum StreamChunk { - /// A single token chunk string + /// A single token chunk string. Token(String), - /// Indicates the stream has fully completed + /// Indicates the stream has fully completed. Done, - /// An error occurred during processing + /// The user explicitly cancelled generation. + Cancelled, + /// An error occurred during processing. Error(String), } -/// Request payload to Ollama server +/// Request payload to Ollama server. #[derive(Serialize)] struct OllamaRequest { model: String, @@ -26,19 +30,66 @@ struct OllamaRequest { stream: bool, } -/// Expected structured response chunk from Ollama +/// Expected structured response chunk from Ollama. #[derive(Deserialize)] struct OllamaResponse { response: Option, done: Option, } +/// Holds the active cancellation token for the current generation request. +/// +/// Only one generation runs at a time — starting a new request replaces the +/// previous token. `cancel_generation` cancels whatever is currently active. +pub struct GenerationState { + token: Mutex>, +} + +impl Default for GenerationState { + fn default() -> Self { + Self::new() + } +} + +impl GenerationState { + /// Creates a new empty generation state with no active token. + pub fn new() -> Self { + Self { + token: Mutex::new(None), + } + } + + /// Stores a new cancellation token, replacing any previous one. + fn set(&self, token: CancellationToken) { + *self.token.lock().unwrap() = Some(token); + } + + /// Cancels the active generation, if any, and clears the stored token. + fn cancel(&self) { + if let Some(token) = self.token.lock().unwrap().take() { + token.cancel(); + } + } + + /// Clears the stored token without cancelling it (used on natural completion). + fn clear(&self) { + *self.token.lock().unwrap() = None; + } +} + /// Core streaming logic, separated from the Tauri command for testability. +/// +/// Streams newline-delimited JSON from the Ollama `/api/generate` endpoint, +/// emitting `StreamChunk` variants via the provided callback. Uses `tokio::select!` +/// to race each chunk read against the cancellation token, ensuring the HTTP +/// connection is dropped immediately when the user cancels — which signals +/// Ollama to stop inference via Go's `context.Context` propagation. pub async fn stream_ollama( endpoint: &str, model: &str, prompt: String, client: &reqwest::Client, + cancel_token: CancellationToken, on_chunk: impl Fn(StreamChunk), ) { let request_payload = OllamaRequest { @@ -66,35 +117,50 @@ pub async fn stream_ollama( let mut stream = response.bytes_stream(); let mut buffer: Vec = Vec::new(); - while let Some(chunk_res) = stream.next().await { - match chunk_res { - Ok(bytes) => { - buffer.extend_from_slice(&bytes); - - while let Some(idx) = buffer.iter().position(|&b| b == b'\n') { - let line_bytes = buffer.drain(..=idx).collect::>(); - if let Ok(line_text) = String::from_utf8(line_bytes) { - let trimmed = line_text.trim(); - if trimmed.is_empty() { - continue; - } + loop { + tokio::select! { + _ = cancel_token.cancelled() => { + // Drop the stream — closes the HTTP connection, + // which signals Ollama to stop inference. + drop(stream); + on_chunk(StreamChunk::Cancelled); + return; + } + chunk_opt = stream.next() => { + match chunk_opt { + Some(Ok(bytes)) => { + buffer.extend_from_slice(&bytes); + + while let Some(idx) = buffer.iter().position(|&b| b == b'\n') { + let line_bytes = buffer.drain(..=idx).collect::>(); + if let Ok(line_text) = String::from_utf8(line_bytes) { + let trimmed = line_text.trim(); + if trimmed.is_empty() { + continue; + } - if let Ok(json) = serde_json::from_str::(trimmed) { - if let Some(token) = json.response { - if !token.is_empty() { - on_chunk(StreamChunk::Token(token)); + if let Ok(json) = + serde_json::from_str::(trimmed) + { + if let Some(token) = json.response { + if !token.is_empty() { + on_chunk(StreamChunk::Token(token)); + } + } + if let Some(true) = json.done { + on_chunk(StreamChunk::Done); + } } } - if let Some(true) = json.done { - on_chunk(StreamChunk::Done); - } } } + Some(Err(e)) => { + on_chunk(StreamChunk::Error(e.to_string())); + return; + } + None => return, } } - Err(e) => { - on_chunk(StreamChunk::Error(e.to_string())); - } } } } @@ -105,31 +171,55 @@ pub async fn stream_ollama( } /// Streams text chunks from the local Ollama backend via `reqwest` to the frontend using `Channel`. -/// Uses `State` to persist the HTTP Client's connection pool. +/// +/// Creates a fresh `CancellationToken` for each request and stores it in `GenerationState`, +/// allowing `cancel_generation` to abort the stream at any time. #[cfg_attr(coverage_nightly, coverage(off))] #[cfg_attr(not(coverage), tauri::command)] pub async fn ask_ollama( prompt: String, on_event: Channel, client: State<'_, reqwest::Client>, + generation: State<'_, GenerationState>, ) -> Result<(), String> { let endpoint = format!("{}/api/generate", DEFAULT_OLLAMA_URL.trim_end_matches('/')); + let cancel_token = CancellationToken::new(); + generation.set(cancel_token.clone()); - stream_ollama(&endpoint, DEFAULT_MODEL_NAME, prompt, &client, |chunk| { - let _ = on_event.send(chunk); - }) + stream_ollama( + &endpoint, + DEFAULT_MODEL_NAME, + prompt, + &client, + cancel_token, + |chunk| { + let _ = on_event.send(chunk); + }, + ) .await; + generation.clear(); + Ok(()) +} + +/// Cancels the currently active generation, if any. +/// +/// Signals the `CancellationToken` stored in `GenerationState`, which causes the +/// `stream_ollama` loop to exit immediately and drop the HTTP connection to Ollama. +#[cfg_attr(coverage_nightly, coverage(off))] +#[cfg_attr(not(coverage), tauri::command)] +pub async fn cancel_generation(generation: State<'_, GenerationState>) -> Result<(), String> { + generation.cancel(); Ok(()) } #[cfg(test)] mod tests { use super::*; - use std::sync::{Arc, Mutex}; + use std::sync::{Arc, Mutex as StdMutex}; - fn collect_chunks() -> (Arc>>, impl Fn(StreamChunk)) { - let chunks: Arc>> = Arc::new(Mutex::new(Vec::new())); + fn collect_chunks() -> (Arc>>, impl Fn(StreamChunk)) { + let chunks: Arc>> = Arc::new(StdMutex::new(Vec::new())); let chunks_clone = chunks.clone(); let callback = move |chunk: StreamChunk| { chunks_clone.lock().unwrap().push(chunk); @@ -151,6 +241,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -158,6 +249,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -180,6 +272,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -187,6 +280,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -200,6 +294,7 @@ mod tests { #[tokio::test] async fn handles_connection_refused() { let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -207,6 +302,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -226,6 +322,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -233,6 +330,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -252,6 +350,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -259,6 +358,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -283,6 +383,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -290,6 +391,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -309,7 +411,6 @@ mod tests { #[tokio::test] async fn handles_invalid_utf8_in_stream() { let mut server = mockito::Server::new_async().await; - // Line 1: invalid UTF-8 bytes + newline, Line 2: valid JSON + newline let body = b"\xFF\xFE\n{\"response\":\"ok\",\"done\":true}\n".to_vec(); let mock = server .mock("POST", "/api/generate") @@ -318,6 +419,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -325,13 +427,13 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; mock.assert_async().await; let chunks = chunks.lock().unwrap(); - // Invalid UTF-8 line is silently skipped; valid line emits Done assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Done))); } @@ -345,7 +447,6 @@ mod tests { tokio::spawn(async move { let (mut stream, _) = listener.accept().await.unwrap(); - // Send valid HTTP headers then partial chunked body and drop let _ = stream .write_all( b"HTTP/1.1 200 OK\r\n\ @@ -354,10 +455,10 @@ mod tests { 4\r\ntest", ) .await; - // Drop stream — connection closed mid-chunk, triggering a reqwest stream error }); let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -365,13 +466,12 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; let chunks = chunks.lock().unwrap(); - // Either an Error chunk from the stream error, or empty if reqwest absorbed the truncation - // Either way, no panic and the function completed cleanly let has_no_tokens = chunks.iter().all(|c| !matches!(c, StreamChunk::Token(_))); assert!(has_no_tokens); } @@ -387,6 +487,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -394,6 +495,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -407,7 +509,6 @@ mod tests { #[tokio::test] async fn whitespace_only_lines_are_skipped() { let mut server = mockito::Server::new_async().await; - // A line of only spaces/tabs followed by a valid JSON line let mock = server .mock("POST", "/api/generate") .with_body(" \n{\"response\":\"hi\",\"done\":true}\n") @@ -415,6 +516,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -422,6 +524,7 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; @@ -434,7 +537,6 @@ mod tests { #[tokio::test] async fn response_field_absent_emits_only_done() { let mut server = mockito::Server::new_async().await; - // JSON where `response` key is missing — exercises the None arm of if-let Some(token) let mock = server .mock("POST", "/api/generate") .with_body("{\"done\":true}\n") @@ -442,6 +544,7 @@ mod tests { .await; let client = reqwest::Client::new(); + let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); stream_ollama( @@ -449,14 +552,141 @@ mod tests { "test-model", "hi".to_string(), &client, + token, callback, ) .await; mock.assert_async().await; let chunks = chunks.lock().unwrap(); - // No Token chunks since response is absent; Done is emitted assert!(chunks.iter().all(|c| !matches!(c, StreamChunk::Token(_)))); assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Done))); } + + #[tokio::test] + async fn cancellation_stops_stream_and_emits_cancelled() { + use tokio::io::AsyncWriteExt; + use tokio::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let _ = stream + .write_all( + b"HTTP/1.1 200 OK\r\n\ + Content-Type: application/x-ndjson\r\n\r\n\ + {\"response\":\"A\",\"done\":false}\n", + ) + .await; + // Keep the connection open so cancellation can interrupt + tokio::time::sleep(tokio::time::Duration::from_secs(10)).await; + }); + + let client = reqwest::Client::new(); + let token = CancellationToken::new(); + let token_clone = token.clone(); + let (chunks, callback) = collect_chunks(); + + tokio::spawn(async move { + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + token_clone.cancel(); + }); + + stream_ollama( + &format!("http://127.0.0.1:{}/api/generate", port), + "test-model", + "hi".to_string(), + &client, + token, + callback, + ) + .await; + + let chunks = chunks.lock().unwrap(); + assert!(chunks + .iter() + .any(|c| matches!(c, StreamChunk::Token(t) if t == "A"))); + assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Cancelled))); + assert!(chunks.iter().all(|c| !matches!(c, StreamChunk::Done))); + } + + #[tokio::test] + async fn pre_cancelled_token_emits_cancelled_immediately() { + let mut server = mockito::Server::new_async().await; + let _mock = server + .mock("POST", "/api/generate") + .with_body( + "{\"response\":\"Hello\",\"done\":false}\n\ + {\"response\":\"\",\"done\":true}\n", + ) + .create_async() + .await; + + let client = reqwest::Client::new(); + let token = CancellationToken::new(); + token.cancel(); + + let (chunks, callback) = collect_chunks(); + + stream_ollama( + &format!("{}/api/generate", server.url()), + "test-model", + "hi".to_string(), + &client, + token, + callback, + ) + .await; + + let chunks = chunks.lock().unwrap(); + assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Cancelled))); + } + + #[test] + fn generation_state_set_and_cancel() { + let state = GenerationState::new(); + let token = CancellationToken::new(); + let token_clone = token.clone(); + + state.set(token); + assert!(!token_clone.is_cancelled()); + + state.cancel(); + assert!(token_clone.is_cancelled()); + } + + #[test] + fn generation_state_cancel_when_empty() { + let state = GenerationState::new(); + state.cancel(); + } + + #[test] + fn generation_state_clear_does_not_cancel() { + let state = GenerationState::new(); + let token = CancellationToken::new(); + let token_clone = token.clone(); + + state.set(token); + state.clear(); + assert!(!token_clone.is_cancelled()); + } + + #[test] + fn generation_state_set_replaces_previous() { + let state = GenerationState::new(); + let first = CancellationToken::new(); + let first_clone = first.clone(); + let second = CancellationToken::new(); + let second_clone = second.clone(); + + state.set(first); + state.set(second); + + state.cancel(); + assert!(!first_clone.is_cancelled()); + assert!(second_clone.is_cancelled()); + } } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index d767e1a8..e7ef9ca8 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -481,11 +481,16 @@ pub fn run() { // ── Persistent HTTP client ──────────────────────────────── app.manage(reqwest::Client::new()); + // ── Generation cancellation state ──────────────────────────── + app.manage(commands::GenerationState::new()); + Ok(()) }) .invoke_handler(tauri::generate_handler![ #[cfg(not(coverage))] commands::ask_ollama, + #[cfg(not(coverage))] + commands::cancel_generation, notify_overlay_hidden, set_window_frame ]) diff --git a/src/App.tsx b/src/App.tsx index 555676df..6351f2d2 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -51,8 +51,15 @@ type OverlayState = 'visible' | 'hidden' | 'hiding'; function App() { const [query, setQuery] = useState(''); const [overlayState, setOverlayState] = useState('hidden'); - const { messages, streamingContent, ask, isGenerating, error, reset } = - useOllama(); + const { + messages, + streamingContent, + ask, + cancel, + isGenerating, + error, + reset, + } = useOllama(); const inputRef = useRef(null); @@ -395,6 +402,7 @@ function App() { isChatMode={isChatMode} isGenerating={isGenerating} onSubmit={handleSubmit} + onCancel={cancel} inputRef={inputRef} selectedText={selectedContext ?? undefined} /> diff --git a/src/hooks/__tests__/useOllama.test.tsx b/src/hooks/__tests__/useOllama.test.tsx index 9b5266bc..d523ceb6 100644 --- a/src/hooks/__tests__/useOllama.test.tsx +++ b/src/hooks/__tests__/useOllama.test.tsx @@ -345,6 +345,103 @@ describe('useOllama', () => { }); }); + // ─── cancel() ─────────────────────────────────────────────────────────────── + + describe('cancel()', () => { + it('invokes cancel_generation on the backend', async () => { + let resolveInvoke!: () => void; + invoke.mockImplementationOnce( + async (_cmd: string, args?: Record) => { + if (args && 'onEvent' in args) { + return new Promise((res) => { + resolveInvoke = res; + }); + } + }, + ); + + const { result } = renderHook(() => useOllama()); + + act(() => { + void result.current.ask('hello', 'hello'); + }); + + expect(result.current.isGenerating).toBe(true); + + await act(async () => { + await result.current.cancel(); + }); + + expect(invoke).toHaveBeenCalledWith('cancel_generation'); + + act(() => { + resolveInvoke?.(); + }); + }); + + it('does nothing when not generating', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.cancel(); + }); + + // cancel_generation should NOT have been called + expect(invoke).not.toHaveBeenCalledWith('cancel_generation'); + }); + }); + + // ─── Cancelled chunk handling ─────────────────────────────────────────────── + + describe('Cancelled chunk', () => { + it('finalizes partial content as assistant message on Cancelled', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.ask('hello', 'hello'); + }); + + const channel = getChannel(); + expect(channel).not.toBeNull(); + + act(() => { + channel!.simulateMessage({ type: 'Token', data: 'Partial ' }); + channel!.simulateMessage({ type: 'Token', data: 'response' }); + channel!.simulateMessage({ type: 'Cancelled' }); + }); + + expect(result.current.streamingContent).toBe(''); + expect(result.current.isGenerating).toBe(false); + expect(result.current.messages).toContainEqual( + expect.objectContaining({ + role: 'assistant', + content: 'Partial response', + }), + ); + }); + + it('does not add empty assistant message when cancelled with no tokens', async () => { + const { result } = renderHook(() => useOllama()); + + await act(async () => { + await result.current.ask('hello', 'hello'); + }); + + const channel = getChannel(); + expect(channel).not.toBeNull(); + + act(() => { + 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 + expect(result.current.messages).toHaveLength(1); + expect(result.current.messages[0].role).toBe('user'); + }); + }); + // ─── reset() ──────────────────────────────────────────────────────────────── describe('reset()', () => { diff --git a/src/hooks/useOllama.ts b/src/hooks/useOllama.ts index a7098f9a..4f2480bb 100644 --- a/src/hooks/useOllama.ts +++ b/src/hooks/useOllama.ts @@ -19,6 +19,7 @@ export interface Message { export type StreamChunk = | { type: 'Token'; data: string } | { type: 'Done' } + | { type: 'Cancelled' } | { type: 'Error'; data: string }; /** @@ -85,6 +86,21 @@ export function useOllama() { ]); setStreamingContent(''); setIsGenerating(false); + } 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, + }, + ]); + } + setStreamingContent(''); + setIsGenerating(false); } else { setError(chunk.data); setMessages((prev) => [ @@ -119,6 +135,12 @@ export function useOllama() { [isGenerating], ); + /** Cancels the currently active generation by signalling the Rust backend. */ + const cancel = useCallback(async () => { + if (!isGenerating) return; + await invoke('cancel_generation'); + }, [isGenerating]); + /** Resets all conversation state to prepare for a fresh session. */ const reset = useCallback(() => { setMessages([]); @@ -127,5 +149,13 @@ export function useOllama() { setError(null); }, []); - return { messages, streamingContent, ask, isGenerating, error, reset }; + return { + messages, + streamingContent, + ask, + cancel, + isGenerating, + error, + reset, + }; } diff --git a/src/view/AskBarView.tsx b/src/view/AskBarView.tsx index 5325e687..f4f03883 100644 --- a/src/view/AskBarView.tsx +++ b/src/view/AskBarView.tsx @@ -28,18 +28,20 @@ const ARROW_UP_ICON = ( ); /** - * Animated spinner rendered in the submit button during response generation. - * Defined as a localized component to guarantee fresh animation state on mount. + * Hoisted static SVG — square stop icon displayed during active generation. */ -function Spinner() { - return ( - - ); -} +const STOP_ICON = ( + +); /** * Props for the AskBarView component. @@ -55,6 +57,8 @@ interface AskBarViewProps { isGenerating: boolean; /** Submit handler fired when the user commits their message. */ onSubmit: () => void; + /** Cancel handler fired when the user stops an active generation. */ + onCancel: () => void; /** Ref to the textarea input element for focus management. */ inputRef: React.RefObject; /** Selected text from the host app captured at activation time, if any. */ @@ -73,6 +77,7 @@ export function AskBarView({ isChatMode, isGenerating, onSubmit, + onCancel, inputRef, selectedText, }: AskBarViewProps) { @@ -145,20 +150,20 @@ export function AskBarView({ - {isGenerating ? : ARROW_UP_ICON} + {isGenerating ? STOP_ICON : ARROW_UP_ICON} diff --git a/src/view/__tests__/AskBarView.test.tsx b/src/view/__tests__/AskBarView.test.tsx index dc28087c..74989459 100644 --- a/src/view/__tests__/AskBarView.test.tsx +++ b/src/view/__tests__/AskBarView.test.tsx @@ -16,6 +16,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -31,6 +32,7 @@ describe('AskBarView', () => { isChatMode={true} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -47,6 +49,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -63,6 +66,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={true} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -79,6 +83,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={onSubmit} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -96,6 +101,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={onSubmit} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -113,6 +119,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={onSubmit} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -128,6 +135,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -146,6 +154,7 @@ describe('AskBarView', () => { isChatMode={true} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -164,6 +173,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); @@ -180,6 +190,7 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} selectedText="some highlighted text" />, @@ -195,12 +206,64 @@ describe('AskBarView', () => { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} />, ); expect(container.querySelector('.whitespace-pre-wrap')).toBeNull(); }); + it('shows stop button with accessible label during generation', () => { + render( + , + ); + expect( + screen.getByRole('button', { name: 'Stop generating' }), + ).toBeInTheDocument(); + }); + + it('calls onCancel when stop button is clicked', () => { + const onCancel = vi.fn(); + render( + , + ); + fireEvent.click(screen.getByRole('button', { name: 'Stop generating' })); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + it('does not call onSubmit when stop button is clicked during generation', () => { + const onSubmit = vi.fn(); + render( + , + ); + fireEvent.click(screen.getByRole('button', { name: 'Stop generating' })); + expect(onSubmit).not.toHaveBeenCalled(); + }); + it('displays selectedText with whitespace-pre-wrap class', () => { const { container } = render( { isChatMode={false} isGenerating={false} onSubmit={vi.fn()} + onCancel={vi.fn()} inputRef={makeRef()} selectedText="context text here" />, From 6931b90ebd4a89ae31213e3a418418fc2021f379 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 15:27:12 -0500 Subject: [PATCH 2/8] feat: add spinning border ring animation to stop generating button Rounded-square stop button now has a CSS spinning border ring during active generation, providing clear visual feedback that inference is in progress and can be interrupted. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/App.css | 20 ++++++++++++++++++++ src/view/AskBarView.tsx | 2 +- src/view/__tests__/AskBarView.test.tsx | 16 ++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/App.css b/src/App.css index 4406c977..88f5d4c9 100644 --- a/src/App.css +++ b/src/App.css @@ -79,6 +79,26 @@ body { inset 0 1px 0 rgba(255, 255, 255, 0.04); } +/* ─── Stop Button Spinner ─── */ +.stop-btn-ring { + position: relative; +} +.stop-btn-ring::before { + content: ''; + position: absolute; + inset: -2px; + border-radius: 14px; + border: 2px solid transparent; + border-top-color: #f87171; + border-right-color: rgba(248, 113, 113, 0.3); + animation: spin-stop 1s linear infinite; +} +@keyframes spin-stop { + to { + transform: rotate(360deg); + } +} + /* ─── Chat Messages Scrollable Area ─── */ .chat-messages-scroll { scrollbar-gutter: stable; diff --git a/src/view/AskBarView.tsx b/src/view/AskBarView.tsx index f4f03883..f819c018 100644 --- a/src/view/AskBarView.tsx +++ b/src/view/AskBarView.tsx @@ -156,7 +156,7 @@ export function AskBarView({ whileTap={canSubmit || isGenerating ? { scale: 0.92 } : undefined} className={`shrink-0 w-9 h-9 rounded-xl flex items-center justify-center transition-colors duration-200 ${ isGenerating - ? 'bg-red-500/20 text-red-400 cursor-pointer' + ? 'stop-btn-ring bg-red-500/10 text-red-400 cursor-pointer' : canSubmit ? 'bg-primary text-neutral cursor-pointer' : 'bg-surface-elevated text-text-secondary cursor-default' diff --git a/src/view/__tests__/AskBarView.test.tsx b/src/view/__tests__/AskBarView.test.tsx index 74989459..b46ec6d1 100644 --- a/src/view/__tests__/AskBarView.test.tsx +++ b/src/view/__tests__/AskBarView.test.tsx @@ -247,6 +247,22 @@ describe('AskBarView', () => { expect(onCancel).toHaveBeenCalledTimes(1); }); + it('applies spinning ring class to stop button', () => { + render( + , + ); + const btn = screen.getByRole('button', { name: 'Stop generating' }); + expect(btn.classList.contains('stop-btn-ring')).toBe(true); + }); + it('does not call onSubmit when stop button is clicked during generation', () => { const onSubmit = vi.fn(); render( From 6db3722593367ce86e5e7ccb8ce366e908a92989 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 15:35:44 -0500 Subject: [PATCH 3/8] fix: smooth border-trace animation with pathLength normalization Replace broken conic-gradient and miscalculated dash-offset approach with pathLength="100" on SVG rects, making all dash math work in clean percentages. Three layered strokes (head/mid/tail) at staggered offsets with decreasing opacity create a seamless comet-tail that follows the actual rounded-rect border path with no stutter at the wrap point. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/App.css | 48 ++++++++++++++++++++++++++++--------- src/view/AskBarView.tsx | 52 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 88 insertions(+), 12 deletions(-) diff --git a/src/App.css b/src/App.css index 88f5d4c9..bf58cde4 100644 --- a/src/App.css +++ b/src/App.css @@ -79,23 +79,49 @@ body { inset 0 1px 0 rgba(255, 255, 255, 0.04); } -/* ─── Stop Button Spinner ─── */ +/* ─── Stop Button Border Trace ─── + * SVG border-trace using pathLength="100" so all dash math is in clean + * percentages. Three layered strokes at staggered offsets + decreasing + * opacity form a smooth comet-tail that follows the rounded-rect path. + * The keyframe shifts by exactly -100 (one full perimeter) for a + * seamless loop with no stutter at the wrap point. + */ .stop-btn-ring { position: relative; } -.stop-btn-ring::before { - content: ''; +.stop-btn-ring > .stop-ring-svg { position: absolute; inset: -2px; - border-radius: 14px; - border: 2px solid transparent; - border-top-color: #f87171; - border-right-color: rgba(248, 113, 113, 0.3); - animation: spin-stop 1s linear infinite; -} -@keyframes spin-stop { + width: calc(100% + 4px); + height: calc(100% + 4px); + pointer-events: none; + overflow: visible; +} +.stop-btn-ring > .stop-ring-svg rect { + fill: none; + stroke: #f87171; + stroke-width: 1.5; + stroke-linecap: round; +} +.stop-trace-head { + stroke-dasharray: 15 85; + animation: trace-border 1.8s linear infinite; +} +.stop-trace-mid { + stroke-dasharray: 12 88; + stroke-dashoffset: -15; + opacity: 0.4; + animation: trace-border 1.8s linear infinite; +} +.stop-trace-tail { + stroke-dasharray: 10 90; + stroke-dashoffset: -28; + opacity: 0.12; + animation: trace-border 1.8s linear infinite; +} +@keyframes trace-border { to { - transform: rotate(360deg); + stroke-dashoffset: -100; } } diff --git a/src/view/AskBarView.tsx b/src/view/AskBarView.tsx index f819c018..a62537bf 100644 --- a/src/view/AskBarView.tsx +++ b/src/view/AskBarView.tsx @@ -43,6 +43,49 @@ const STOP_ICON = ( ); +/** + * SVG overlay that traces a glowing comet-tail along the button's border. + * Uses `pathLength="100"` so dash math is in clean percentages regardless + * of the actual rect perimeter. Three layered strokes at staggered offsets + * create a smooth fade-out tail that follows the rounded-rect path exactly. + */ +const BORDER_TRACE_RING = ( + +); + /** * Props for the AskBarView component. */ @@ -163,7 +206,14 @@ export function AskBarView({ }`} aria-label={isGenerating ? 'Stop generating' : 'Send message'} > - {isGenerating ? STOP_ICON : ARROW_UP_ICON} + {isGenerating ? ( + <> + {BORDER_TRACE_RING} + {STOP_ICON} + + ) : ( + ARROW_UP_ICON + )} From 7957259358a95c9d457d722513d5beb0354dc370 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 15:38:35 -0500 Subject: [PATCH 4/8] fix: use animation-delay for border trace stagger instead of dashoffset The three trace layers had different initial stroke-dashoffset values, causing each to animate a different distance per cycle and drift out of sync. Replacing with animation-delay ensures all layers share the same keyframe distance (0 to -100) and move in lockstep, with the delay creating a fixed spatial offset for the comet-tail effect. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/App.css | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/App.css b/src/App.css index bf58cde4..1ab07586 100644 --- a/src/App.css +++ b/src/App.css @@ -81,10 +81,10 @@ body { /* ─── Stop Button Border Trace ─── * SVG border-trace using pathLength="100" so all dash math is in clean - * percentages. Three layered strokes at staggered offsets + decreasing - * opacity form a smooth comet-tail that follows the rounded-rect path. - * The keyframe shifts by exactly -100 (one full perimeter) for a - * seamless loop with no stutter at the wrap point. + * percentages. Three layered strokes staggered via animation-delay (NOT + * initial dashoffset — that causes different travel distances and speed + * drift). All three share the same keyframe and duration so they move + * in lockstep, with the delay creating a fixed spatial offset. */ .stop-btn-ring { position: relative; @@ -103,23 +103,29 @@ body { stroke-width: 1.5; stroke-linecap: round; } +/* Head — bright leading segment, furthest along the path */ .stop-trace-head { stroke-dasharray: 15 85; - animation: trace-border 1.8s linear infinite; + animation: trace-border 2s linear infinite; + animation-delay: -0.56s; } +/* Mid — dimmer segment trailing the head */ .stop-trace-mid { stroke-dasharray: 12 88; - stroke-dashoffset: -15; opacity: 0.4; - animation: trace-border 1.8s linear infinite; + animation: trace-border 2s linear infinite; + animation-delay: -0.3s; } +/* Tail — faintest segment, furthest behind */ .stop-trace-tail { stroke-dasharray: 10 90; - stroke-dashoffset: -28; opacity: 0.12; - animation: trace-border 1.8s linear infinite; + animation: trace-border 2s linear infinite; } @keyframes trace-border { + from { + stroke-dashoffset: 0; + } to { stroke-dashoffset: -100; } From 680618af8574079ce369850ff2abaf7fe01cc3e5 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 15:54:46 -0500 Subject: [PATCH 5/8] fix: add biased select for cancellation priority, update CLAUDE.md Add `biased;` to tokio::select! so the cancellation branch is always polled first when both futures are ready. Without this, a pre-cancelled token could non-deterministically lose to an immediately-ready stream chunk, making the pre_cancelled test flaky. Update CLAUDE.md architecture to document the new Cancelled variant. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Logan Nguyen --- CLAUDE.md | 2 +- src-tauri/src/commands.rs | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index b8fe7eaf..559ed435 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,7 +50,7 @@ Thuki is a macOS-only desktop app — a floating AI secretary activated by doubl The UI morphs between two states: a compact spotlight-style input bar → an expanded chat window. This morphing is driven by Framer Motion and a single `isChatMode` boolean in `App.tsx`. - **`App.tsx`** — orchestrates all state: messages, streaming, window resizing via ResizeObserver + Tauri `setSize()` -- **`hooks/useOllama.ts`** — Tauri Channel-based streaming hook; emits `Token`, `Done`, `Error` variants +- **`hooks/useOllama.ts`** — Tauri Channel-based streaming hook; emits `Token`, `Done`, `Cancelled`, `Error` variants - **`view/ConversationView.tsx`** — smart auto-scroll (pins to bottom unless user scrolls up) - **`view/AskBarView.tsx`** — auto-expanding textarea (max 144px), morphs logo size - **`components/ChatBubble.tsx`** — markdown rendering with DOMPurify sanitization diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 28611696..4d00bc25 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -119,6 +119,7 @@ pub async fn stream_ollama( loop { tokio::select! { + biased; _ = cancel_token.cancelled() => { // Drop the stream — closes the HTTP connection, // which signals Ollama to stop inference. From 88531e883185e9445ec99d3975711e11e1643f59 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 16:01:37 -0500 Subject: [PATCH 6/8] fix: derive Default for GenerationState to satisfy CI coverage The manual Default impl was uncovered by cargo-llvm-cov since no code path called it. Replace with #[derive(Default)] so the generated impl doesn't appear in coverage instrumentation. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Logan Nguyen --- src-tauri/src/commands.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 4d00bc25..88cb27c8 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -41,22 +41,15 @@ struct OllamaResponse { /// /// Only one generation runs at a time — starting a new request replaces the /// previous token. `cancel_generation` cancels whatever is currently active. +#[derive(Default)] pub struct GenerationState { token: Mutex>, } -impl Default for GenerationState { - fn default() -> Self { - Self::new() - } -} - impl GenerationState { /// Creates a new empty generation state with no active token. pub fn new() -> Self { - Self { - token: Mutex::new(None), - } + Self::default() } /// Stores a new cancellation token, replacing any previous one. From afad18afea070f80d39b7efa8e5695f482c39ef5 Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 16:09:38 -0500 Subject: [PATCH 7/8] fix: use direct construction in GenerationState::new() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revert new() from delegating to Self::default() — the derived Default impl is still instrumented by cargo-llvm-cov and shows as an uncovered line. Direct construction ensures new() (which is tested) covers its own body without depending on the derived impl. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Logan Nguyen --- src-tauri/src/commands.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 88cb27c8..44d2e391 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -49,7 +49,9 @@ pub struct GenerationState { impl GenerationState { /// Creates a new empty generation state with no active token. pub fn new() -> Self { - Self::default() + Self { + token: Mutex::new(None), + } } /// Stores a new cancellation token, replacing any previous one. From 44d2efafe46c3383324ea80a03b5138138e2275b Mon Sep 17 00:00:00 2001 From: Logan Nguyen Date: Wed, 1 Apr 2026 17:02:00 -0500 Subject: [PATCH 8/8] fix: ensure spawned server task exits cleanly for coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cancellation test spawned a TCP server with a 10-second sleep to hold the connection open. When the test cancelled after 100ms, the tokio runtime aborted the spawned task — its closure never ran to completion, leaving the implicit return at the closing }); uncovered. Replace the long sleep with a Notify that the test signals after assertions, so the spawned task exits cleanly and all regions are instrumented by cargo-llvm-cov. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Logan Nguyen --- src-tauri/src/commands.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index 44d2e391..0f0967e3 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -561,12 +561,18 @@ mod tests { #[tokio::test] async fn cancellation_stops_stream_and_emits_cancelled() { + use std::sync::Arc; use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); + // Notify lets the spawned server task exit cleanly after the test + // completes, ensuring its closure runs to completion for coverage. + let server_done = Arc::new(tokio::sync::Notify::new()); + let server_done_clone = server_done.clone(); + tokio::spawn(async move { let (mut stream, _) = listener.accept().await.unwrap(); let _ = stream @@ -576,8 +582,8 @@ mod tests { {\"response\":\"A\",\"done\":false}\n", ) .await; - // Keep the connection open so cancellation can interrupt - tokio::time::sleep(tokio::time::Duration::from_secs(10)).await; + // Keep the connection open until the test signals completion. + server_done_clone.notified().await; }); let client = reqwest::Client::new(); @@ -606,6 +612,10 @@ mod tests { .any(|c| matches!(c, StreamChunk::Token(t) if t == "A"))); assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Cancelled))); assert!(chunks.iter().all(|c| !matches!(c, StreamChunk::Done))); + + // Signal the server task to exit cleanly. + server_done.notify_one(); + tokio::task::yield_now().await; } #[tokio::test]