diff --git a/.env.example b/.env.example index e03d61c5..ab099562 100644 --- a/.env.example +++ b/.env.example @@ -14,3 +14,8 @@ VITE_QUOTE_MAX_DISPLAY_CHARS=300 # Maximum length of selected context text included in the Ollama prompt. VITE_QUOTE_MAX_CONTEXT_LENGTH=4096 + +# ── System Prompt ──────────────────────────────────────────────────────────── +# Custom system prompt prepended to every Ollama conversation. +# If unset or empty, defaults to: "You are Thuki, a concise desktop assistant. Keep responses short and direct." +THUKI_SYSTEM_PROMPT= diff --git a/.gitignore b/.gitignore index 27c34db7..37d161fe 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,4 @@ coverage # Git worktrees .worktrees .claude/worktrees +.gstack/ diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index dd7eecf9..47678093 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -715,6 +715,12 @@ dependencies = [ "tendril 0.5.0", ] +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + [[package]] name = "dpi" version = "0.1.2" @@ -4110,6 +4116,7 @@ version = "0.1.0" dependencies = [ "core-foundation 0.10.1", "core-graphics", + "dotenvy", "futures-util", "mockito", "reqwest", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index f04c49b2..d328a8ad 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -25,6 +25,7 @@ reqwest = { version = "0.13.2", features = ["json", "stream"] } tokio = { version = "1.50.0", features = ["macros"] } futures-util = "0.3.32" tokio-util = "0.7" +dotenvy = "0.15.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 0f0967e3..cc0e9cfb 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -1,12 +1,24 @@ +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Mutex; + 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"; +const DEFAULT_SYSTEM_PROMPT: &str = "You are Thuki (thư ký), a personal desktop secretary that \ +lives as a floating overlay on macOS. You are fast, sharp, and helpful.\n\nResponse style:\n- Be \ +concise. You appear in a small floating window — keep responses scannable.\n- Use short paragraphs, \ +bullet points, and code blocks where appropriate.\n- Lead with the answer, then explain if needed. \ +Never pad with filler.\n- Match the user's tone — casual if they're casual, precise if they're \ +technical.\n\nWhen the user provides context (quoted text from another app):\n- Treat it as the \ +subject of their question unless they say otherwise.\n- Summarize, explain, fix, or transform it as \ +asked.\n- Don't repeat the context back unless specifically helpful.\n\nYou excel at: quick answers, \ +summarizing text, explaining code, drafting messages, brainstorming ideas, and catching errors. You \ +are the user's second brain — always ready, never in the way."; /// Payload emitted back to the frontend per token chunk. #[derive(Clone, Serialize)] @@ -22,18 +34,31 @@ pub enum StreamChunk { Error(String), } -/// Request payload to Ollama server. +/// A single message in the Ollama `/api/chat` conversation format. +#[derive(Clone, Serialize, Deserialize)] +pub struct ChatMessage { + pub role: String, + pub content: String, +} + +/// Request payload for Ollama `/api/chat` endpoint. #[derive(Serialize)] -struct OllamaRequest { +struct OllamaChatRequest { model: String, - prompt: String, + messages: Vec, stream: bool, } -/// Expected structured response chunk from Ollama. +/// Nested message object in Ollama `/api/chat` response chunks. +#[derive(Deserialize)] +struct OllamaChatResponseMessage { + content: Option, +} + +/// Expected structured response chunk from Ollama `/api/chat`. #[derive(Deserialize)] -struct OllamaResponse { - response: Option, +struct OllamaChatResponse { + message: Option, done: Option, } @@ -72,27 +97,64 @@ impl GenerationState { } } -/// 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( +/// Backend-managed conversation history with an epoch counter to prevent +/// stale writes after a reset. The Rust side is the source of truth; the +/// frontend sends only new user messages and receives streamed tokens. +pub struct ConversationHistory { + messages: Mutex>, + epoch: AtomicU64, +} + +impl Default for ConversationHistory { + fn default() -> Self { + Self { + messages: Mutex::new(Vec::new()), + epoch: AtomicU64::new(0), + } + } +} + +impl ConversationHistory { + /// Creates a new empty conversation history at epoch 0. + pub fn new() -> Self { + Self::default() + } +} + +/// System prompt loaded once at startup from the `THUKI_SYSTEM_PROMPT` +/// environment variable, falling back to a built-in default. +pub struct SystemPrompt(pub String); + +/// Reads `THUKI_SYSTEM_PROMPT` from the environment, falling back to the +/// built-in default when unset or empty. +pub fn load_system_prompt() -> String { + std::env::var("THUKI_SYSTEM_PROMPT") + .ok() + .filter(|s| !s.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_SYSTEM_PROMPT.to_string()) +} + +/// Core streaming logic for Ollama `/api/chat`, separated from the Tauri +/// command for testability. 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. +/// Returns the accumulated assistant response so the caller can persist it. +pub async fn stream_ollama_chat( endpoint: &str, model: &str, - prompt: String, + messages: Vec, client: &reqwest::Client, cancel_token: CancellationToken, on_chunk: impl Fn(StreamChunk), -) { - let request_payload = OllamaRequest { +) -> String { + let request_payload = OllamaChatRequest { model: model.to_string(), - prompt, + messages, stream: true, }; + let mut accumulated = String::new(); + let res = client.post(endpoint).json(&request_payload).send().await; match res { @@ -106,7 +168,7 @@ pub async fn stream_ollama( format!("HTTP {}: {}", status, err_body) }; on_chunk(StreamChunk::Error(err_msg)); - return; + return accumulated; } let mut stream = response.bytes_stream(); @@ -120,7 +182,7 @@ pub async fn stream_ollama( // which signals Ollama to stop inference. drop(stream); on_chunk(StreamChunk::Cancelled); - return; + return accumulated; } chunk_opt = stream.next() => { match chunk_opt { @@ -136,11 +198,14 @@ pub async fn stream_ollama( } if let Ok(json) = - serde_json::from_str::(trimmed) + serde_json::from_str::(trimmed) { - if let Some(token) = json.response { - if !token.is_empty() { - on_chunk(StreamChunk::Token(token)); + if let Some(msg) = json.message { + if let Some(token) = msg.content { + if !token.is_empty() { + accumulated.push_str(&token); + on_chunk(StreamChunk::Token(token)); + } } } if let Some(true) = json.done { @@ -152,9 +217,9 @@ pub async fn stream_ollama( } Some(Err(e)) => { on_chunk(StreamChunk::Error(e.to_string())); - return; + return accumulated; } - None => return, + None => return accumulated, } } } @@ -164,36 +229,80 @@ pub async fn stream_ollama( on_chunk(StreamChunk::Error(e.to_string())); } } + + accumulated } -/// Streams text chunks from the local Ollama backend via `reqwest` to the frontend using `Channel`. -/// -/// Creates a fresh `CancellationToken` for each request and stores it in `GenerationState`, -/// allowing `cancel_generation` to abort the stream at any time. +/// Streams a chat response from the local Ollama backend. Appends the user +/// message and assistant response to conversation history only after successful +/// completion. Uses an epoch counter to prevent stale writes after a reset. #[cfg_attr(coverage_nightly, coverage(off))] #[cfg_attr(not(coverage), tauri::command)] pub async fn ask_ollama( - prompt: String, + message: String, + quoted_text: Option, on_event: Channel, client: State<'_, reqwest::Client>, generation: State<'_, GenerationState>, + history: State<'_, ConversationHistory>, + system_prompt: State<'_, SystemPrompt>, ) -> Result<(), String> { - let endpoint = format!("{}/api/generate", DEFAULT_OLLAMA_URL.trim_end_matches('/')); + let endpoint = format!("{}/api/chat", DEFAULT_OLLAMA_URL.trim_end_matches('/')); let cancel_token = CancellationToken::new(); generation.set(cancel_token.clone()); - stream_ollama( + // Build user message content, prepending quoted context when present. + let content = match quoted_text { + Some(ref qt) if !qt.trim().is_empty() => format!("Context: \"{}\"\n\n{}", qt, message), + _ => message, + }; + + let user_msg = ChatMessage { + role: "user".to_string(), + content, + }; + + // Snapshot the current epoch and build the messages array for Ollama. + // The user message is NOT yet committed to history — it is only added + // after a successful response to prevent orphaned messages on errors. + let (epoch_at_start, messages) = { + let conv = history.messages.lock().unwrap(); + let epoch = history.epoch.load(Ordering::SeqCst); + let mut msgs = vec![ChatMessage { + role: "system".to_string(), + content: system_prompt.0.clone(), + }]; + msgs.extend(conv.clone()); + msgs.push(user_msg.clone()); + (epoch, msgs) + }; + + let accumulated = stream_ollama_chat( &endpoint, DEFAULT_MODEL_NAME, - prompt, + messages, &client, - cancel_token, + cancel_token.clone(), |chunk| { let _ = on_event.send(chunk); }, ) .await; + // Persist user + assistant messages only when the stream completed + // naturally (not cancelled or error), the epoch has not changed + // (no reset during streaming), and we received content. + let was_cancelled = cancel_token.is_cancelled(); + let current_epoch = history.epoch.load(Ordering::SeqCst); + if !was_cancelled && current_epoch == epoch_at_start && !accumulated.is_empty() { + let mut conv = history.messages.lock().unwrap(); + conv.push(user_msg); + conv.push(ChatMessage { + role: "assistant".to_string(), + content: accumulated, + }); + } + generation.clear(); Ok(()) } @@ -201,7 +310,7 @@ pub async fn ask_ollama( /// 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. +/// `stream_ollama_chat` loop to exit immediately and drop the HTTP connection. #[cfg_attr(coverage_nightly, coverage(off))] #[cfg_attr(not(coverage), tauri::command)] pub async fn cancel_generation(generation: State<'_, GenerationState>) -> Result<(), String> { @@ -209,6 +318,16 @@ pub async fn cancel_generation(generation: State<'_, GenerationState>) -> Result Ok(()) } +/// Clears the backend conversation history and increments the epoch counter. +/// The epoch increment prevents any in-flight `ask_ollama` from writing stale +/// messages into the freshly cleared history. +#[cfg_attr(coverage_nightly, coverage(off))] +#[cfg_attr(not(coverage), tauri::command)] +pub fn reset_conversation(history: State<'_, ConversationHistory>) { + history.epoch.fetch_add(1, Ordering::SeqCst); + history.messages.lock().unwrap().clear(); +} + #[cfg(test)] mod tests { use super::*; @@ -223,27 +342,41 @@ mod tests { (chunks, callback) } + /// Helper: builds a `/api/chat` response line from content + done flag. + fn chat_line(content: &str, done: bool) -> String { + format!( + "{{\"message\":{{\"role\":\"assistant\",\"content\":\"{}\"}},\"done\":{}}}\n", + content, done + ) + } + #[tokio::test] async fn streams_tokens_from_valid_response() { let mut server = mockito::Server::new_async().await; + let body = format!( + "{}{}{}", + chat_line("Hello", false), + chat_line(" world", false), + chat_line("", true), + ); let mock = server - .mock("POST", "/api/generate") - .with_body( - "{\"response\":\"Hello\",\"done\":false}\n\ - {\"response\":\" world\",\"done\":false}\n\ - {\"response\":\"\",\"done\":true}\n", - ) + .mock("POST", "/api/chat") + .with_body(body) .create_async() .await; let client = reqwest::Client::new(); let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); + let messages = vec![ChatMessage { + role: "user".to_string(), + content: "hi".to_string(), + }]; - stream_ollama( - &format!("{}/api/generate", server.url()), + let accumulated = stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + messages, &client, token, callback, @@ -255,13 +388,14 @@ mod tests { assert!(matches!(&chunks[0], StreamChunk::Token(t) if t == "Hello")); assert!(matches!(&chunks[1], StreamChunk::Token(t) if t == " world")); assert!(matches!(&chunks[2], StreamChunk::Done)); + assert_eq!(accumulated, "Hello world"); } #[tokio::test] async fn handles_http_500() { let mut server = mockito::Server::new_async().await; let mock = server - .mock("POST", "/api/generate") + .mock("POST", "/api/chat") .with_status(500) .with_body("Internal Server Error") .create_async() @@ -271,10 +405,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + let accumulated = stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -285,6 +419,7 @@ mod tests { let chunks = chunks.lock().unwrap(); assert_eq!(chunks.len(), 1); assert!(matches!(&chunks[0], StreamChunk::Error(e) if e.contains("500"))); + assert!(accumulated.is_empty()); } #[tokio::test] @@ -293,10 +428,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - "http://127.0.0.1:1/api/generate", + let accumulated = stream_ollama_chat( + "http://127.0.0.1:1/api/chat", "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -306,14 +441,16 @@ mod tests { let chunks = chunks.lock().unwrap(); assert_eq!(chunks.len(), 1); assert!(matches!(&chunks[0], StreamChunk::Error(_))); + assert!(accumulated.is_empty()); } #[tokio::test] async fn handles_malformed_json() { let mut server = mockito::Server::new_async().await; + let body = format!("not json at all\n{}", chat_line("ok", true)); let mock = server - .mock("POST", "/api/generate") - .with_body("not json at all\n{\"response\":\"ok\",\"done\":true}\n") + .mock("POST", "/api/chat") + .with_body(body) .create_async() .await; @@ -321,10 +458,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -340,7 +477,7 @@ mod tests { async fn handles_empty_response_body() { let mut server = mockito::Server::new_async().await; let mock = server - .mock("POST", "/api/generate") + .mock("POST", "/api/chat") .with_body("") .create_async() .await; @@ -349,10 +486,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + let accumulated = stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -362,19 +499,22 @@ mod tests { mock.assert_async().await; let chunks = chunks.lock().unwrap(); assert!(chunks.is_empty()); + assert!(accumulated.is_empty()); } #[tokio::test] async fn tokens_arrive_in_order() { let mut server = mockito::Server::new_async().await; + let body = format!( + "{}{}{}{}", + chat_line("A", false), + chat_line("B", false), + chat_line("C", false), + chat_line("", true), + ); let mock = server - .mock("POST", "/api/generate") - .with_body( - "{\"response\":\"A\",\"done\":false}\n\ - {\"response\":\"B\",\"done\":false}\n\ - {\"response\":\"C\",\"done\":false}\n\ - {\"response\":\"\",\"done\":true}\n", - ) + .mock("POST", "/api/chat") + .with_body(body) .create_async() .await; @@ -382,10 +522,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + let accumulated = stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -402,14 +542,16 @@ mod tests { }) .collect(); assert_eq!(tokens, vec!["A", "B", "C"]); + assert_eq!(accumulated, "ABC"); } #[tokio::test] async fn handles_invalid_utf8_in_stream() { let mut server = mockito::Server::new_async().await; - let body = b"\xFF\xFE\n{\"response\":\"ok\",\"done\":true}\n".to_vec(); + let mut body = b"\xFF\xFE\n".to_vec(); + body.extend_from_slice(chat_line("ok", true).as_bytes()); let mock = server - .mock("POST", "/api/generate") + .mock("POST", "/api/chat") .with_body(body) .create_async() .await; @@ -418,10 +560,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -457,10 +599,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("http://127.0.0.1:{}/api/generate", port), + stream_ollama_chat( + &format!("http://127.0.0.1:{}/api/chat", port), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -476,7 +618,7 @@ mod tests { async fn http_500_with_empty_body() { let mut server = mockito::Server::new_async().await; let mock = server - .mock("POST", "/api/generate") + .mock("POST", "/api/chat") .with_status(500) .with_body("") .create_async() @@ -486,10 +628,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -505,9 +647,10 @@ mod tests { #[tokio::test] async fn whitespace_only_lines_are_skipped() { let mut server = mockito::Server::new_async().await; + let body = format!(" \n{}", chat_line("hi", true)); let mock = server - .mock("POST", "/api/generate") - .with_body(" \n{\"response\":\"hi\",\"done\":true}\n") + .mock("POST", "/api/chat") + .with_body(body) .create_async() .await; @@ -515,10 +658,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -531,10 +674,10 @@ mod tests { } #[tokio::test] - async fn response_field_absent_emits_only_done() { + async fn message_field_absent_emits_only_done() { let mut server = mockito::Server::new_async().await; let mock = server - .mock("POST", "/api/generate") + .mock("POST", "/api/chat") .with_body("{\"done\":true}\n") .create_async() .await; @@ -543,10 +686,10 @@ mod tests { let token = CancellationToken::new(); let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -568,21 +711,17 @@ mod tests { 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 - .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 until the test signals completion. + let first_line = chat_line("A", false); + let header = format!( + "HTTP/1.1 200 OK\r\nContent-Type: application/x-ndjson\r\n\r\n{}", + first_line + ); + let _ = stream.write_all(header.as_bytes()).await; server_done_clone.notified().await; }); @@ -596,10 +735,10 @@ mod tests { token_clone.cancel(); }); - stream_ollama( - &format!("http://127.0.0.1:{}/api/generate", port), + stream_ollama_chat( + &format!("http://127.0.0.1:{}/api/chat", port), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -613,7 +752,6 @@ mod tests { 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; } @@ -622,11 +760,8 @@ mod tests { 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", - ) + .mock("POST", "/api/chat") + .with_body(chat_line("Hello", true)) .create_async() .await; @@ -636,10 +771,10 @@ mod tests { let (chunks, callback) = collect_chunks(); - stream_ollama( - &format!("{}/api/generate", server.url()), + stream_ollama_chat( + &format!("{}/api/chat", server.url()), "test-model", - "hi".to_string(), + vec![], &client, token, callback, @@ -650,6 +785,74 @@ mod tests { assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Cancelled))); } + #[tokio::test] + async fn sends_messages_array_in_request() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/api/chat") + .match_body(mockito::Matcher::PartialJsonString( + r#"{"messages":[{"role":"system","content":"Be helpful"},{"role":"user","content":"hi"}]}"#.to_string(), + )) + .with_body(chat_line("", true)) + .create_async() + .await; + + let client = reqwest::Client::new(); + let token = CancellationToken::new(); + let (_, callback) = collect_chunks(); + let messages = vec![ + ChatMessage { + role: "system".to_string(), + content: "Be helpful".to_string(), + }, + ChatMessage { + role: "user".to_string(), + content: "hi".to_string(), + }, + ]; + + stream_ollama_chat( + &format!("{}/api/chat", server.url()), + "test-model", + messages, + &client, + token, + callback, + ) + .await; + + mock.assert_async().await; + } + + #[tokio::test] + async fn message_content_absent_emits_only_done() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("POST", "/api/chat") + .with_body("{\"message\":{\"role\":\"assistant\"},\"done\":true}\n") + .create_async() + .await; + + let client = reqwest::Client::new(); + let token = CancellationToken::new(); + let (chunks, callback) = collect_chunks(); + + stream_ollama_chat( + &format!("{}/api/chat", server.url()), + "test-model", + vec![], + &client, + token, + callback, + ) + .await; + + mock.assert_async().await; + let chunks = chunks.lock().unwrap(); + assert!(chunks.iter().all(|c| !matches!(c, StreamChunk::Token(_)))); + assert!(chunks.iter().any(|c| matches!(c, StreamChunk::Done))); + } + #[test] fn generation_state_set_and_cancel() { let state = GenerationState::new(); @@ -695,4 +898,54 @@ mod tests { assert!(!first_clone.is_cancelled()); assert!(second_clone.is_cancelled()); } + + #[test] + fn load_system_prompt_returns_default_when_unset() { + std::env::remove_var("THUKI_SYSTEM_PROMPT"); + + let prompt = load_system_prompt(); + assert_eq!(prompt, DEFAULT_SYSTEM_PROMPT); + } + + #[test] + fn load_system_prompt_reads_env_var() { + std::env::set_var("THUKI_SYSTEM_PROMPT", "Custom prompt"); + + let prompt = load_system_prompt(); + assert_eq!(prompt, "Custom prompt"); + + std::env::remove_var("THUKI_SYSTEM_PROMPT"); + } + + #[test] + fn load_system_prompt_ignores_empty_env_var() { + std::env::set_var("THUKI_SYSTEM_PROMPT", " "); + + let prompt = load_system_prompt(); + assert_eq!(prompt, DEFAULT_SYSTEM_PROMPT); + + std::env::remove_var("THUKI_SYSTEM_PROMPT"); + } + + #[test] + fn conversation_history_new_starts_at_epoch_zero() { + let h = ConversationHistory::new(); + assert_eq!(h.epoch.load(Ordering::SeqCst), 0); + assert!(h.messages.lock().unwrap().is_empty()); + } + + #[test] + fn conversation_history_epoch_increments_on_clear() { + let h = ConversationHistory::new(); + h.messages.lock().unwrap().push(ChatMessage { + role: "user".to_string(), + content: "hi".to_string(), + }); + + h.epoch.fetch_add(1, Ordering::SeqCst); + h.messages.lock().unwrap().clear(); + + assert_eq!(h.epoch.load(Ordering::SeqCst), 1); + assert!(h.messages.lock().unwrap().is_empty()); + } } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e7ef9ca8..1f665b1a 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -407,6 +407,10 @@ fn init_panel(app_handle: &tauri::AppHandle) { /// Panics if the Tauri runtime fails to initialise. #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { + // Load .env file so THUKI_SYSTEM_PROMPT and future backend env vars + // work the same way as Vite's VITE_* vars for the frontend. + dotenvy::dotenv().ok(); + let mut builder = tauri::Builder::default(); #[cfg(target_os = "macos")] @@ -481,8 +485,10 @@ pub fn run() { // ── Persistent HTTP client ──────────────────────────────── app.manage(reqwest::Client::new()); - // ── Generation cancellation state ──────────────────────────── + // ── Generation + conversation state ───────────────────── app.manage(commands::GenerationState::new()); + app.manage(commands::ConversationHistory::new()); + app.manage(commands::SystemPrompt(commands::load_system_prompt())); Ok(()) }) @@ -491,6 +497,8 @@ pub fn run() { commands::ask_ollama, #[cfg(not(coverage))] commands::cancel_generation, + #[cfg(not(coverage))] + commands::reset_conversation, notify_overlay_hidden, set_window_frame ]) diff --git a/src/App.tsx b/src/App.tsx index 38f461a3..307629df 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -222,10 +222,7 @@ function App() { ?.replace(CONTROL_CHARS, '') .slice(0, quote.maxContextLength); const hasContext = sanitized && sanitized.trim().length > 0; - const ollamaPrompt = hasContext - ? `Context: "${sanitized}"\n\n${query}` - : query; - ask(query, ollamaPrompt, hasContext ? sanitized : undefined); + ask(query, hasContext ? sanitized : undefined); setSelectedContext(null); setQuery(''); if (inputRef.current) { diff --git a/src/__tests__/App.test.tsx b/src/__tests__/App.test.tsx index 8b3c778b..837553dd 100644 --- a/src/__tests__/App.test.tsx +++ b/src/__tests__/App.test.tsx @@ -281,7 +281,7 @@ describe('App', () => { expect(__mockWindow.startDragging).not.toHaveBeenCalled(); }); - it('submits query with context prepended when selectedContext is set', async () => { + it('submits query with quoted text when selectedContext is set', async () => { render(); await act(async () => {}); @@ -299,11 +299,12 @@ describe('App', () => { await act(async () => {}); - // Prompt should be context-wrapped + // Backend receives the message and quoted text separately expect(invoke).toHaveBeenCalledWith( 'ask_ollama', expect.objectContaining({ - prompt: expect.stringContaining('Context:'), + message: 'my question', + quotedText: 'selected snippet', }), ); }); diff --git a/src/hooks/__tests__/useOllama.test.tsx b/src/hooks/__tests__/useOllama.test.tsx index d523ceb6..923d1524 100644 --- a/src/hooks/__tests__/useOllama.test.tsx +++ b/src/hooks/__tests__/useOllama.test.tsx @@ -24,16 +24,19 @@ describe('useOllama', () => { // ─── ask() ────────────────────────────────────────────────────────────────── describe('ask()', () => { - it('sends prompt via invoke with correct command name and args', async () => { + it('sends message via invoke with correct command name and args', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello world', 'hello world'); + await result.current.ask('hello world'); }); expect(invoke).toHaveBeenCalledWith( 'ask_ollama', - expect.objectContaining({ prompt: 'hello world' }), + expect.objectContaining({ + message: 'hello world', + quotedText: null, + }), ); }); @@ -56,7 +59,7 @@ describe('useOllama', () => { // Start ask but don't await so we can read state while in-flight act(() => { - void result.current.ask('test prompt', 'test prompt'); + void result.current.ask('test prompt'); }); // isGenerating should be true right after ask sets it @@ -72,7 +75,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('my question', 'my question'); + await result.current.ask('my question'); }); expect(result.current.messages[0]).toEqual( @@ -88,11 +91,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask( - 'what is this?', - 'Context: "code snippet"\n\nwhat is this?', - 'code snippet', - ); + await result.current.ask('what is this?', 'code snippet'); }); expect(result.current.messages[0]).toEqual( @@ -104,20 +103,18 @@ describe('useOllama', () => { ); }); - it('sends ollamaPrompt (not displayContent) to invoke', async () => { + it('sends quotedText to invoke when provided', async () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask( - 'summarize', - 'Context: "selected text"\n\nsummarize', - ); + await result.current.ask('summarize', 'selected text'); }); expect(invoke).toHaveBeenCalledWith( 'ask_ollama', expect.objectContaining({ - prompt: 'Context: "selected text"\n\nsummarize', + message: 'summarize', + quotedText: 'selected text', }), ); }); @@ -126,7 +123,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -144,7 +141,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -169,7 +166,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('', ''); + await result.current.ask(''); }); expect(invoke).not.toHaveBeenCalled(); @@ -180,7 +177,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask(' ', ' '); + await result.current.ask(' '); }); expect(invoke).not.toHaveBeenCalled(); @@ -203,7 +200,7 @@ describe('useOllama', () => { // Start the first ask (stalls) act(() => { - void result.current.ask('first', 'first'); + void result.current.ask('first'); }); expect(result.current.isGenerating).toBe(true); @@ -211,7 +208,7 @@ describe('useOllama', () => { // Try a second ask while generating await act(async () => { - await result.current.ask('second', 'second'); + await result.current.ask('second'); }); // invoke should NOT have been called again @@ -231,7 +228,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('test', 'test'); + await result.current.ask('test'); }); const channel = getChannel(); @@ -251,7 +248,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('test', 'test'); + await result.current.ask('test'); }); expect(result.current.error).toBe('Error: connection refused'); @@ -262,7 +259,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('test', 'test'); + await result.current.ask('test'); }); const channel = getChannel(); @@ -287,7 +284,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('test', 'test'); + await result.current.ask('test'); }); const assistantMsg = result.current.messages.find( @@ -302,14 +299,14 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('first ask', 'first ask'); + await result.current.ask('first ask'); }); expect(result.current.error).toBe('Error: first error'); // Second ask — succeeds, channel sends Done await act(async () => { - await result.current.ask('second ask', 'second ask'); + await result.current.ask('second ask'); }); const channel = getChannel(); @@ -330,7 +327,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -363,7 +360,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); act(() => { - void result.current.ask('hello', 'hello'); + void result.current.ask('hello'); }); expect(result.current.isGenerating).toBe(true); @@ -398,7 +395,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -424,7 +421,7 @@ describe('useOllama', () => { const { result } = renderHook(() => useOllama()); await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -450,7 +447,7 @@ describe('useOllama', () => { // Build up some state await act(async () => { - await result.current.ask('hello', 'hello'); + await result.current.ask('hello'); }); const channel = getChannel(); @@ -469,6 +466,8 @@ describe('useOllama', () => { expect(result.current.streamingContent).toBe(''); expect(result.current.isGenerating).toBe(false); expect(result.current.error).toBeNull(); + // Should also reset backend conversation history + expect(invoke).toHaveBeenCalledWith('reset_conversation'); }); }); @@ -480,7 +479,7 @@ describe('useOllama', () => { // First ask + response await act(async () => { - await result.current.ask('first question', 'first question'); + await result.current.ask('first question'); }); let channel = getChannel(); act(() => { @@ -494,7 +493,7 @@ describe('useOllama', () => { // Second ask + response await act(async () => { - await result.current.ask('second question', 'second question'); + await result.current.ask('second question'); }); channel = getChannel(); act(() => { diff --git a/src/hooks/useOllama.ts b/src/hooks/useOllama.ts index 4f2480bb..f304844b 100644 --- a/src/hooks/useOllama.ts +++ b/src/hooks/useOllama.ts @@ -35,22 +35,17 @@ export function useOllama() { const [error, setError] = useState(null); /** - * Submits a prompt to the Ollama backend and initiates the streaming response. - * Modifies local state optimistically. + * 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. * * @param displayContent The user's query as it should appear in the chat bubble. - * @param ollamaPrompt The full prompt sent to Ollama (may include context preamble). * @param quotedText Optional selected text quoted alongside this message. */ const ask = useCallback( - async ( - displayContent: string, - ollamaPrompt: string, - quotedText?: string, - ) => { + async (displayContent: string, quotedText?: string) => { if (!displayContent.trim() || isGenerating) return; setMessages((prev) => [ @@ -117,7 +112,11 @@ export function useOllama() { }; try { - await invoke('ask_ollama', { prompt: ollamaPrompt, onEvent: channel }); + await invoke('ask_ollama', { + message: displayContent, + quotedText: quotedText ?? null, + onEvent: channel, + }); } catch (err) { setError(String(err)); setMessages((prev) => [ @@ -147,6 +146,7 @@ export function useOllama() { setStreamingContent(''); setIsGenerating(false); setError(null); + void invoke('reset_conversation'); }, []); return {