Merge pull request #301 from blopker/background-dictionary-prefetch

Move dictionary downloads off the spell-check path

Author: blopker

Cargo.lock

@@ -59,7 +59,7 @@ tree-sitter-cpp = "<0.25.0"
 tree-sitter-css = "<0.26.0"
 tree-sitter-dart = "<1"
 tree-sitter-elixir = "<0.4.0"
-tree-sitter-erlang = "<0.20.0"
+tree-sitter-erlang = "<0.21.0"
 tree-sitter-go = "<0.26.0"
 tree-sitter-haskell = "<0.25.0"
 tree-sitter-html = "<0.25.0"

Cargo.toml

@@ -47,6 +47,11 @@ pub trait CodebookConfig: Sync + Send + Debug {
     fn add_ignore(&self, file: &str) -> bool;
     fn add_include(&self, file: &str) -> bool;
     fn get_dictionary_ids(&self) -> Vec<String>;
+    /// Every dictionary ID this config can resolve to for any file,
+    /// including per-path override blocks. Used for prefetching.
+    fn all_dictionary_ids(&self) -> Vec<String> {
+        self.get_dictionary_ids()
+    }
     fn should_ignore_path(&self, path: &Path) -> bool;
     fn should_include_path(&self, path: &Path) -> bool;
     fn is_allowed_word(&self, word: &str) -> bool;
@@ -508,6 +513,10 @@ impl CodebookConfig for CodebookConfigFile {
         snapshot.dictionary_ids()
     }
 
+    fn all_dictionary_ids(&self) -> Vec<String> {
+        self.snapshot().all_dictionary_ids()
+    }
+
     /// Check if a path is included based on the effective configuration
     fn should_include_path(&self, path: &Path) -> bool {
         let snapshot = self.snapshot();
@@ -624,6 +633,10 @@ impl CodebookConfig for CodebookConfigMemory {
         snapshot.dictionary_ids()
     }
 
+    fn all_dictionary_ids(&self) -> Vec<String> {
+        self.snapshot().all_dictionary_ids()
+    }
+
     fn should_include_path(&self, path: &Path) -> bool {
         let snapshot = self.snapshot();
         snapshot.should_include_path(path)

crates/codebook-config/src/lib.rs

@@ -492,6 +492,21 @@ impl ConfigSettings {
         }
     }
 
+    /// Every dictionary ID this config can resolve to for any file: the base
+    /// ids plus each override block's `dictionaries` and `extra_dictionaries`.
+    /// Sorted and deduped. Used to prefetch dictionaries before a matching
+    /// file is ever opened.
+    pub fn all_dictionary_ids(&self) -> Vec<String> {
+        let mut ids = self.dictionary_ids();
+        for block in &self.overrides {
+            ids.extend(block.dictionaries.iter().flatten().cloned());
+            ids.extend(block.extra_dictionaries.iter().flatten().cloned());
+        }
+        ids.sort();
+        ids.dedup();
+        ids
+    }
+
     /// Determine whether a path should be included based on the configured glob patterns.
     pub fn should_include_path(&self, path: &Path) -> bool {
         if self.include_paths.is_empty() {
@@ -1286,6 +1301,42 @@ mod tests {
         assert_eq!(resolved.words, vec!["codebook"]);
     }
 
+    #[test]
+    fn test_all_dictionary_ids_unions_base_and_overrides() {
+        let settings = ConfigSettings {
+            dictionaries: vec!["en_us".to_string(), "es".to_string()],
+            overrides: vec![
+                OverrideBlock {
+                    paths: vec![glob("docs/**")],
+                    dictionaries: Some(vec!["de".to_string(), "en_us".to_string()]),
+                    ..OverrideBlock::default_for_test()
+                },
+                OverrideBlock {
+                    paths: vec![glob("**/*.md")],
+                    extra_dictionaries: Some(vec!["fr".to_string()]),
+                    ..OverrideBlock::default_for_test()
+                },
+            ],
+            ..Default::default()
+        };
+
+        assert_eq!(settings.all_dictionary_ids(), ["de", "en_us", "es", "fr"]);
+    }
+
+    #[test]
+    fn test_all_dictionary_ids_includes_implicit_default() {
+        let settings = ConfigSettings {
+            overrides: vec![OverrideBlock {
+                paths: vec![glob("**/*.md")],
+                extra_dictionaries: Some(vec!["es".to_string()]),
+                ..OverrideBlock::default_for_test()
+            }],
+            ..Default::default()
+        };
+
+        assert_eq!(settings.all_dictionary_ids(), ["en_us", "es"]);
+    }
+
     #[test]
     fn test_merge_preserves_override_order() {
         let mut global = ConfigSettings {

crates/codebook-config/src/settings.rs

@@ -2,3 +2,4 @@ pub mod file_cache;
 mod init_options;
 pub mod lsp;
 pub mod lsp_logger;
+mod prefetch;

crates/codebook-lsp/src/lib.rs

@@ -104,6 +104,24 @@ pub fn run_lint(files: &[String], root: &Path, unique: bool, suggest: bool) -> L
 
     let codebook = Codebook::new(config.clone());
 
+    // The lint run is one-shot, so unlike the LSP it blocks on downloads.
+    // Individual failures degrade to checking with whatever is available —
+    // but with no primary dictionary at all, every check noops, and
+    // reporting that as a clean run would be a false green in CI.
+    let warmup = codebook.warm_dictionaries();
+    if warmup.network_disabled {
+        eprintln!("NO_NETWORK set; checking with cached dictionaries only");
+    }
+    for (id, e) in &warmup.failures {
+        err!("could not download dictionary '{id}': {e} (checking with cached/available dictionaries)");
+    }
+    if !warmup.primary_available {
+        err!(
+            "no configured dictionary is available (downloads failed, or the cache is empty with NO_NETWORK set); cannot check spelling"
+        );
+        return LintResult::Failure;
+    }
+
     // Canonicalize the root once here rather than once per file.
     let root_canonical = root.canonicalize().ok();
 

crates/codebook-lsp/src/lint.rs

@@ -24,6 +24,7 @@ use log::{debug, info};
 use crate::file_cache::TextDocumentCache;
 use crate::init_options::ClientInitializationOptions;
 use crate::lsp_logger;
+use crate::prefetch::{self, PrefetchHandle};
 
 const SOURCE_NAME: &str = "Codebook";
 
@@ -62,6 +63,23 @@ fn compute_relative_path(
 }
 
 pub struct Backend {
+    state: Arc<BackendState>,
+}
+
+/// `Backend` is a pointer newtype over the shared state; Deref keeps every
+/// handler spelling `self.field` instead of `self.state.field`.
+impl std::ops::Deref for Backend {
+    type Target = BackendState;
+    fn deref(&self) -> &Self::Target {
+        &self.state
+    }
+}
+
+/// Shared server state, `Arc`'d so tasks outside the LSP dispatch loop (the
+/// prefetch worker's recheck listener) can reach the same document cache,
+/// config, and client. Public only because `Deref<Target = BackendState>`
+/// requires it; all fields are private.
+pub struct BackendState {
     client: Client,
     workspace_dir: PathBuf,
     /// Cached canonicalized workspace directory for efficient relative path computation
@@ -72,6 +90,9 @@ pub struct Backend {
     initialize_options: RwLock<Arc<ClientInitializationOptions>>,
     /// When the config files were last polled for changes (None = never)
     last_config_poll: Mutex<Option<Instant>>,
+    /// Handle to the background dictionary prefetch worker (None until
+    /// `initialized`, or forever when NO_NETWORK is set)
+    prefetch: OnceLock<PrefetchHandle>,
 }
 
 enum CodebookCommand {
@@ -169,6 +190,7 @@ impl LanguageServer for Backend {
             "Global config: {}",
             config.global_config_path().unwrap_or_default().display()
         );
+        self.spawn_prefetch();
     }
 
     async fn shutdown(&self) -> RpcResult<()> {
@@ -388,17 +410,45 @@ impl Backend {
     pub fn new(client: Client, workspace_dir: &Path) -> Self {
         let workspace_dir_canonical = workspace_dir.canonicalize().ok();
         Self {
-            client,
-            workspace_dir: workspace_dir.to_path_buf(),
-            workspace_dir_canonical,
-            codebook: OnceLock::new(),
-            config: OnceLock::new(),
-            document_cache: TextDocumentCache::default(),
-            initialize_options: RwLock::new(Arc::new(ClientInitializationOptions::default())),
-            last_config_poll: Mutex::new(None),
+            state: Arc::new(BackendState {
+                client,
+                workspace_dir: workspace_dir.to_path_buf(),
+                workspace_dir_canonical,
+                codebook: OnceLock::new(),
+                config: OnceLock::new(),
+                document_cache: TextDocumentCache::default(),
+                initialize_options: RwLock::new(Arc::new(ClientInitializationOptions::default())),
+                last_config_poll: Mutex::new(None),
+                prefetch: OnceLock::new(),
+            }),
+        }
+    }
+
+    /// Start the background dictionary prefetch worker and the task that
+    /// re-checks open documents when a new dictionary lands.
+    fn spawn_prefetch(&self) {
+        if self.prefetch.get().is_some() {
+            return;
         }
+        if codebook::network_disabled() {
+            info!("NO_NETWORK set; dictionary downloads disabled");
+            return;
+        }
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+        let state = self.state.clone();
+        tokio::spawn(async move {
+            while rx.recv().await.is_some() {
+                // Coalesce a burst of completions into one recheck
+                while rx.try_recv().is_ok() {}
+                state.recheck_all().await;
+            }
+        });
+        let handle = prefetch::spawn(self.codebook_handle(), tx);
+        let _ = self.prefetch.set(handle);
     }
+}
 
+impl BackendState {
     fn config_handle(&self) -> Arc<CodebookConfigFile> {
         self.config
             .get_or_init(|| {
@@ -556,6 +606,11 @@ impl Backend {
 
         if did_reload {
             debug!("Config reloaded, rechecking all files.");
+            // The change may add dictionaries: wake the prefetch worker and
+            // reset its backoff so they download immediately.
+            if let Some(handle) = self.prefetch.get() {
+                handle.kick();
+            }
             self.recheck_all().await;
         } else {
             debug!("Checking file: {uri:?}");

crates/codebook-lsp/src/lsp.rs

@@ -3,6 +3,7 @@ mod init_options;
 mod lint;
 mod lsp;
 mod lsp_logger;
+mod prefetch;
 
 use clap::{Parser, Subcommand};
 use codebook_config::{CodebookConfig, CodebookConfigFile, ConfigError};