From c000ae41e84dc8f7fdfe3613b1cf03b8cc0c6df6 Mon Sep 17 00:00:00 2001 From: donfaquir <1458918806@qq.com> Date: Wed, 24 Jun 2026 11:58:50 +0800 Subject: [PATCH 1/6] perf(llm): share HTTP client and cache embedding vectors Spec-1.0: replace per-request reqwest::Client with a shared Arc registered as Tauri managed state, eliminating redundant TLS init and connection-pool setup on every LLM call. Spec-2.0: add EmbeddingCache (RwLock + generation-based invalidation) so load_all_doc/thought_embeddings runs once then serves from memory, cutting 100-500ms of SQLite full-table-scan IO per chat message. --- src-tauri/src/challenge_review.rs | 8 +- src-tauri/src/lib.rs | 12 ++- src-tauri/src/link_recommendation.rs | 22 +++-- src-tauri/src/llm/agent_loop.rs | 5 +- src-tauri/src/llm/mod.rs | 19 ++++- src-tauri/src/llm/provider.rs | 19 ++++- src-tauri/src/llm/provider_impl.rs | 48 ++++++----- src-tauri/src/passive_highlight.rs | 5 +- src-tauri/src/semantic_index.rs | 86 ++++++++++++++++++-- src-tauri/src/skills/commands.rs | 5 +- src-tauri/src/skills/skill_tool.rs | 3 +- src-tauri/src/tools/built_in/graph_ops.rs | 14 +++- src-tauri/src/tools/built_in/link_ops.rs | 12 ++- src-tauri/src/tools/built_in/vault_search.rs | 11 ++- src-tauri/src/tools/commands.rs | 3 + src-tauri/src/tools/context.rs | 3 + src-tauri/src/topic_network.rs | 16 ++-- src-tauri/src/writing_coach.rs | 5 +- 18 files changed, 235 insertions(+), 61 deletions(-) diff --git a/src-tauri/src/challenge_review.rs b/src-tauri/src/challenge_review.rs index 20e2536..ea69839 100644 --- a/src-tauri/src/challenge_review.rs +++ b/src-tauri/src/challenge_review.rs @@ -7,6 +7,8 @@ use std::fs; use std::path::{Path, PathBuf}; use uuid::Uuid; +use std::sync::Arc; + use crate::llm::{create_provider, CompletionOverrides}; use crate::llm::LlmChatMessage; use crate::thought_parser; @@ -272,6 +274,7 @@ fn normalize_template_kind(raw: Option<&str>) -> String { #[tauri::command] pub async fn generate_challenge_question( workspace: tauri::State<'_, crate::WorkspaceState>, + http_client: tauri::State<'_, Arc>, args: GenerateChallengeQuestionArgs, ) -> Result { let root = crate::lock_workspace_root(&workspace)?; @@ -282,7 +285,7 @@ pub async fn generate_challenge_question( .await .map_err(|e| e.to_string())??; - let provider = match create_provider(&ai, None) { + let provider = match create_provider(&ai, None, http_client.inner()) { Ok(p) => p, Err(_) => { return Ok(GenerateChallengeQuestionResponse { @@ -407,6 +410,7 @@ pub async fn generate_challenge_question( #[tauri::command] pub async fn evaluate_challenge_answer( workspace: tauri::State<'_, crate::WorkspaceState>, + http_client: tauri::State<'_, Arc>, args: EvaluateChallengeAnswerArgs, ) -> Result { let root = crate::lock_workspace_root(&workspace)?; @@ -424,7 +428,7 @@ pub async fn evaluate_challenge_answer( } }; - let provider = match create_provider(&ai, None) { + let provider = match create_provider(&ai, None, http_client.inner()) { Ok(p) => p, Err(_) => { return Ok(EvaluateChallengeAnswerResponse { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index b9e3b70..50d2363 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1540,11 +1540,13 @@ async fn semantic_search( args: semantic_index::SemanticSearchArgs, app: tauri::AppHandle, state: tauri::State<'_, WorkspaceState>, + embed_cache: tauri::State<'_, Arc>, ) -> Result, String> { let root = lock_workspace_root(&state)?; let cache = semantic_index::default_model_cache_dir(); let bundle = semantic_index::resolve_bundle_model_dir(&app); - tauri::async_runtime::spawn_blocking(move || semantic_index::run_semantic_search(&root, &cache, &bundle, args)) + let embed_cache_arc = Arc::clone(embed_cache.inner()); + tauri::async_runtime::spawn_blocking(move || semantic_index::run_semantic_search(&root, &cache, &bundle, args, &embed_cache_arc)) .await .map_err(|e| e.to_string())? } @@ -1557,6 +1559,8 @@ async fn suggest_related_notes( include_reasons: Option, editor_markdown_override: Option, state: tauri::State<'_, WorkspaceState>, + http_client: tauri::State<'_, Arc>, + embed_cache: tauri::State<'_, Arc>, ) -> Result, String> { let canonical_root = lock_workspace_root(&state)?; let rel_path = rel_path.trim().to_string(); @@ -1569,6 +1573,7 @@ async fn suggest_related_notes( let root = canonical_root.clone(); let rel = rel_path.clone(); let md_override = editor_markdown_override.clone(); + let embed_cache_arc = Arc::clone(embed_cache.inner()); let mut out = tauri::async_runtime::spawn_blocking(move || -> Result, String> { let emb = semantic_index::open_embedding_db(&root)?; let thoughts = vault_thoughts_db::open_thoughts_db(&root)?; @@ -1579,6 +1584,7 @@ async fn suggest_related_notes( &thoughts, max_results, md_override.as_deref(), + &embed_cache_arc, ) }) .await @@ -1593,7 +1599,7 @@ async fn suggest_related_notes( }) .await .map_err(|e| e.to_string())??; - link_recommendation::enrich_recommendations_with_reasons(&mut out, &excerpt, &ai).await?; + link_recommendation::enrich_recommendations_with_reasons(&mut out, &excerpt, &ai, http_client.inner()).await?; } Ok(out) @@ -1655,6 +1661,8 @@ async fn add_manual_topic_semantic( pub fn run() { tauri::Builder::default() .manage(WorkspaceState::default()) + .manage(Arc::new(llm::build_shared_http_client())) + .manage(Arc::new(semantic_index::EmbeddingCache::new())) .manage(Arc::new(llm::LlmSessionState::default())) .manage(Arc::new(llm::approval::ToolApprovalState::new())) .manage(Arc::new(tools::ToolRegistry::new())) diff --git a/src-tauri/src/link_recommendation.rs b/src-tauri/src/link_recommendation.rs index 4705edb..1448a18 100644 --- a/src-tauri/src/link_recommendation.rs +++ b/src-tauri/src/link_recommendation.rs @@ -1,6 +1,8 @@ //! 文档级语义相似度 → 双向链接推荐候选（迭代 6.3 步骤 15）。 //! 无语义索引时不提供推荐（无关键词兜底）。 +use std::sync::Arc; + use crate::llm::create_provider; use crate::llm::LlmChatMessage; use crate::note_privacy; @@ -435,6 +437,7 @@ pub fn suggest_related_notes( _thoughts_db: &Connection, max_results: usize, editor_markdown_override: Option<&str>, + embed_cache: &semantic_index::EmbeddingCache, ) -> Result, String> { if max_results == 0 { return Ok(Vec::new()); @@ -456,7 +459,7 @@ pub fn suggest_related_notes( return Err("semantic_index_not_ready: kf-private notes are excluded from link recommendations".to_string()); } - let all_chunks = semantic_index::load_all_doc_embeddings(embedding_db)?; + let all_chunks = embed_cache.get_docs(embedding_db); if all_chunks.is_empty() { return Err( "semantic_index_not_ready: no document chunks in embedding index; rebuild embeddings first" @@ -623,8 +626,9 @@ async fn link_reason_completion_body( ai: &ResolvedAiConfig, messages: &[LlmChatMessage], timeout_ms: u64, + http_client: &Arc, ) -> Option { - let provider = create_provider(ai, None).ok()?; + let provider = create_provider(ai, None, http_client).ok()?; let overrides = crate::llm::CompletionOverrides { timeout_ms: Some(timeout_ms), ..Default::default() @@ -639,6 +643,7 @@ pub async fn enrich_recommendations_with_reasons( candidates: &mut [LinkRecommendation], current_doc_excerpt: &str, config: &ResolvedAiConfig, + http_client: &Arc, ) -> Result<(), String> { if candidates.is_empty() { return Ok(()); @@ -680,7 +685,7 @@ pub async fn enrich_recommendations_with_reasons( }, ]; - let raw = match link_reason_completion_body(config, &msgs, timeout_ms).await { + let raw = match link_reason_completion_body(config, &msgs, timeout_ms, http_client).await { Some(s) => { let t = s.trim(); if t.is_empty() { @@ -813,7 +818,8 @@ mod tests { semantic_index::upsert_doc_chunk(&conn, "far.md#0", "far.md", 0, "t", &far, "m").unwrap(); let tconn = Connection::open_in_memory().unwrap(); - let rec = suggest_related_notes(root, "cur.md", &conn, &tconn, 5, None).unwrap(); + let ec = semantic_index::EmbeddingCache::new(); + let rec = suggest_related_notes(root, "cur.md", &conn, &tconn, 5, None, &ec).unwrap(); assert_eq!(rec.len(), 2); assert_eq!(rec[0].target_rel_path, "near.md"); assert_eq!(rec[1].target_rel_path, "mid.md"); @@ -837,7 +843,8 @@ mod tests { semantic_index::upsert_doc_chunk(&conn, "linked#0", "linked", 0, "t", &near, "m").unwrap(); let tconn = Connection::open_in_memory().unwrap(); - let rec = suggest_related_notes(root, "cur.md", &conn, &tconn, 5, None).expect("chunks align"); + let ec = semantic_index::EmbeddingCache::new(); + let rec = suggest_related_notes(root, "cur.md", &conn, &tconn, 5, None, &ec).expect("chunks align"); assert_eq!(rec.len(), 1); assert_eq!(rec[0].target_rel_path, "near.md"); } @@ -862,6 +869,7 @@ mod tests { semantic_index::upsert_doc_chunk(&conn, "far.md#0", "far.md", 0, "t", &far, "m").unwrap(); let tconn = Connection::open_in_memory().unwrap(); + let ec = semantic_index::EmbeddingCache::new(); let rec = suggest_related_notes( root, "cur.md", @@ -869,6 +877,7 @@ mod tests { &tconn, 5, Some("x [[linked]]\n"), + &ec, ) .unwrap(); assert_eq!(rec.len(), 1); @@ -931,7 +940,8 @@ mod tests { existing_link: false, reason: None, }]; - super::enrich_recommendations_with_reasons(&mut c, "excerpt", &ai) + let http_client = Arc::new(reqwest::Client::new()); + super::enrich_recommendations_with_reasons(&mut c, "excerpt", &ai, &http_client) .await .unwrap(); assert!(c[0].reason.is_none()); diff --git a/src-tauri/src/llm/agent_loop.rs b/src-tauri/src/llm/agent_loop.rs index e9e6000..72c4034 100644 --- a/src-tauri/src/llm/agent_loop.rs +++ b/src-tauri/src/llm/agent_loop.rs @@ -11,7 +11,7 @@ use std::time::Duration; use futures_util::future::join_all; use serde_json::{json, Value}; -use tauri::{AppHandle, Emitter}; +use tauri::{AppHandle, Emitter, Manager}; use tokio_util::sync::CancellationToken; use super::approval::ToolApprovalState; @@ -484,6 +484,9 @@ pub(crate) async fn execute_tool( ); ctx.call_id = Some(tc.id.clone()); ctx.provider = provider; + if let Some(ec) = app.try_state::>() { + ctx.embed_cache = Some(Arc::clone(&*ec)); + } let manifest = tool.manifest().clone(); let start = std::time::Instant::now(); diff --git a/src-tauri/src/llm/mod.rs b/src-tauri/src/llm/mod.rs index b103440..a656143 100644 --- a/src-tauri/src/llm/mod.rs +++ b/src-tauri/src/llm/mod.rs @@ -8,7 +8,10 @@ pub(crate) mod provider; pub(crate) mod provider_impl; pub mod memory; -pub use provider::{create_provider, create_provider_by_id, CompletionOverrides, LlmProvider}; +pub use provider::{ + build_shared_http_client, create_provider, create_provider_by_id, CompletionOverrides, + LlmProvider, +}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum AgentMode { @@ -413,6 +416,7 @@ fn assemble_messages( args: &ChatStreamStartArgs, embed_cache_bundle: Option<(PathBuf, PathBuf)>, auto_invocable_skills: &[(String, String, Option)], + embed_cache: &semantic_index::EmbeddingCache, ) -> Result { let mut out: Vec = Vec::new(); let mut reply_context_sources = ReplyContextSources::default(); @@ -557,6 +561,7 @@ fn assemble_messages( &sem_cfg, &kw_paths, &omit_semantic_docs, + embed_cache, ) { reply_context_sources.semantic = ReplySemanticSource { injected: true, @@ -689,6 +694,7 @@ pub struct ListModelsArgs { #[tauri::command] pub async fn list_models( state: State<'_, crate::WorkspaceState>, + http_client: State<'_, Arc>, args: ListModelsArgs, ) -> Result, String> { let root = lock_workspace_root(&state)?; @@ -709,6 +715,7 @@ pub async fn list_models( }; let provider = provider_impl::UnifiedProvider::new( + Arc::clone(http_client.inner()), base, key, String::new(), @@ -730,6 +737,8 @@ pub async fn start_chat_stream( ctx_factory: State<'_, Arc>, approval: State<'_, Arc>, skills: State<'_, Arc>, + http_client: State<'_, Arc>, + embed_cache_state: State<'_, Arc>, args: ChatStreamStartArgs, ) -> Result { let root = lock_workspace_root(&workspace)?; @@ -757,7 +766,8 @@ pub async fn start_chat_stream( &active_profile.default_model, )) .unwrap_or_default(); - let provider = create_provider(&ai, model_override.map(|s| s))?; + let http_client_arc = Arc::clone(http_client.inner()); + let provider = create_provider(&ai, model_override.map(|s| s), &http_client_arc)?; let cache = semantic_index::default_model_cache_dir(); let bundle = semantic_index::resolve_bundle_model_dir(&app); @@ -773,7 +783,8 @@ pub async fn start_chat_stream( } else { Vec::new() }; - let outcome = assemble_messages(&root, &ai, &args, embed_paths, &skills_for_prompt)?; + let embed_cache_arc = Arc::clone(embed_cache_state.inner()); + let outcome = assemble_messages(&root, &ai, &args, embed_paths, &skills_for_prompt, &embed_cache_arc)?; let mut messages = outcome.messages; let resolved_depth = outcome.resolved_depth; let reply_context_sources = outcome.reply_context_sources; @@ -821,7 +832,7 @@ pub async fn start_chat_stream( tokio::spawn(async move { let memory_manager: agent_loop::SharedMemoryManager = if memory_enabled { - let extraction_provider = provider::create_provider(&ai_for_memory, None).ok(); + let extraction_provider = provider::create_provider(&ai_for_memory, None, &http_client_arc).ok(); let mgr = memory::MemoryManager::new(workspace_root.clone(), extraction_provider); if let Some(mem_msg) = mgr.format_for_injection() { let pos = if messages.is_empty() { 0 } else { 1 }; diff --git a/src-tauri/src/llm/provider.rs b/src-tauri/src/llm/provider.rs index a766633..854b8ef 100644 --- a/src-tauri/src/llm/provider.rs +++ b/src-tauri/src/llm/provider.rs @@ -78,15 +78,27 @@ pub fn resolve_model_name(last_used: Option<&str>, default_model: &str) -> Optio }) } +/// Build a default shared HTTP client for LLM providers. +pub fn build_shared_http_client() -> reqwest::Client { + reqwest::Client::builder() + .pool_max_idle_per_host(4) + .pool_idle_timeout(std::time::Duration::from_secs(90)) + .connect_timeout(std::time::Duration::from_secs(15)) + .use_rustls_tls() + .build() + .expect("Failed to create shared HTTP client") +} + /// Create a provider from the active profile in the config. pub fn create_provider( config: &AiConfig, model_override: Option<&str>, + http_client: &Arc, ) -> Result, String> { let profile = config .active_profile() .ok_or("No active provider configured. Choose a provider in settings.")?; - create_provider_from_profile(profile, config, model_override) + create_provider_from_profile(profile, config, model_override, http_client) } /// Create a provider for a specific profile identified by `provider_id`. @@ -94,19 +106,21 @@ pub fn create_provider_by_id( config: &AiConfig, provider_id: &str, model_override: Option<&str>, + http_client: &Arc, ) -> Result, String> { let profile = config .providers .iter() .find(|p| p.id == provider_id) .ok_or_else(|| format!("Provider '{}' not found in config.", provider_id))?; - create_provider_from_profile(profile, config, model_override) + create_provider_from_profile(profile, config, model_override, http_client) } fn create_provider_from_profile( profile: &ProviderProfile, config: &AiConfig, model_override: Option<&str>, + http_client: &Arc, ) -> Result, String> { let model = model_override .map(str::trim) @@ -117,6 +131,7 @@ fn create_provider_from_profile( Ok(Arc::new( super::provider_impl::UnifiedProvider::new( + Arc::clone(http_client), profile.base_url.clone(), profile.api_key.clone(), model, diff --git a/src-tauri/src/llm/provider_impl.rs b/src-tauri/src/llm/provider_impl.rs index 5dec24f..00a2584 100644 --- a/src-tauri/src/llm/provider_impl.rs +++ b/src-tauri/src/llm/provider_impl.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use async_trait::async_trait; use futures_util::StreamExt; use serde::Deserialize; @@ -9,6 +11,7 @@ use super::provider::{ChatStreamResult, CompletionOverrides, LlmProvider, Normal use super::{emit_chunk, emit_done, emit_error, LlmChatMessage}; pub struct UnifiedProvider { + client: Arc, base_url: String, api_key: String, model: String, @@ -21,6 +24,7 @@ pub struct UnifiedProvider { impl UnifiedProvider { pub fn new( + client: Arc, base_url: String, api_key: String, model: String, @@ -31,6 +35,7 @@ impl UnifiedProvider { is_remote: bool, ) -> Self { Self { + client, base_url, api_key, model, @@ -42,15 +47,6 @@ impl UnifiedProvider { } } - fn http_client(&self) -> Result { - reqwest::Client::builder() - .timeout(std::time::Duration::from_millis(self.timeout_ms)) - .connect_timeout(std::time::Duration::from_secs(15)) - .use_rustls_tls() - .build() - .map_err(|e| format!("Failed to create HTTP client: {e}")) - } - fn build_auth_headers( &self, builder: reqwest::RequestBuilder, @@ -200,7 +196,6 @@ impl LlmProvider for UnifiedProvider { tools: Option>, cancel: CancellationToken, ) -> Result { - let client = self.http_client()?; let url = format!( "{}/chat/completions", self.base_url.trim_end_matches('/') @@ -223,7 +218,12 @@ impl LlmProvider for UnifiedProvider { } } - let req = self.build_auth_headers(client.post(&url)).json(&body); + let req = self.build_auth_headers( + self.client + .post(&url) + .timeout(std::time::Duration::from_millis(self.timeout_ms)), + ) + .json(&body); let resp = match req.send().await { Ok(r) => r, Err(e) => { @@ -388,13 +388,6 @@ impl LlmProvider for UnifiedProvider { .and_then(|o| o.timeout_ms) .unwrap_or(self.timeout_ms); - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_millis(timeout.max(3000).min(45_000))) - .connect_timeout(std::time::Duration::from_secs(15)) - .use_rustls_tls() - .build() - .map_err(|e| format!("Failed to create HTTP client: {e}"))?; - let url = format!( "{}/chat/completions", self.base_url.trim_end_matches('/') @@ -413,7 +406,12 @@ impl LlmProvider for UnifiedProvider { body["response_format"] = json!({"type": "json_object"}); } - let req = self.build_auth_headers(client.post(&url)).json(&body); + let req = self.build_auth_headers( + self.client + .post(&url) + .timeout(std::time::Duration::from_millis(timeout.max(3000).min(45_000))), + ) + .json(&body); let resp = req .send() .await @@ -442,9 +440,12 @@ impl LlmProvider for UnifiedProvider { } async fn list_models(&self) -> Result, String> { - let client = self.http_client()?; let url = format!("{}/models", self.base_url.trim_end_matches('/')); - let req = self.build_auth_headers(client.get(&url)); + let req = self.build_auth_headers( + self.client + .get(&url) + .timeout(std::time::Duration::from_millis(self.timeout_ms)), + ); let resp = req .send() .await @@ -508,9 +509,14 @@ impl LlmProvider for UnifiedProvider { mod tests { use super::*; + fn test_client() -> Arc { + Arc::new(reqwest::Client::new()) + } + #[test] fn serialize_tool_result_message() { let provider = UnifiedProvider::new( + test_client(), "https://api.openai.com/v1".to_string(), "test-key".to_string(), "gpt-4o".to_string(), diff --git a/src-tauri/src/passive_highlight.rs b/src-tauri/src/passive_highlight.rs index 3802195..91c5f25 100644 --- a/src-tauri/src/passive_highlight.rs +++ b/src-tauri/src/passive_highlight.rs @@ -1,5 +1,7 @@ //! Passive highlight: sidecar value detection on user messages (non-streaming JSON). +use std::sync::Arc; + use crate::llm::create_provider; use crate::llm::LlmChatMessage; use crate::lock_workspace_root; @@ -98,6 +100,7 @@ fn normalize_kind(k: &str) -> Option<&'static str> { #[tauri::command] pub async fn detect_passive_highlight( workspace: State<'_, crate::WorkspaceState>, + http_client: State<'_, Arc>, args: DetectPassiveHighlightArgs, ) -> Result { let root = lock_workspace_root(&workspace)?; @@ -122,7 +125,7 @@ pub async fn detect_passive_highlight( return Ok(empty()); } - let provider = match create_provider(&ai, None) { + let provider = match create_provider(&ai, None, http_client.inner()) { Ok(p) => p, Err(_) => return Ok(empty()), }; diff --git a/src-tauri/src/semantic_index.rs b/src-tauri/src/semantic_index.rs index 64fe8b7..6b4ed0f 100644 --- a/src-tauri/src/semantic_index.rs +++ b/src-tauri/src/semantic_index.rs @@ -13,7 +13,8 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::io::Read; use std::path::{Path, PathBuf}; -use std::sync::{Arc, Mutex}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; use tauri::{AppHandle, Emitter, Manager}; use uuid::Uuid; @@ -170,6 +171,68 @@ pub struct ThoughtEmbeddingRow { pub model_id: String, } +pub struct EmbeddingCache { + docs: RwLock>)>>, + thoughts: RwLock>)>>, + generation: AtomicU64, +} + +impl EmbeddingCache { + pub fn new() -> Self { + Self { + docs: RwLock::new(None), + thoughts: RwLock::new(None), + generation: AtomicU64::new(0), + } + } + + pub fn invalidate(&self) { + self.generation.fetch_add(1, Ordering::Relaxed); + } + + pub fn get_docs(&self, conn: &Connection) -> Arc> { + let current_gen = self.generation.load(Ordering::Relaxed); + { + let guard = self.docs.read().unwrap(); + if let Some((cached_gen, ref data)) = *guard { + if cached_gen == current_gen { + return Arc::clone(data); + } + } + } + let mut guard = self.docs.write().unwrap(); + if let Some((cached_gen, ref data)) = *guard { + if cached_gen == current_gen { + return Arc::clone(data); + } + } + let rows = Arc::new(load_all_doc_embeddings(conn).unwrap_or_default()); + *guard = Some((current_gen, Arc::clone(&rows))); + rows + } + + pub fn get_thoughts(&self, conn: &Connection) -> Arc> { + let current_gen = self.generation.load(Ordering::Relaxed); + { + let guard = self.thoughts.read().unwrap(); + if let Some((cached_gen, ref data)) = *guard { + if cached_gen == current_gen { + return Arc::clone(data); + } + } + } + let mut guard = self.thoughts.write().unwrap(); + if let Some((cached_gen, ref data)) = *guard { + if cached_gen == current_gen { + return Arc::clone(data); + } + } + let rows = Arc::new(load_all_thought_embeddings(conn).unwrap_or_default()); + *guard = Some((current_gen, Arc::clone(&rows))); + rows + } +} + pub fn upsert_doc_chunk( conn: &Connection, chunk_id: &str, @@ -1147,6 +1210,10 @@ fn rebuild_index_impl(vault_root: &Path, app: &AppHandle, resume: bool) -> Resul rebuild_progress::write_rebuild_progress(vault_root, &rp)?; rebuild_progress::emit_checkpoint(app, &rp); + if let Some(ec) = app.try_state::>() { + ec.invalidate(); + } + semantic_rebuild_log(&format!( "rebuild_index: ok indexed_chunks={indexed_chunks} indexed_thoughts={indexed_thoughts} elapsed_ms={elapsed_ms}" )); @@ -1171,6 +1238,7 @@ pub fn run_semantic_search( cache_dir: &Path, bundle_dir: &Path, args: SemanticSearchArgs, + embed_cache: &EmbeddingCache, ) -> Result, String> { let model = get_cached_or_load_model(cache_dir, bundle_dir)?; let conn = open_embedding_db(vault_root)?; @@ -1184,11 +1252,11 @@ pub fn run_semantic_search( let exclude = HashSet::new(); let mut out = Vec::new(); if scope == "docs" || scope == "all" { - let rows = load_all_doc_embeddings(&conn)?; + let rows = embed_cache.get_docs(&conn); out.extend(semantic_search_docs(&qv, &rows, args.top_k, &exclude)); } if scope == "thoughts" || scope == "all" { - let thoughts = load_all_thought_embeddings(&conn)?; + let thoughts = embed_cache.get_thoughts(&conn); let mut h = semantic_search_thoughts(&qv, &thoughts, args.top_k); let tconn = vault_thoughts_db::open_thoughts_db(vault_root)?; for hit in &mut h { @@ -1267,6 +1335,7 @@ pub fn build_semantic_context_for_llm( cfg: &SemanticConfig, keyword_snippet_paths: &[String], omit_doc_rel_paths: &[String], + embed_cache: &EmbeddingCache, ) -> Option { if !cfg.enabled { return None; @@ -1284,8 +1353,8 @@ pub fn build_semantic_context_for_llm( let conn = open_embedding_db(vault_root).ok()?; let tconn = vault_thoughts_db::open_thoughts_db(vault_root).ok()?; let qv = encode_single(model.as_ref(), q).ok()?; - let doc_rows = load_all_doc_embeddings(&conn).ok()?; - let thought_rows = load_all_thought_embeddings(&conn).ok()?; + let doc_rows = embed_cache.get_docs(&conn); + let thought_rows = embed_cache.get_thoughts(&conn); let sem_docs = semantic_search_docs(&qv, &doc_rows, 12, &omit_set); let mut sem_thoughts = semantic_search_thoughts(&qv, &thought_rows, 12); for hit in &mut sem_thoughts { @@ -1462,7 +1531,12 @@ pub fn incremental_reindex_note(vault_root: &Path, app: &AppHandle, rel_path: &s tx.commit().map_err(|e| e.to_string())?; Ok(()) })(); - if let Err(e) = res { + if let Err(e) = &res { eprintln!("[semantic_index] incremental reindex skipped: {e}"); } + if res.is_ok() { + if let Some(ec) = app.try_state::>() { + ec.invalidate(); + } + } } diff --git a/src-tauri/src/skills/commands.rs b/src-tauri/src/skills/commands.rs index 5423eee..8c7ce81 100644 --- a/src-tauri/src/skills/commands.rs +++ b/src-tauri/src/skills/commands.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; -use tauri::{AppHandle, State}; +use tauri::{AppHandle, Manager, State}; use tokio_util::sync::CancellationToken; use crate::llm::approval::ToolApprovalState; @@ -326,7 +326,8 @@ pub async fn invoke_skill( .as_deref() .map(str::trim) .filter(|s| !s.is_empty()); - let provider = create_provider(&ai, model_override)?; + let http_client: Arc = Arc::clone(&*app.state::>()); + let provider = create_provider(&ai, model_override, &http_client)?; let workspace_name = root .file_name() diff --git a/src-tauri/src/skills/skill_tool.rs b/src-tauri/src/skills/skill_tool.rs index 11d657b..5314949 100644 --- a/src-tauri/src/skills/skill_tool.rs +++ b/src-tauri/src/skills/skill_tool.rs @@ -199,7 +199,8 @@ impl Tool for SkillAsTool { return tool_err(ToolErrorCode::Internal, &format!("config join: {e}")); } }; - let provider = match create_provider(&ai, None) { + let http_client: std::sync::Arc = std::sync::Arc::clone(&*self.app.state::>()); + let provider = match create_provider(&ai, None, &http_client) { Ok(p) => p, Err(e) => { drop(permit); diff --git a/src-tauri/src/tools/built_in/graph_ops.rs b/src-tauri/src/tools/built_in/graph_ops.rs index a7c2ee6..69343b2 100644 --- a/src-tauri/src/tools/built_in/graph_ops.rs +++ b/src-tauri/src/tools/built_in/graph_ops.rs @@ -166,13 +166,21 @@ impl Tool for IndexStatusTool { let root = ctx.workspace_root.clone(); let cache_dir_opt = ctx.app_cache_dir.clone(); let bundle_dir_opt = ctx.app_bundle_resource_dir.clone(); + let embed_cache = ctx.embed_cache.clone(); let result = tauri::async_runtime::spawn_blocking( move || -> Result<(usize, usize, bool), String> { let conn = crate::semantic_index::open_embedding_db(&root)?; - let doc_chunks = crate::semantic_index::load_all_doc_embeddings(&conn)?; - let thought_embeddings = - crate::semantic_index::load_all_thought_embeddings(&conn)?; + let fallback_cache; + let cache_ref = match embed_cache.as_ref() { + Some(c) => c.as_ref(), + None => { + fallback_cache = crate::semantic_index::EmbeddingCache::new(); + &fallback_cache + } + }; + let doc_chunks = cache_ref.get_docs(&conn); + let thought_embeddings = cache_ref.get_thoughts(&conn); let doc_chunk_count = doc_chunks.len(); let thought_embedding_count = thought_embeddings.len(); diff --git a/src-tauri/src/tools/built_in/link_ops.rs b/src-tauri/src/tools/built_in/link_ops.rs index 2ebbc61..37ee9a6 100644 --- a/src-tauri/src/tools/built_in/link_ops.rs +++ b/src-tauri/src/tools/built_in/link_ops.rs @@ -96,23 +96,30 @@ impl Tool for LinkSuggestRelatedTool { let root = ctx.workspace_root.clone(); let rel = rel_path.clone(); + let embed_cache = ctx.embed_cache.clone(); let result = tauri::async_runtime::spawn_blocking( move || -> Result, String> { let full_path = root.join(&rel); - // 检查文件是否存在 if !full_path.exists() { return Err("__NOT_FOUND__".to_string()); } - // 检查是否私密 if crate::note_privacy::peek_kf_private_from_md_file(&full_path) { return Err("__PRIVACY_BLOCKED__".to_string()); } let emb_conn = crate::semantic_index::open_embedding_db(&root)?; let thoughts_conn = crate::vault_thoughts_db::open_thoughts_db(&root)?; + let fallback_cache; + let cache_ref = match embed_cache.as_ref() { + Some(c) => c.as_ref(), + None => { + fallback_cache = crate::semantic_index::EmbeddingCache::new(); + &fallback_cache + } + }; crate::link_recommendation::suggest_related_notes( &root, &rel, @@ -120,6 +127,7 @@ impl Tool for LinkSuggestRelatedTool { &thoughts_conn, max_results, None, + cache_ref, ) }, ) diff --git a/src-tauri/src/tools/built_in/vault_search.rs b/src-tauri/src/tools/built_in/vault_search.rs index 580457a..d901348 100644 --- a/src-tauri/src/tools/built_in/vault_search.rs +++ b/src-tauri/src/tools/built_in/vault_search.rs @@ -293,6 +293,7 @@ impl Tool for VaultSemanticSearchTool { let root = ctx.workspace_root.clone(); let privacy_filter = Arc::clone(&ctx.privacy_filter); let workspace_root_for_filter = root.clone(); + let embed_cache = ctx.embed_cache.clone(); let args = SemanticSearchArgs { query, @@ -301,7 +302,15 @@ impl Tool for VaultSemanticSearchTool { }; let result = tauri::async_runtime::spawn_blocking(move || { - crate::semantic_index::run_semantic_search(&root, &cache_dir, &bundle_dir, args) + let fallback_cache; + let cache_ref = match embed_cache.as_ref() { + Some(c) => c.as_ref(), + None => { + fallback_cache = crate::semantic_index::EmbeddingCache::new(); + &fallback_cache + } + }; + crate::semantic_index::run_semantic_search(&root, &cache_dir, &bundle_dir, args, cache_ref) }) .await; diff --git a/src-tauri/src/tools/commands.rs b/src-tauri/src/tools/commands.rs index e88d6b8..593dcbb 100644 --- a/src-tauri/src/tools/commands.rs +++ b/src-tauri/src/tools/commands.rs @@ -2,6 +2,7 @@ use serde_json::Value; use std::sync::Arc; use tauri::State; +use crate::semantic_index::EmbeddingCache; use crate::tools::context::ToolContextFactory; use crate::tools::registry::{ToolRegistry, ToolScope}; use crate::WorkspaceState; @@ -23,6 +24,7 @@ pub async fn invoke_tool( registry: State<'_, Arc>, ctx_factory: State<'_, Arc>, ws_state: State<'_, WorkspaceState>, + embed_cache_state: State<'_, Arc>, app: tauri::AppHandle, ) -> Result { let workspace_root = crate::lock_workspace_root(&ws_state)?; @@ -43,6 +45,7 @@ pub async fn invoke_tool( Some(bundle_dir), ); ctx.call_id = Some(uuid::Uuid::now_v7().to_string()); + ctx.embed_cache = Some(Arc::clone(embed_cache_state.inner())); let manifest = tool.manifest().clone(); let start = std::time::Instant::now(); diff --git a/src-tauri/src/tools/context.rs b/src-tauri/src/tools/context.rs index 2d348a2..e3ee2cf 100644 --- a/src-tauri/src/tools/context.rs +++ b/src-tauri/src/tools/context.rs @@ -5,6 +5,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use crate::llm::provider::LlmProvider; +use crate::semantic_index::EmbeddingCache; // ─── ToolContext ─────────────────────────────────────────────────────────────── @@ -24,6 +25,7 @@ pub struct ToolContext { /// (or another tool that recursed into agent_loop). Stage 1 caps this at 1. pub nesting_depth: u8, pub provider: Option>, + pub embed_cache: Option>, } // ─── AuditSink trait ─────────────────────────────────────────────────────────── @@ -107,6 +109,7 @@ impl ToolContextFactory { app_bundle_resource_dir, nesting_depth, provider: None, + embed_cache: None, } } } diff --git a/src-tauri/src/topic_network.rs b/src-tauri/src/topic_network.rs index d94b521..095245f 100644 --- a/src-tauri/src/topic_network.rs +++ b/src-tauri/src/topic_network.rs @@ -1,5 +1,7 @@ //! 主题网络（迭代 6.4）：LLM 提取主题、SQLite 缓存、二部图构建与 Markdown 导出快照。 +use std::sync::Arc; + use crate::llm::{create_provider, CompletionOverrides}; use crate::llm::LlmChatMessage; use crate::note_privacy; @@ -18,8 +20,7 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::io::Read; use std::path::{Path, PathBuf}; -use tauri::AppHandle; -use tauri::Emitter; +use tauri::{AppHandle, Emitter, Manager}; const MAX_FILES: usize = 600; const READ_CAP: usize = 512 * 1024; @@ -346,8 +347,9 @@ pub async fn extract_topics_for_document( doc_content_for_prompt: &str, existing_topics_list: &[String], ai: &AiConfig, + http_client: &Arc, ) -> Result, String> { - let provider = create_provider(ai, None)?; + let provider = create_provider(ai, None, http_client)?; let list_block = if existing_topics_list.is_empty() { "(none)".to_string() } else { @@ -683,7 +685,8 @@ pub fn add_manual_topic_semantic_blocking( let query_vec = crate::builtin_embed::encode_single(&emb_model, display)?; let emb_conn = semantic_index::open_embedding_db(vault_root)?; - let chunks = semantic_index::load_all_doc_embeddings(&emb_conn)?; + let embed_cache: Arc = Arc::clone(&*app.state::>()); + let chunks = embed_cache.get_docs(&emb_conn); if chunks.is_empty() { return Err("语义索引中尚无文档向量，请先在设置中重建嵌入索引后再新增主题。".to_string()); } @@ -905,7 +908,8 @@ struct TopicExtractProgressPayload { /// 全量：扫描 vault、增量提取、构图 pub async fn build_topic_network(vault_root: &Path, app: &AppHandle) -> Result { let ai = crate::vault_config::load_ai_config_internal(vault_root)?; - let llm_available = create_provider(&ai, None).is_ok(); + let http_client: Arc = Arc::clone(&*app.state::>()); + let llm_available = create_provider(&ai, None, &http_client).is_ok(); let topic_conn = open_topic_db(vault_root)?; let mut paths: Vec = Vec::new(); @@ -968,7 +972,7 @@ pub async fn build_topic_network(vault_root: &Path, app: &AppHandle) -> Result t, Err(_) => continue, }; diff --git a/src-tauri/src/writing_coach.rs b/src-tauri/src/writing_coach.rs index d38ed2c..506a755 100644 --- a/src-tauri/src/writing_coach.rs +++ b/src-tauri/src/writing_coach.rs @@ -1,5 +1,7 @@ //! Writing coach: paragraph-level argumentation check + vault keyword linkage (JSON mode). +use std::sync::Arc; + use crate::llm::create_provider; use crate::llm::LlmChatMessage; use crate::lock_workspace_root; @@ -367,6 +369,7 @@ fn filter_response( #[tauri::command] pub async fn analyze_writing_coach( workspace: State<'_, crate::WorkspaceState>, + http_client: State<'_, Arc>, args: AnalyzeWritingCoachArgs, ) -> Result { let root = lock_workspace_root(&workspace)?; @@ -385,7 +388,7 @@ pub async fn analyze_writing_coach( let ai = vault_config::load_ai_config_internal(&root) .map_err(|e| e.to_string())?; - let provider = create_provider(&ai, None)?; + let provider = create_provider(&ai, None, http_client.inner())?; let msgs = vec![ LlmChatMessage { From 965d85191b16233a63e24084c082f8973c58e74d Mon Sep 17 00:00:00 2001 From: donfaquir <1458918806@qq.com> Date: Wed, 24 Jun 2026 14:04:14 +0800 Subject: [PATCH 2/6] perf(llm): move vault and semantic search out of IPC handler --- src-tauri/src/llm/mod.rs | 422 ++++++++++++------ src-tauri/src/vault_context_search.rs | 2 +- src/components/AiConversationPanel.tsx | 133 +++--- .../ThoughtMgmtAiConversationPanel.tsx | 115 +++-- src/hooks/useThoughtMgmtAiConversations.ts | 3 - src/hooks/useWorkspaceAiConversations.ts | 4 - src/types/vaultContextSearch.ts | 2 +- 7 files changed, 385 insertions(+), 296 deletions(-) diff --git a/src-tauri/src/llm/mod.rs b/src-tauri/src/llm/mod.rs index a656143..f1d8faf 100644 --- a/src-tauri/src/llm/mod.rs +++ b/src-tauri/src/llm/mod.rs @@ -41,7 +41,6 @@ use chrono::Utc; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::{HashMap, HashSet}; -use std::path::Path; use std::sync::{Arc, Mutex}; use tauri::{AppHandle, Emitter, State}; use tokio_util::sync::CancellationToken; @@ -120,12 +119,6 @@ pub struct NoteContextIn { pub markdown_for_gate: String, } -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct VaultContextIn { - pub snippets: Vec, -} - /// 想法聚焦对话：由前端与会话持久化传入（迭代 6.1） #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] @@ -144,9 +137,8 @@ pub struct ChatStreamStartArgs { pub model: Option, #[serde(default)] pub note_context: Option, - /// 由 `search_workspace_context` 结果回传；出站前在 `assemble_messages` 内按磁盘重算摘录。 #[serde(default)] - pub vault_context: Option, + pub include_vault_context: Option, /// 当前深度模式（迭代 3）；影响系统提示的详细程度与风格。 #[serde(default)] pub depth_mode: Option, @@ -264,14 +256,33 @@ fn resolve_auto_depth_heuristic(query: &str) -> (DepthMode, String) { (DepthMode::Deep, "long_query".to_string()) } -struct AssembleOutcome { +struct AssembleFastOutcome { messages: Vec, - /// `Some` 仅当本次请求显式为 Auto：解析结果供前端展示与决策日志。 + context_insert_pos: usize, resolved_depth: Option, auto_resolve_reason: Option, reply_context_sources: ReplyContextSources, } +#[derive(Default)] +struct ContextBlocks { + vault_block: Option, + vault_snippets: Vec, + vault_meta: Option, + vault_reply_sources: Option, + semantic_block: Option, + semantic_reply_sources: ReplySemanticSource, +} + +#[derive(Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct ContextReadyPayload { + session_id: String, + snippets: Vec, + meta: Option, + reply_context_sources: ReplyContextSources, +} + // --- 事件载荷（与前端 listen 对齐） --- #[derive(Clone, Serialize)] @@ -409,15 +420,11 @@ fn build_skills_system_block( Some(s) } -/// 合并 `note_context`、可选 `vault_context` 与对话轮次，并校验角色。 -fn assemble_messages( - canonical_root: &Path, +fn assemble_messages_fast( ai: &vault_config::AiConfig, args: &ChatStreamStartArgs, - embed_cache_bundle: Option<(PathBuf, PathBuf)>, auto_invocable_skills: &[(String, String, Option)], - embed_cache: &semantic_index::EmbeddingCache, -) -> Result { +) -> Result { let mut out: Vec = Vec::new(); let mut reply_context_sources = ReplyContextSources::default(); @@ -459,6 +466,8 @@ fn assemble_messages( } } + let context_insert_pos = out.len(); + let last_user_query: Option = args .messages .iter() @@ -480,106 +489,6 @@ fn assemble_messages( Some(d) => Some(d), }; - // Vault 摘录：预算与 `maxContextTokens` 粗挂钩（字符≈4×token），当前笔记 system 优先占满后再给摘录。 - if let (Some(vc), Some(q)) = (&args.vault_context, last_user_query.as_deref()) { - if !vc.snippets.is_empty() { - let rebuilt = vault_context_search::rebuild_vault_snippets_for_llm( - canonical_root, - ai, - &vc.snippets, - q, - 1200, - 96 * 1024, - ); - if !rebuilt.is_empty() { - let default_total: usize = 32_000; - let total_budget = ai - .request - .max_context_tokens - .map(|m| (m as usize).saturating_mul(4)) - .unwrap_or(default_total) - .clamp(4_000, 100_000); - let used_note: usize = out.first().map(|m| m.content.chars().count()).unwrap_or(0); - let used_msgs: usize = args.messages.iter().map(|m| m.content.chars().count() + 8).sum(); - let vault_cap = total_budget - .saturating_sub(used_note) - .saturating_sub(used_msgs) - .min(12_000) - .max(400); - if let Some((vault_block, truncated, used_rel_paths)) = - vault_context_search::build_vault_context_system_block(&rebuilt, vault_cap) - { - let used_set: HashSet<&str> = used_rel_paths.iter().map(String::as_str).collect(); - let entries: Vec = rebuilt - .iter() - .filter(|s| used_set.contains(s.rel_path.as_str())) - .map(|s| ReplyVaultKeywordEntry { - rel_path: s.rel_path.clone(), - kind: s.kind.clone(), - }) - .collect(); - reply_context_sources.vault_keyword = Some(ReplyVaultKeywordSource { - entries, - truncated, - }); - out.push(LlmChatMessage { - role: "system".to_string(), - content: vault_block, - ..Default::default() - }); - } else { - reply_context_sources.vault_keyword = Some(ReplyVaultKeywordSource { - entries: Vec::new(), - truncated: true, - }); - } - } - } - } - - let sem_ctx_on = args.semantic_context_enabled.unwrap_or(true); - if sem_ctx_on { - if let Ok(sem_cfg) = vault_config::load_semantic_merged(canonical_root) { - if sem_cfg.enabled { - if let (Some(q), Some(paths)) = (last_user_query.as_deref(), embed_cache_bundle.as_ref()) { - let (cache, bundle) = paths; - let kw_paths: Vec = args - .vault_context - .as_ref() - .map(|vc| vc.snippets.iter().map(|s| s.rel_path.clone()).collect()) - .unwrap_or_default(); - let omit_semantic_docs: Vec = args - .note_context - .as_ref() - .map(|n| vec![n.rel_path.clone()]) - .unwrap_or_default(); - if let Some(sem_res) = semantic_index::build_semantic_context_for_llm( - canonical_root, - cache, - bundle, - q, - &sem_cfg, - &kw_paths, - &omit_semantic_docs, - embed_cache, - ) { - reply_context_sources.semantic = ReplySemanticSource { - injected: true, - document_paths: sem_res.used.document_paths.clone(), - thought_ids: sem_res.used.thought_ids.clone(), - }; - out.push(LlmChatMessage { - role: "system".to_string(), - content: sem_res.block, - ..Default::default() - }); - } - } - } - } - } - - // 深度模式系统指令（迭代 3）；Auto 在入模前已解析为浅/中/深 if let Some(depth) = depth_for_prompt { let depth_instruction = build_depth_system_instruction(depth); out.push(LlmChatMessage { @@ -589,7 +498,6 @@ fn assemble_messages( }); } - // 深化子轮次上下文（迭代 3 Phase 4） if let Some(invite) = &args.invite_context { let mut ctx = format!( "The user accepted a deepening question: \"{}\". ", @@ -613,7 +521,6 @@ fn assemble_messages( }); } - // 语言匹配指令：确保 LLM 使用用户相同的语言回复 out.push(LlmChatMessage { role: "system".to_string(), content: "IMPORTANT: Always respond in the same language the user writes in. \ @@ -624,14 +531,12 @@ fn assemble_messages( ..Default::default() }); - // 叠词 / 重复词抑制：无笔记上下文时仅靠本条生效；与深度指令互补 out.push(LlmChatMessage { role: "system".to_string(), content: CHAT_ANTI_REPETITION_SYSTEM.to_string(), ..Default::default() }); - // Iter 3.5 P0-2：开启工具调用时,明确告诉 LLM "先发现后读",避免按训练直觉假设文件在根目录。 let tools_enabled_eff = args.tools_enabled.unwrap_or(ai.tools_enabled); if tools_enabled_eff { out.push(LlmChatMessage { @@ -639,8 +544,6 @@ fn assemble_messages( content: agent_loop::TOOL_USE_DISCOVERY_HINT.to_string(), ..Default::default() }); - // Iter 5 #4: when auto_invocable skills are registered, surface them so - // the model can pick `skill.` over recreating the workflow inline. if let Some(block) = build_skills_system_block(auto_invocable_skills) { out.push(LlmChatMessage { role: "system".to_string(), @@ -661,14 +564,184 @@ fn assemble_messages( ..Default::default() }); } - Ok(AssembleOutcome { + Ok(AssembleFastOutcome { messages: out, + context_insert_pos, resolved_depth, auto_resolve_reason, reply_context_sources, }) } +fn inject_context_blocks(outcome: &mut AssembleFastOutcome, blocks: &ContextBlocks) { + let pos = outcome.context_insert_pos; + if let Some(ref sb) = blocks.semantic_block { + outcome.messages.insert( + pos, + LlmChatMessage { + role: "system".to_string(), + content: sb.clone(), + ..Default::default() + }, + ); + outcome.reply_context_sources.semantic = blocks.semantic_reply_sources.clone(); + } + if let Some(ref vb) = blocks.vault_block { + outcome.messages.insert( + pos, + LlmChatMessage { + role: "system".to_string(), + content: vb.clone(), + ..Default::default() + }, + ); + if let Some(ref vs) = blocks.vault_reply_sources { + outcome.reply_context_sources.vault_keyword = Some(vs.clone()); + } + } +} + +async fn build_context_pipeline( + canonical_root: PathBuf, + ai: vault_config::AiConfig, + last_user_query: Option, + note_context_rel_path: Option, + include_vault_context: bool, + semantic_enabled: bool, + embed_cache: Arc, + cache_dir: PathBuf, + bundle_dir: PathBuf, + note_chars_count: usize, + conversation_chars_count: usize, +) -> ContextBlocks { + let mut blocks = ContextBlocks::default(); + let query = match last_user_query { + Some(q) if !q.trim().is_empty() => q, + _ => return blocks, + }; + + let mut kw_paths: Vec = Vec::new(); + + if include_vault_context { + let root = canonical_root.clone(); + let query_clone = query.clone(); + let exclude = note_context_rel_path + .clone() + .map(|p| vec![p]) + .unwrap_or_default(); + + let vault_result = tauri::async_runtime::spawn_blocking(move || { + vault_context_search::search_workspace_context_blocking( + &root, + vault_context_search::SearchWorkspaceContextArgs { + query: query_clone, + exclude_rel_paths: exclude, + limits: None, + }, + ) + .ok() + }) + .await + .ok() + .flatten(); + + if let Some(ref vr) = vault_result { + blocks.vault_meta = Some(vr.meta.clone()); + blocks.vault_snippets = vr.snippets.clone(); + kw_paths = vr.snippets.iter().map(|s| s.rel_path.clone()).collect(); + + let root = canonical_root.clone(); + let snippets = vr.snippets.clone(); + let q = query.clone(); + let ai_for_rebuild = ai.clone(); + let rebuilt = tauri::async_runtime::spawn_blocking(move || { + vault_context_search::rebuild_vault_snippets_for_llm( + &root, + &ai_for_rebuild, + &snippets, + &q, + 1200, + 96 * 1024, + ) + }) + .await + .unwrap_or_default(); + + if !rebuilt.is_empty() { + let default_total: usize = 32_000; + let total_budget = ai + .request + .max_context_tokens + .map(|m| (m as usize).saturating_mul(4)) + .unwrap_or(default_total) + .clamp(4_000, 100_000); + let vault_cap = total_budget + .saturating_sub(note_chars_count) + .saturating_sub(conversation_chars_count) + .min(12_000) + .max(400); + if let Some((vault_block, truncated, used_rel_paths)) = + vault_context_search::build_vault_context_system_block(&rebuilt, vault_cap) + { + let used_set: HashSet<&str> = + used_rel_paths.iter().map(String::as_str).collect(); + let entries: Vec = rebuilt + .iter() + .filter(|s| used_set.contains(s.rel_path.as_str())) + .map(|s| ReplyVaultKeywordEntry { + rel_path: s.rel_path.clone(), + kind: s.kind.clone(), + }) + .collect(); + blocks.vault_reply_sources = Some(ReplyVaultKeywordSource { + entries, + truncated, + }); + blocks.vault_block = Some(vault_block); + } else { + blocks.vault_reply_sources = Some(ReplyVaultKeywordSource { + entries: Vec::new(), + truncated: true, + }); + } + } + } + } + + if semantic_enabled { + let root = canonical_root.clone(); + let q = query.clone(); + let omit = note_context_rel_path.map(|p| vec![p]).unwrap_or_default(); + let ec = Arc::clone(&embed_cache); + let cd = cache_dir; + let bd = bundle_dir; + + let semantic = tauri::async_runtime::spawn_blocking(move || { + let sem_cfg = vault_config::load_semantic_merged(&root).ok()?; + if !sem_cfg.enabled { + return None; + } + semantic_index::build_semantic_context_for_llm( + &root, &cd, &bd, &q, &sem_cfg, &kw_paths, &omit, &ec, + ) + }) + .await + .ok() + .flatten(); + + if let Some(sem) = semantic { + blocks.semantic_block = Some(sem.block); + blocks.semantic_reply_sources = ReplySemanticSource { + injected: true, + document_paths: sem.used.document_paths, + thought_ids: sem.used.thought_ids, + }; + } + } + + blocks +} + pub(super) fn emit_error(app: &AppHandle, session_id: &str, code: Option<&str>, message: &str) { let payload = LlmStreamErrorPayload { session_id: session_id.to_string(), @@ -771,7 +844,6 @@ pub async fn start_chat_stream( let cache = semantic_index::default_model_cache_dir(); let bundle = semantic_index::resolve_bundle_model_dir(&app); - let embed_paths = Some((cache.clone(), bundle.clone())); let tools_enabled = args.tools_enabled.unwrap_or(ai.tools_enabled); let skills_for_prompt: Vec<(String, String, Option)> = if tools_enabled { skills @@ -783,13 +855,12 @@ pub async fn start_chat_stream( } else { Vec::new() }; - let embed_cache_arc = Arc::clone(embed_cache_state.inner()); - let outcome = assemble_messages(&root, &ai, &args, embed_paths, &skills_for_prompt, &embed_cache_arc)?; - let mut messages = outcome.messages; - let resolved_depth = outcome.resolved_depth; - let reply_context_sources = outcome.reply_context_sources; - if let (Some(d), Some(reason)) = (resolved_depth, outcome.auto_resolve_reason.as_ref()) { + let mut fast_outcome = assemble_messages_fast(&ai, &args, &skills_for_prompt)?; + let resolved_depth = fast_outcome.resolved_depth; + let reply_context_sources = fast_outcome.reply_context_sources.clone(); + + if let (Some(d), Some(reason)) = (resolved_depth, fast_outcome.auto_resolve_reason.as_ref()) { if matches!(args.depth_mode, Some(DepthMode::Auto)) { let entry = depth_decisions::DepthDecisionEntry { timestamp: Utc::now(), @@ -830,22 +901,85 @@ pub async fn start_chat_stream( let reflection_mode = ai.memory_reflection_mode.clone(); let ai_for_memory = ai.clone(); + let include_vault_context = args.include_vault_context.unwrap_or(true); + let semantic_enabled = args.semantic_context_enabled.unwrap_or(true); + let note_context_rel_path = args.note_context.as_ref().map(|nc| nc.rel_path.clone()); + let last_user_query: Option = args + .messages + .iter() + .rev() + .find(|m| m.role.trim() == "user") + .map(|m| m.content.clone()); + let note_chars_count: usize = fast_outcome + .messages + .first() + .map(|m| m.content.chars().count()) + .unwrap_or(0); + let conversation_chars_count: usize = args + .messages + .iter() + .map(|m| m.content.chars().count() + 8) + .sum(); + let embed_cache_arc = Arc::clone(embed_cache_state.inner()); + tokio::spawn(async move { - let memory_manager: agent_loop::SharedMemoryManager = if memory_enabled { - let extraction_provider = provider::create_provider(&ai_for_memory, None, &http_client_arc).ok(); - let mgr = memory::MemoryManager::new(workspace_root.clone(), extraction_provider); + let (context_blocks, memory_manager) = tokio::join!( + build_context_pipeline( + workspace_root.clone(), + ai.clone(), + last_user_query, + note_context_rel_path, + include_vault_context, + semantic_enabled, + embed_cache_arc, + cache.clone(), + bundle.clone(), + note_chars_count, + conversation_chars_count, + ), + async { + if memory_enabled { + let extraction_provider = + provider::create_provider(&ai_for_memory, None, &http_client_arc).ok(); + let mgr = + memory::MemoryManager::new(workspace_root.clone(), extraction_provider); + Some(Arc::new(tokio::sync::Mutex::new(mgr))) + } else { + None + } + }, + ); + + inject_context_blocks(&mut fast_outcome, &context_blocks); + + if let Some(ref mm) = memory_manager { + let mgr = mm.lock().await; if let Some(mem_msg) = mgr.format_for_injection() { - let pos = if messages.is_empty() { 0 } else { 1 }; - messages.insert(pos, LlmChatMessage { - role: "system".to_string(), - content: mem_msg, - ..Default::default() - }); + let pos = if fast_outcome.messages.is_empty() { 0 } else { 1 }; + fast_outcome.messages.insert( + pos, + LlmChatMessage { + role: "system".to_string(), + content: mem_msg, + ..Default::default() + }, + ); } - Some(Arc::new(tokio::sync::Mutex::new(mgr))) - } else { - None - }; + drop(mgr); + } + + let _ = app_h.emit( + "llm:context-ready", + ContextReadyPayload { + session_id: sid.clone(), + snippets: context_blocks.vault_snippets, + meta: context_blocks.vault_meta, + reply_context_sources: fast_outcome.reply_context_sources.clone(), + }, + ); + + let messages = fast_outcome.messages; + let memory_manager: agent_loop::SharedMemoryManager = memory_manager; if tools_enabled { let manifests = registry_arc.list_for_llm_filtered(&ToolFilter::all()); diff --git a/src-tauri/src/vault_context_search.rs b/src-tauri/src/vault_context_search.rs index 8c7bf0e..2cbb08f 100644 --- a/src-tauri/src/vault_context_search.rs +++ b/src-tauri/src/vault_context_search.rs @@ -1,6 +1,6 @@ //! Vault 级关键词检索（任务 08 MVP）：与 `build_md_tree` 相同遍历规则，供 AI 上下文摘录。 //! -//! Private hits are tagged `privateOmitted` at search time; `assemble_messages` re-derives excerpts from disk to avoid trusting frontend payloads. +//! Private hits are tagged `privateOmitted` at search time; `build_context_pipeline` re-derives excerpts from disk. use crate::note_privacy; use crate::vault_config; diff --git a/src/components/AiConversationPanel.tsx b/src/components/AiConversationPanel.tsx index 6239895..b6e14b7 100644 --- a/src/components/AiConversationPanel.tsx +++ b/src/components/AiConversationPanel.tsx @@ -8,7 +8,6 @@ import { useAiNoteContext } from "../contexts/AiNoteContext"; import type { ChatMessage, ToolCallDisplayInfo } from "../hooks/useWorkspaceAiConversations"; import type { ReplyContextSources } from "../types/replyContextSources"; import { hasReplyContextSourcesToShow } from "../types/replyContextSources"; -import type { SearchWorkspaceContextResponse } from "../types/vaultContextSearch"; import type { AutoResolvedDepth, ThoughtRetrievalResult, @@ -237,7 +236,6 @@ export function AiConversationPanel() { setIncludeVaultContext, isVaultSearching, setIsVaultSearching, - vaultSearchEpochRef, depthMode, setDepthMode, autoResolved, @@ -638,6 +636,7 @@ export function AiConversationPanel() { if (p.sessionId !== activeSessionRef.current) { return; } + setIsVaultSearching(false); // P2 Tool Calling Loop：agent 模式下 stream-done 只是轮次中的中间信号， // 后续还会有 tool-call-* 与后续文本输出，只能由 llm:agent-done 最终化。 if (isAgentModeRef.current) { @@ -834,6 +833,7 @@ export function AiConversationPanel() { activeSessionRef.current = null; isAgentModeRef.current = false; setIsStreaming(false); + setIsVaultSearching(false); if (p.code === "cancelled") { setMessages((prev) => { if (composerInputRef.current.trim().length > 0) { @@ -1093,6 +1093,50 @@ export function AiConversationPanel() { }); }, ), + listen<{ + sessionId: string; + snippets: import("../types/vaultContextSearch").VaultSnippetRecord[]; + meta: import("../types/vaultContextSearch").SearchWorkspaceContextMeta | null; + replyContextSources: ReplyContextSources; + }>("llm:context-ready", (e) => { + const { sessionId, snippets, meta, replyContextSources } = e.payload; + if (sessionId !== activeSessionRef.current) return; + + if (snippets.length > 0 && meta) { + const paths = snippets.map((s) => s.relPath).join(", "); + const priv = snippets.filter((s) => s.kind === "privateOmitted").length; + let line = t("aiPanel.vaultLine", { + paths, + scannedFiles: meta.scannedFiles, + elapsedMs: meta.elapsedMs, + }); + if (priv > 0) { + line += ` ${t("aiPanel.vaultPrivateOmitted", { count: priv })}`; + } + if (meta.stoppedEarly) { + line += ` ${t("aiPanel.vaultStoppedEarly")}`; + } + setVaultSearchSummary(line); + } else { + setVaultSearchSummary(null); + } + + for (const s of snippets) { + if (s.kind !== "privateOmitted") { + sharedDocPathsRef.current.add(s.relPath); + } + } + + setMessages((prev) => { + const idx = prev.findIndex((m) => m.role === "assistant" && m.streaming); + if (idx < 0) return prev; + const next = [...prev]; + next[idx] = { ...next[idx], meta: { ...next[idx].meta, replyContextSources } }; + return next; + }); + + setIsVaultSearching(false); + }), // P2 Tool Calling Loop：Agent 轮次结束 → 最终化助手消息、清理 streaming 状态 listen<{ sessionId: string }>("llm:agent-done", (e) => { const sid = e.payload.sessionId; @@ -1136,6 +1180,7 @@ export function AiConversationPanel() { activeSessionRef.current = null; isAgentModeRef.current = false; setIsStreaming(false); + setIsVaultSearching(false); setIsPlanning(false); setMessages((prev) => { const next = [...prev]; @@ -1223,12 +1268,7 @@ export function AiConversationPanel() { }, []); const handleStop = useCallback(async () => { - if (isVaultSearching) { - vaultSearchEpochRef.current += 1; - setIsVaultSearching(false); - setVaultSearchSummary(null); - return; - } + setIsVaultSearching(false); const sid = activeSessionRef.current; if (!sid || !isTauri()) { @@ -1576,80 +1616,27 @@ export function AiConversationPanel() { setPrivacyHint(null); } - vaultSearchEpochRef.current += 1; - const searchEpoch = vaultSearchEpochRef.current; - - let vaultSearchResult: SearchWorkspaceContextResponse | null = null; - if (includeVaultContext && workspaceReady) { - setIsVaultSearching(true); - setVaultSearchSummary(null); - try { - const excludeRelPaths = - noteContext != null ? [noteContext.relPath] : ([] as string[]); - vaultSearchResult = await invoke("search_workspace_context", { - args: { - query: trimmed, - excludeRelPaths, - }, - }); - } catch (e) { - console.error(e); - vaultSearchResult = null; - } finally { - if (vaultSearchEpochRef.current === searchEpoch) { - setIsVaultSearching(false); - } - } - } - - if (vaultSearchEpochRef.current !== searchEpoch) { - return; - } - thoughtInviteExcludeRef.current = noteContext != null ? [noteContext.relPath.replace(/\\/g, "/")] : []; - if (vaultSearchResult != null && vaultSearchResult.snippets.length > 0) { - const paths = vaultSearchResult.snippets.map((s) => s.relPath).join(", "); - const priv = vaultSearchResult.snippets.filter((s) => s.kind === "privateOmitted").length; - const m = vaultSearchResult.meta; - let line = t("aiPanel.vaultLine", { - paths, - scannedFiles: m.scannedFiles, - elapsedMs: m.elapsedMs, - }); - if (priv > 0) { - line += ` ${t("aiPanel.vaultPrivateOmitted", { count: priv })}`; - } - if (m.stoppedEarly) { - line += ` ${t("aiPanel.vaultStoppedEarly")}`; - } - setVaultSearchSummary(line); - } else { - setVaultSearchSummary(null); - } - - // 记录本轮发送涉及的文档路径，用于隐私变更检测 if (noteContext && !markdownTreatAsKfPrivateForUi(noteContext.markdownForGate)) { sharedDocPathsRef.current.add(noteContext.relPath); } - if (vaultSearchResult) { - for (const s of vaultSearchResult.snippets) { - if (s.kind !== "privateOmitted") { - sharedDocPathsRef.current.add(s.relPath); - } - } - } setErrorBanner(null); setInput(""); setMessages(nextChat); + if (includeVaultContext && workspaceReady) { + setIsVaultSearching(true); + setVaultSearchSummary(null); + } + try { const streamArgs: { messages: { role: string; content: string }[]; noteContext?: { relPath: string; markdownForGate: string }; - vaultContext?: { snippets: SearchWorkspaceContextResponse["snippets"] }; + includeVaultContext?: boolean; depthMode: typeof depthMode; thoughtFocusContext?: ThoughtFocusContext; toolsEnabled?: boolean; @@ -1658,18 +1645,12 @@ export function AiConversationPanel() { messages: chatTurns, depthMode, toolsEnabled, + includeVaultContext: includeVaultContext && workspaceReady, conversationId: conversationId ?? undefined, }; if (noteContext) { streamArgs.noteContext = noteContext; } - if ( - includeVaultContext && - vaultSearchResult != null && - vaultSearchResult.snippets.length > 0 - ) { - streamArgs.vaultContext = { snippets: vaultSearchResult.snippets }; - } if ( thoughtFocusContext != null && thoughtFocusContext.thoughtId.trim() !== "" && @@ -1680,10 +1661,6 @@ export function AiConversationPanel() { const res = await invoke("start_chat_stream", { args: streamArgs, }); - if (vaultSearchEpochRef.current !== searchEpoch) { - void invoke("abort_llm_stream", { sessionId: res.sessionId }).catch(() => {}); - return; - } if (depthMode === "auto" && res.resolvedDepth) { setAutoResolved(res.resolvedDepth); } else if (depthMode !== "auto") { diff --git a/src/components/ThoughtMgmtAiConversationPanel.tsx b/src/components/ThoughtMgmtAiConversationPanel.tsx index 35532f8..3b7804b 100644 --- a/src/components/ThoughtMgmtAiConversationPanel.tsx +++ b/src/components/ThoughtMgmtAiConversationPanel.tsx @@ -8,7 +8,6 @@ import { useCallback, useEffect, useRef, useState, type KeyboardEvent } from "re import { useTranslation } from "react-i18next"; import { useThoughtMgmtAiConversationSession } from "../contexts/ThoughtMgmtAiConversationSessionContext"; import type { ThoughtFocusContext } from "../types/aiConversation"; -import type { SearchWorkspaceContextResponse } from "../types/vaultContextSearch"; import type { ProviderProfileForUi } from "../types/vaultAiConfig"; import { markdownTreatAsKfPrivateForUi } from "../utils/kfPrivateMarkdown"; import { useAiNoteContext } from "../contexts/AiNoteContext"; @@ -168,7 +167,6 @@ export function ThoughtMgmtAiConversationPanel({ tauriRuntime, isVaultSearching, setIsVaultSearching, - vaultSearchEpochRef, setThoughtFocusContext, } = useThoughtMgmtAiConversationSession(); @@ -254,6 +252,7 @@ export function ThoughtMgmtAiConversationPanel({ if (p.sessionId !== activeSessionRef.current) { return; } + setIsVaultSearching(false); markNeedPersist(); activeSessionRef.current = null; setIsStreaming(false); @@ -280,6 +279,7 @@ export function ThoughtMgmtAiConversationPanel({ markNeedPersist(); activeSessionRef.current = null; setIsStreaming(false); + setIsVaultSearching(false); if (p.code === "cancelled") { setMessages((prev) => { if (composerInputRef.current.trim().length > 0) { @@ -310,6 +310,44 @@ export function ThoughtMgmtAiConversationPanel({ }); setErrorBanner(p.message); }), + listen<{ + sessionId: string; + snippets: import("../types/vaultContextSearch").VaultSnippetRecord[]; + meta: import("../types/vaultContextSearch").SearchWorkspaceContextMeta | null; + replyContextSources: ReplyContextSources; + }>("llm:context-ready", (e) => { + const { sessionId, snippets, meta, replyContextSources } = e.payload; + if (sessionId !== activeSessionRef.current) return; + + if (snippets.length > 0 && meta) { + const paths = snippets.map((s) => s.relPath).join(", "); + const priv = snippets.filter((s) => s.kind === "privateOmitted").length; + let line = t("aiPanel.vaultLine", { + paths, + scannedFiles: meta.scannedFiles, + elapsedMs: meta.elapsedMs, + }); + if (priv > 0) { + line += ` ${t("aiPanel.vaultPrivateOmitted", { count: priv })}`; + } + if (meta.stoppedEarly) { + line += ` ${t("aiPanel.vaultStoppedEarly")}`; + } + setVaultSearchSummary(line); + } else { + setVaultSearchSummary(null); + } + + setMessages((prev) => { + const idx = prev.findIndex((m) => m.role === "assistant" && m.streaming); + if (idx < 0) return prev; + const next = [...prev]; + next[idx] = { ...next[idx], meta: { ...next[idx].meta, replyContextSources } }; + return next; + }); + + setIsVaultSearching(false); + }), ]).then((unlisteners) => { if (disposed) { unlisteners.forEach((u) => void u()); @@ -325,12 +363,7 @@ export function ThoughtMgmtAiConversationPanel({ }, [markNeedPersist, setMessages, setIsStreaming]); const handleStop = useCallback(async () => { - if (isVaultSearching) { - vaultSearchEpochRef.current += 1; - setIsVaultSearching(false); - setVaultSearchSummary(null); - return; - } + setIsVaultSearching(false); const sid = activeSessionRef.current; if (!sid || !isTauri()) { @@ -422,86 +455,38 @@ export function ThoughtMgmtAiConversationPanel({ setPrivacyHint(null); } - vaultSearchEpochRef.current += 1; - const searchEpoch = vaultSearchEpochRef.current; + setErrorBanner(null); + setInput(""); + setMessages(nextChat); - let vaultSearchResult: SearchWorkspaceContextResponse | null = null; if (workspaceReady) { setIsVaultSearching(true); setVaultSearchSummary(null); - try { - const excludeRelPaths = - noteContext != null ? [noteContext.relPath] : ([] as string[]); - vaultSearchResult = await invoke("search_workspace_context", { - args: { - query: trimmed, - excludeRelPaths, - }, - }); - } catch (e) { - console.error(e); - vaultSearchResult = null; - } finally { - if (vaultSearchEpochRef.current === searchEpoch) { - setIsVaultSearching(false); - } - } - } - - if (vaultSearchEpochRef.current !== searchEpoch) { - return; } - if (vaultSearchResult != null && vaultSearchResult.snippets.length > 0) { - const paths = vaultSearchResult.snippets.map((s) => s.relPath).join(", "); - const priv = vaultSearchResult.snippets.filter((s) => s.kind === "privateOmitted").length; - const m = vaultSearchResult.meta; - let line = t("aiPanel.vaultLine", { - paths, - scannedFiles: m.scannedFiles, - elapsedMs: m.elapsedMs, - }); - if (priv > 0) { - line += ` ${t("aiPanel.vaultPrivateOmitted", { count: priv })}`; - } - if (m.stoppedEarly) { - line += ` ${t("aiPanel.vaultStoppedEarly")}`; - } - setVaultSearchSummary(line); - } else { - setVaultSearchSummary(null); - } - - setErrorBanner(null); - setInput(""); - setMessages(nextChat); - const tf = thoughtFocusFromDetail; try { const streamArgs: { messages: { role: string; content: string }[]; noteContext?: { relPath: string; markdownForGate: string }; - vaultContext?: { snippets: SearchWorkspaceContextResponse["snippets"] }; + includeVaultContext?: boolean; depthMode: typeof DEPTH_MODE; thoughtFocusContext?: ThoughtFocusContext; - } = { messages: chatTurns, depthMode: DEPTH_MODE }; + } = { + messages: chatTurns, + depthMode: DEPTH_MODE, + includeVaultContext: workspaceReady, + }; if (noteContext) { streamArgs.noteContext = noteContext; } - if (vaultSearchResult != null && vaultSearchResult.snippets.length > 0) { - streamArgs.vaultContext = { snippets: vaultSearchResult.snippets }; - } if (tf != null && tf.thoughtId.trim() !== "" && tf.thoughtBody.trim() !== "") { streamArgs.thoughtFocusContext = tf; } const res = await invoke("start_chat_stream", { args: streamArgs, }); - if (vaultSearchEpochRef.current !== searchEpoch) { - void invoke("abort_llm_stream", { sessionId: res.sessionId }).catch(() => {}); - return; - } activeSessionRef.current = res.sessionId; setIsStreaming(true); setMessages((prev) => [ diff --git a/src/hooks/useThoughtMgmtAiConversations.ts b/src/hooks/useThoughtMgmtAiConversations.ts index ec9d81d..588cdd3 100644 --- a/src/hooks/useThoughtMgmtAiConversations.ts +++ b/src/hooks/useThoughtMgmtAiConversations.ts @@ -60,7 +60,6 @@ export function useThoughtMgmtAiConversations(opts: { const [sessionReady, setSessionReady] = useState(false); const [thoughtFocusContext, setThoughtFocusContext] = useState(null); const [isVaultSearching, setIsVaultSearching] = useState(false); - const vaultSearchEpochRef = useRef(0); const saveChainRef = useRef(Promise.resolve()); const pendingPersistRef = useRef(false); @@ -76,7 +75,6 @@ export function useThoughtMgmtAiConversations(opts: { }, [thoughtFocusContext]); useEffect(() => { - vaultSearchEpochRef.current += 1; setIsVaultSearching(false); }, [conversationId]); @@ -317,7 +315,6 @@ export function useThoughtMgmtAiConversations(opts: { setThoughtFocusContext, isVaultSearching, setIsVaultSearching, - vaultSearchEpochRef, switchConversation, createConversation, deleteConversation, diff --git a/src/hooks/useWorkspaceAiConversations.ts b/src/hooks/useWorkspaceAiConversations.ts index 45681f4..01c3657 100644 --- a/src/hooks/useWorkspaceAiConversations.ts +++ b/src/hooks/useWorkspaceAiConversations.ts @@ -111,8 +111,6 @@ export function useWorkspaceAiConversations(opts: { const [includeVaultContext, setIncludeVaultContext] = useState(false); const [thoughtFocusContext, setThoughtFocusContext] = useState(null); const [isVaultSearching, setIsVaultSearching] = useState(false); - /** 与会话切换对齐：丢弃过时的 vault 检索/发送链 */ - const vaultSearchEpochRef = useRef(0); const saveChainRef = useRef(Promise.resolve()); const pendingPersistRef = useRef(false); @@ -138,7 +136,6 @@ export function useWorkspaceAiConversations(opts: { }, [thoughtFocusContext]); useEffect(() => { - vaultSearchEpochRef.current += 1; setIsVaultSearching(false); }, [conversationId]); @@ -408,7 +405,6 @@ export function useWorkspaceAiConversations(opts: { setThoughtFocusContext, isVaultSearching, setIsVaultSearching, - vaultSearchEpochRef, switchConversation, createConversation, deleteConversation, diff --git a/src/types/vaultContextSearch.ts b/src/types/vaultContextSearch.ts index a898d0f..e6aae02 100644 --- a/src/types/vaultContextSearch.ts +++ b/src/types/vaultContextSearch.ts @@ -1,4 +1,4 @@ -/** Aligned with `search_workspace_context` / `start_chat_stream` vaultContext (camelCase) */ +/** Aligned with `llm:context-ready` event payload (camelCase) */ export type VaultSnippetKind = "excerpt" | "privateOmitted"; From ae905a58c8683960b192d462ee60d15534d26f42 Mon Sep 17 00:00:00 2001 From: donfaquir <1458918806@qq.com> Date: Wed, 24 Jun 2026 16:13:20 +0800 Subject: [PATCH 3/6] perf(llm): pre-summarize context async to reduce long-conversation TTFT --- src-tauri/src/llm/agent_loop.rs | 17 ++++- src-tauri/src/llm/context_guard.rs | 111 +++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/src-tauri/src/llm/agent_loop.rs b/src-tauri/src/llm/agent_loop.rs index 72c4034..a29fde7 100644 --- a/src-tauri/src/llm/agent_loop.rs +++ b/src-tauri/src/llm/agent_loop.rs @@ -15,7 +15,7 @@ use tauri::{AppHandle, Emitter, Manager}; use tokio_util::sync::CancellationToken; use super::approval::ToolApprovalState; -use super::context_guard::ContextGuard; +use super::context_guard::{ContextGuard, PrecomputedSummary}; use super::memory; use super::provider::{LlmProvider, NormalizedToolCall}; use super::{LlmChatMessage, LlmToolCall, LlmToolCallFunction}; @@ -137,6 +137,7 @@ pub async fn run_agent_stream( ContextGuard::with_provider(config.max_context_tokens, provider.clone()) }; let mut loop_detector = LoopDetector::new(); + let mut pending_summary: Option>> = None; loop { if cancel.is_cancelled() { @@ -154,6 +155,11 @@ pub async fn run_agent_stream( } } + if let Some(handle) = pending_summary.take() { + if let Ok(Some(cached)) = handle.await { + context_guard.apply_cached_summary(&mut messages, &cached); + } + } context_guard.trim_with_summary(&mut messages).await; // 1. 流式请求（携带 tools 字段；本轮文字会通过 emit_chunk/emit_done 推给前端） @@ -384,7 +390,14 @@ pub async fn run_agent_stream( store_extraction_msgs(&memory_manager, &messages).await; return String::new(); } - // 继续 loop：下一轮流式请求会带上完整的 messages 历史 + + if context_guard.budget_pressure(&messages) > 0.7 { + let msgs_snapshot = messages.clone(); + let guard_clone = context_guard.clone(); + pending_summary = Some(tokio::spawn(async move { + guard_clone.pre_summarize(&msgs_snapshot).await + })); + } } } diff --git a/src-tauri/src/llm/context_guard.rs b/src-tauri/src/llm/context_guard.rs index 52b38b9..65ed013 100644 --- a/src-tauri/src/llm/context_guard.rs +++ b/src-tauri/src/llm/context_guard.rs @@ -16,12 +16,18 @@ Summarize the following conversation excerpt in 2-3 concise sentences. \ Focus on: what the user asked, what tools were called, and key findings. \ Output only the summary, nothing else."; +#[derive(Clone)] pub struct ContextGuard { max_tokens: usize, reserve: usize, provider: Option>, } +pub struct PrecomputedSummary { + pub summary_text: String, + pub original_msg_count: usize, +} + impl ContextGuard { pub fn new(max_context_tokens: Option) -> Self { let max_tokens = max_context_tokens @@ -73,6 +79,91 @@ impl ContextGuard { self.max_tokens.saturating_sub(self.reserve) } + pub fn budget_pressure(&self, messages: &[LlmChatMessage]) -> f64 { + let budget = self.budget(); + if budget == 0 { + return 0.0; + } + let used = Self::estimate_total(messages); + (used as f64 / budget as f64).min(1.0) + } + + pub async fn pre_summarize( + &self, + messages: &[LlmChatMessage], + ) -> Option { + let provider = self.provider.as_ref()?; + + let tail_boundary = Self::find_tail_boundary(messages); + let removable_indices: Vec = (0..tail_boundary.min(messages.len())) + .filter(|&i| messages[i].role != "system") + .collect(); + + if removable_indices.len() < MIN_MESSAGES_FOR_SUMMARY { + return None; + } + + let removable_msgs: Vec<&LlmChatMessage> = + removable_indices.iter().map(|&i| &messages[i]).collect(); + + let summary_input = build_summary_input(&removable_msgs); + let overrides = CompletionOverrides { + temperature: Some(0.0), + ..Default::default() + }; + + let summary_text = provider + .chat_completion(&summary_input, Some(&overrides)) + .await + .ok() + .filter(|t| !t.trim().is_empty())?; + + Some(PrecomputedSummary { + summary_text, + original_msg_count: messages.len(), + }) + } + + pub fn apply_cached_summary( + &self, + messages: &mut Vec, + cached: &PrecomputedSummary, + ) -> bool { + if messages.len() != cached.original_msg_count { + return false; + } + + let tail_boundary = Self::find_tail_boundary(messages); + let removable_indices: Vec = (0..tail_boundary.min(messages.len())) + .filter(|&i| messages[i].role != "system") + .collect(); + + for &i in removable_indices.iter().rev() { + if i < messages.len() { + messages.remove(i); + } + } + + let insert_pos = messages + .iter() + .position(|m| m.role != "system") + .unwrap_or(messages.len()); + + messages.insert( + insert_pos, + LlmChatMessage { + role: "system".to_string(), + content: format!( + "[Earlier conversation summary]\n{}", + cached.summary_text.trim() + ), + ..Default::default() + }, + ); + + true + } + pub async fn trim_with_summary(&self, messages: &mut Vec) { let budget = self.budget(); if Self::estimate_total(messages) <= budget { @@ -424,6 +515,26 @@ mod tests { assert!(truncated.len() <= 10); // 6 bytes (2 chars) + "..." } + #[test] + fn budget_pressure_under() { + let guard = ContextGuard::new(Some(4096)); + let msgs = vec![sys("hello"), user("hi"), assistant("hey")]; + let pressure = guard.budget_pressure(&msgs); + assert!(pressure < 0.1, "expected low pressure, got {}", pressure); + } + + #[test] + fn budget_pressure_over() { + let guard = ContextGuard::new(Some(600)); + let msgs = vec![ + sys("system prompt"), + user(&"long message ".repeat(50)), + assistant(&"long reply ".repeat(50)), + ]; + let pressure = guard.budget_pressure(&msgs); + assert!(pressure >= 0.7, "expected high pressure, got {}", pressure); + } + #[tokio::test] async fn trim_with_summary_no_provider_falls_back() { let guard = ContextGuard::new(Some(20)); From 0cbb1c50d5a93c99013eacfa3dfe6381786f33c1 Mon Sep 17 00:00:00 2001 From: donfaquir <1458918806@qq.com> Date: Wed, 24 Jun 2026 16:45:53 +0800 Subject: [PATCH 4/6] feat(llm): stream planning tokens to frontend for visible progress --- src-tauri/src/llm/planning.rs | 4 +--- src/components/AiConversationPanel.css | 17 ++++++++++++++++ src/components/AiConversationPanel.tsx | 25 +++++++++++++++++++++++- src/hooks/useWorkspaceAiConversations.ts | 1 + 4 files changed, 43 insertions(+), 4 deletions(-) diff --git a/src-tauri/src/llm/planning.rs b/src-tauri/src/llm/planning.rs index 62f10d4..d7c2256 100644 --- a/src-tauri/src/llm/planning.rs +++ b/src-tauri/src/llm/planning.rs @@ -104,8 +104,6 @@ pub async fn run_planned_agent( memory_manager: SharedMemoryManager, ) -> String { // Phase A: Planning (no tools, text-only output) - // Use a separate session_id so plan tokens don't stream into the user's chat - let planning_sid = format!("plan-{}", uuid::Uuid::new_v4()); let tool_desc = build_tool_descriptions(&tools_json); let plan_messages = build_planning_messages(&initial_messages, &tool_desc); @@ -114,7 +112,7 @@ pub async fn run_planned_agent( let plan_result = provider .chat_stream( &app, - &planning_sid, + &session_id, plan_messages, None, cancel.clone(), diff --git a/src/components/AiConversationPanel.css b/src/components/AiConversationPanel.css index 72081bd..94a9c7d 100644 --- a/src/components/AiConversationPanel.css +++ b/src/components/AiConversationPanel.css @@ -289,3 +289,20 @@ background: color-mix(in srgb, var(--kf-accent, #6366f1) 24%, transparent); } } + +/* Planning phase: live streaming + fold */ +.ai-chat__planning-live { + opacity: 0.5; +} +.ai-chat__planning-fold { + margin-bottom: 0.5em; + border: 1px solid var(--border-subtle, #e0e0e0); + border-radius: 4px; + padding: 2px 8px; + font-size: 0.9em; +} +.ai-chat__planning-fold summary { + cursor: pointer; + color: var(--fg-muted, #888); + user-select: none; +} diff --git a/src/components/AiConversationPanel.tsx b/src/components/AiConversationPanel.tsx index b6e14b7..f6ddc70 100644 --- a/src/components/AiConversationPanel.tsx +++ b/src/components/AiConversationPanel.tsx @@ -1071,6 +1071,21 @@ export function AiConversationPanel() { listen<{ sessionId: string; planText: string }>("llm:planning-done", (e) => { if (e.payload.sessionId !== activeSessionRef.current) return; setIsPlanning(false); + const planText = e.payload.planText; + if (planText && planText.trim()) { + setMessages((prev) => { + const next = [...prev]; + const last = next[next.length - 1]; + if (last?.role === "assistant" && last.streaming) { + next[next.length - 1] = { + ...last, + content: "", + meta: { ...last.meta, planningText: planText }, + }; + } + return next; + }); + } }), // Budget warning: tool call budget reaching 80% listen<{ sessionId: string; used: number; limit: number; type: string }>( @@ -2209,7 +2224,15 @@ export function AiConversationPanel() { Agent {m.meta.budgetWarning.used}/{m.meta.budgetWarning.limit} tool calls used )} - + {m.meta?.planningText && ( +
+ Plan + +
+ )} +
+ +
{!m.streaming && m.meta?.thoughtCitation && !m.meta.thoughtCitation.privateOmitted ? (