diff --git a/GEMINI.md b/GEMINI.md index 20bae2a2..78a04e76 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -23,5 +23,6 @@ Available subsystem targets: | `codemark::http` | HTTP client, remote tour listing, pack download/upload, sync | | `codemark::shell` | External editor spawning (program, args, errors) | | `codemark::auth` | Server authorization decisions: repo read/write access checks, public-repo visibility, allow/deny outcomes | +| `codemark::embeddings` | Embedding provider lifecycle: shared-provider cache reuse vs. creation, model loading | When adding new functionality, instrument it with `tracing::debug!` (or `info!`/`warn!`/`error!` as appropriate) using the matching subsystem target. If adding a new subsystem, define a new `codemark::` target and add a row to the table above. diff --git a/crates/codemark-core/src/embeddings/local.rs b/crates/codemark-core/src/embeddings/local.rs index 24975f29..2477095e 100644 --- a/crates/codemark-core/src/embeddings/local.rs +++ b/crates/codemark-core/src/embeddings/local.rs @@ -9,8 +9,9 @@ use candle_core::{DType, Device, Result as CandleResult, Tensor}; use candle_nn::VarBuilder; use candle_transformers::models::bert::{BertModel, Config as BertConfig}; use hf_hub::{Repo, RepoType, api::sync::Api}; +use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::Mutex; +use std::sync::{Arc, Mutex, OnceLock}; /// Mean pooling layer for sentence embeddings. struct MeanPooling; @@ -267,6 +268,51 @@ impl LocalEmbeddingProvider { } } +/// Process-wide cache of loaded embedding providers, keyed by model id + cache +/// directory. +/// +/// Loading a model reads its weights from disk and builds the BERT graph, which +/// took seconds on *every* semantic search because each search built a fresh +/// [`LocalEmbeddingProvider`] with an empty model slot. A provider embeds through +/// `&self` (interior `Mutex`es over the model/tokenizer), so one instance is +/// safely shared across threads — concurrent embeds simply serialize on the lock. +/// Caching it keeps the model resident so only the first search in a process pays +/// the load cost; later searches (and the tab switches that run alongside them) +/// no longer contend with a re-load. +static PROVIDER_CACHE: OnceLock>>> = + OnceLock::new(); + +/// Return a shared, cached [`LocalEmbeddingProvider`] for the given model and +/// cache directory, constructing (but not yet loading) it on first use. The model +/// weights load lazily on the first `embed`/`embed_batch` call and then stay +/// resident for the life of the process. +pub fn shared_local_provider( + model: EmbeddingModel, + cache_dir: Option, +) -> EmbeddingResult> { + let key = format!("{}|{:?}", model.model_id(), cache_dir); + let cache = PROVIDER_CACHE.get_or_init(|| Mutex::new(HashMap::new())); + // A poisoned lock only means a prior holder panicked while inserting; the map + // itself is a plain HashMap that can't be left half-updated, so recover it. + let mut map = cache.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(provider) = map.get(&key) { + tracing::debug!( + target: "codemark::embeddings", + model = %model.model_id(), + "reusing cached embedding provider" + ); + return Ok(Arc::clone(provider)); + } + tracing::debug!( + target: "codemark::embeddings", + model = %model.model_id(), + "creating embedding provider (model weights load on first embed)" + ); + let provider = Arc::new(LocalEmbeddingProvider::new(model, cache_dir)?); + map.insert(key, Arc::clone(&provider)); + Ok(provider) +} + #[async_trait] impl EmbeddingProvider for LocalEmbeddingProvider { async fn embed(&self, text: &str) -> EmbeddingResult> { diff --git a/crates/codemark-core/src/embeddings/mod.rs b/crates/codemark-core/src/embeddings/mod.rs index 43483366..02466f68 100644 --- a/crates/codemark-core/src/embeddings/mod.rs +++ b/crates/codemark-core/src/embeddings/mod.rs @@ -14,6 +14,6 @@ pub mod vec_store; pub use config::{DistanceMetric, EmbeddingConfig, EmbeddingModel}; #[cfg(feature = "semantic")] -pub use local::LocalEmbeddingProvider; +pub use local::{LocalEmbeddingProvider, shared_local_provider}; pub use provider::EmbeddingProvider; pub use vec_store::{SearchResult, VecStore, VecStoreEntry}; diff --git a/crates/codemark-core/src/storage/semantic_repo.rs b/crates/codemark-core/src/storage/semantic_repo.rs index c3dd6687..6be3027a 100644 --- a/crates/codemark-core/src/storage/semantic_repo.rs +++ b/crates/codemark-core/src/storage/semantic_repo.rs @@ -42,11 +42,20 @@ impl SemanticRepo { Self { cache_dir, model, distance_metric, threshold } } - /// Get or create the embedding provider. - fn provider(&self) -> Result { - LocalEmbeddingProvider::new(self.model.clone(), self.cache_dir.clone()).map_err(|e| { - crate::error::Error::Operation(format!("Failed to create embedding provider: {}", e)) - }) + /// Get the shared embedding provider. + /// + /// Returns a process-cached provider (keyed by model + cache dir) so the model + /// loads once and stays resident, instead of re-loading its weights on every + /// search. The provider embeds through `&self`, so sharing it across searches + /// and threads is safe. + fn provider(&self) -> Result> { + crate::embeddings::shared_local_provider(self.model.clone(), self.cache_dir.clone()) + .map_err(|e| { + crate::error::Error::Operation(format!( + "Failed to create embedding provider: {}", + e + )) + }) } /// Generate an embedding for a bookmark's searchable text. diff --git a/crates/codemark-tui/src/browser/events.rs b/crates/codemark-tui/src/browser/events.rs index 3916357d..132bc757 100644 --- a/crates/codemark-tui/src/browser/events.rs +++ b/crates/codemark-tui/src/browser/events.rs @@ -533,12 +533,18 @@ impl BrowserLayout { } Some(true) } - Event::PreviewReady { request_id, payload } => { + Event::PreviewReady { request_id, repo_root, payload } => { // Drop stale results: the selection moved on since this resolve // was spawned, so a newer request supersedes it. if *request_id == self.active_preview_request { self.inflight_preview = None; self.right_pane.apply_preview((**payload).clone()); + // `apply_preview` has no db, so record the previewed bookmark's + // owning repo here (mirrors the synchronous `load_bookmark_live` + // path) so a later right-pane refresh/action resolves against the + // right db under multi-select instead of a stale repo. + self.right_pane.active_repo_root = + super::RightPane::repo_root_of(self.workspace.db_for(repo_root.as_deref())); } Some(true) } @@ -1157,7 +1163,7 @@ impl BrowserLayout { /// `preview_cache` across tasks so repeat visits to the same file reuse the /// parse tree. Debounce serializes these tasks, so the lock is effectively /// uncontended. - fn spawn_preview_task(&self, bookmark_id: String, repo_root: Option) { + pub(super) fn spawn_preview_task(&self, bookmark_id: String, repo_root: Option) { use codemark_core::storage::db::Database; let request_id = self.active_preview_request; @@ -1206,7 +1212,8 @@ impl BrowserLayout { false, ) { Ok(Some(payload)) => { - let _ = event_handler.send(Event::PreviewReady { request_id, payload }); + let _ = + event_handler.send(Event::PreviewReady { request_id, repo_root, payload }); } Ok(None) => { let _ = event_handler.send(Event::PreviewFailed { @@ -1915,8 +1922,11 @@ impl BrowserLayout { self.fetch_remote_tours(); } // Refresh the preview for the newly active tab so it doesn't - // linger on content from the previous tab. - self.update_content_live_preview(); + // linger on content from the previous tab. Bookmarks resolve + // in the background so a slow live resolve (e.g. under CPU + // contention from a concurrent semantic search) can't freeze + // the switch. + self.preview_after_tab_change(); } // A click on the sort glyph reordered the list; refresh the @@ -1976,8 +1986,11 @@ impl BrowserLayout { self.fetch_remote_tours(); } // Refresh the preview for the newly active tab so it doesn't - // linger on content from the previous tab. - self.update_content_live_preview(); + // linger on content from the previous tab. Bookmarks resolve + // in the background so a slow live resolve (e.g. under CPU + // contention from a concurrent semantic search) can't freeze + // the switch. + self.preview_after_tab_change(); } handled diff --git a/crates/codemark-tui/src/browser/mod.rs b/crates/codemark-tui/src/browser/mod.rs index ad49f921..16958bc3 100644 --- a/crates/codemark-tui/src/browser/mod.rs +++ b/crates/codemark-tui/src/browser/mod.rs @@ -451,6 +451,55 @@ impl BrowserLayout { } } + /// Refresh the right-pane preview after the active Content tab changed. + /// + /// Unlike [`update_content_live_preview`](Self::update_content_live_preview) + /// — used for init and focus-enter, which resolve synchronously for immediate + /// feedback and the e2e snapshots — a tab switch routes the *Bookmarks* + /// preview through the background resolve. Its live tree-sitter resolve can be + /// slow (a cold/large file, or CPU contention from a concurrent semantic + /// search still loading its model / resolving list health), and doing it + /// synchronously here blocked the event loop, so switching back to Bookmarks + /// after a search felt like a freeze. Collections/Tours overviews are DB-only + /// (no parse) and stay synchronous. + fn preview_after_tab_change(&mut self) { + if self.focus != FocusArea::ContentPanel { + return; + } + + let Some(tab) = ContentTab::from_index(self.left_pane.content_panel.tabs.selected_index()) + else { + return; + }; + + // A background Bookmarks preview spawned on a prior tab must not land on + // top of another tab's content. The Bookmarks arm re-bumps the request id + // itself (via `request_bookmark_preview_now`); for the synchronously + // rendered overview tabs, invalidate any in-flight/pending bookmark + // preview up front so a late `PreviewReady` is dropped — even when the new + // tab has no selection to render below (the early return would otherwise + // leave the stale request active). + if tab != ContentTab::Bookmarks { + self.cancel_inflight_preview(); + } + + let selected_id = self + .left_pane + .content_panel + .active_panel() + .and_then(|panel| panel.selected()) + .and_then(|selected| selected.user_data.clone()); + + let Some(id) = selected_id else { + return; + }; + + match tab { + ContentTab::Bookmarks => self.request_bookmark_preview_now(&id), + ContentTab::Collections | ContentTab::Tours => self.preview_content_item(tab, &id), + } + } + /// Whether the user has usable sync credentials (a resolvable server *and* a /// token). Drives whether the Tours tab — which lists remote tours — is shown. /// @@ -638,6 +687,41 @@ impl BrowserLayout { Some((id.to_string(), label, repo_root, self.tick_count + PREVIEW_DEBOUNCE_TICKS)); } + /// Resolve the selected bookmark's preview on a background task *immediately*, + /// without the [`request_bookmark_preview`](Self::request_bookmark_preview) + /// debounce. + /// + /// Used for discrete actions (a tab switch) where there's no fast scroll to + /// coalesce, so the resolve should start at once — but still off the event + /// loop, so a slow live resolve or CPU contention from a concurrent search + /// can't freeze the UI the way the old synchronous preview path did. The + /// loading indicator shows right away (replacing any overview from the tab we + /// left) so the pane doesn't linger on stale content while the resolve runs. + fn request_bookmark_preview_now(&mut self, id: &str) { + // Bump the request id so any in-flight/older result is treated as stale, + // and drop any queued debounced request for a prior selection. + self.preview_seq = self.preview_seq.wrapping_add(1); + self.active_preview_request = self.preview_seq; + self.pending_preview = None; + + let (label, repo_root) = self + .left_pane + .content_panel + .active_panel() + .and_then(|panel| panel.selected()) + .map(|selected| { + (Some(selected.text().to_string()), selected.repo_root().map(str::to_string)) + }) + .unwrap_or((None, None)); + + // Show the loading indicator immediately (a tab switch is discrete, not a + // scroll, so nothing flashes on rapid moves), then resolve in the + // background; the result arrives via `PreviewReady`. + self.right_pane.begin_bookmark_loading(label.clone()); + self.inflight_preview = Some((label, self.tick_count)); + self.spawn_preview_task(id.to_string(), repo_root); + } + /// Invalidate any pending or in-flight bookmark preview so a stale async /// result can't clobber a subsequent *synchronous* render. /// diff --git a/crates/codemark-tui/src/browser/right_pane.rs b/crates/codemark-tui/src/browser/right_pane.rs index 26fd07ec..99a4b104 100644 --- a/crates/codemark-tui/src/browser/right_pane.rs +++ b/crates/codemark-tui/src/browser/right_pane.rs @@ -322,6 +322,19 @@ impl RightPane { self.loading_label = label; } + /// Enter the loading state for a single-bookmark preview, replacing any active + /// collection/tour overview. + /// + /// A bare [`begin_loading`](Self::begin_loading) keeps `overview_active` set, + /// so the render path (which shows the overview *instead of* the loading + /// block while it's active) would leave the previous tab's overview on screen. + /// Clearing it here makes the loading indicator show immediately when a tab + /// switch hands a bookmark resolve to the background. + pub fn begin_bookmark_loading(&mut self, label: Option) { + self.overview_active = false; + self.begin_loading(label); + } + /// Leave the loading state without applying a preview (e.g. on failure). pub fn finish_loading(&mut self) { self.loading = false; @@ -1077,7 +1090,7 @@ impl RightPane { /// Repo root of a database (the parent of its `.codemark` dir) as a string /// matching `PanelItem::repo_root()`. `None` for in-memory or degenerate paths /// (which resolve to the focused repo). - fn repo_root_of(db: &Database) -> Option { + pub(super) fn repo_root_of(db: &Database) -> Option { db.path().parent().and_then(|p| p.parent()).map(|p| p.to_string_lossy().to_string()) } diff --git a/crates/codemark-tui/src/event.rs b/crates/codemark-tui/src/event.rs index 34f319e2..af769d53 100644 --- a/crates/codemark-tui/src/event.rs +++ b/crates/codemark-tui/src/event.rs @@ -116,6 +116,10 @@ pub enum Event { PreviewReady { /// Monotonic id of the preview request this result answers. request_id: u64, + /// The previewed bookmark's owning repo tag (None → focused repo), so the + /// right pane records the right db as active — a later refresh/action then + /// resolves against the correct repo under multi-select. + repo_root: Option, /// Fully-computed preview, boxed to keep the enum small. payload: Box, },