ssg/plugins/
search_index.rs

1// Copyright © 2023 - 2026 Static Site Generator (SSG). All rights reserved.
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! Vector-search artifact emitter (issue #545).
5//!
6//! Walks the compiled site for HTML files, runs the `ssg-search`
7//! encoder over each page's visible text, and writes the four-file
8//! WASM-loadable bundle to `<site>/search/`:
9//!
10//!   - `embeddings.bin`  — pre-normalised f32 vectors, little-endian.
11//!   - `manifest.json`   — row index → `{url, title, excerpt}` map.
12//!   - `model.bin`       — encoder weights / config (with magic header).
13//!   - `tokenizer.bin`   — tokeniser config (with magic header).
14//!
15//! Intended to run **alongside** [`crate::search::SearchPlugin`]
16//! (which is a separate keystroke-based lexical search). The two
17//! coexist — `search-index.json` powers the modal autocomplete,
18//! `search/embeddings.bin` powers the semantic ranked results
19//! returned by the WASM engine.
20//!
21//! ## Why a separate plugin
22//!
23//! The vector bundle is large (one f32 per dim per doc) so it ships
24//! to the browser separately and lazily — never inlined into HTML —
25//! and the worker fetches it the first time the user types. The
26//! existing `SearchPlugin` is on every page; this plugin is opt-in
27//! via [`crate::plugin::PluginManager::register`].
28
29use crate::error::{PathErrorExt, SsgError};
30use crate::plugin::{Plugin, PluginContext};
31use crate::search::SearchIndex;
32use ssg_search::{
33    paths::{EMBEDDINGS_FILE, MANIFEST_FILE, MODEL_FILE, TOKENIZER_FILE},
34    ArtifactsBuilder,
35};
36use std::fs;
37
38/// Short excerpt length in characters — fits a single search-result
39/// snippet line in typical layouts.
40const EXCERPT_CHARS: usize = 160;
41
42/// Plugin that builds the `<site>/search/` vector bundle.
43///
44/// # Examples
45///
46/// ```
47/// use ssg::plugin::Plugin;
48/// use ssg::search_index::VectorSearchPlugin;
49/// assert_eq!(VectorSearchPlugin.name(), "vector-search");
50/// ```
51#[derive(Debug, Clone, Copy, Default)]
52pub struct VectorSearchPlugin;
53
54impl Plugin for VectorSearchPlugin {
55    fn name(&self) -> &'static str {
56        "vector-search"
57    }
58
59    fn after_compile(&self, ctx: &PluginContext) -> Result<(), SsgError> {
60        if !ctx.site_dir.exists() {
61            return Ok(());
62        }
63        if ctx.dry_run {
64            return Ok(());
65        }
66
67        // Re-use SearchIndex's HTML-walking logic so we get title +
68        // body text via lol_html (same filters that exclude
69        // <script>/<style>/<nav>/<footer>/<head>).
70        let index = SearchIndex::build(&ctx.site_dir)?;
71        if index.is_empty() {
72            return Ok(());
73        }
74
75        let mut builder = ArtifactsBuilder::default();
76        for e in &index.entries {
77            let excerpt = truncate_chars(&e.content, EXCERPT_CHARS);
78            let _ = builder.add_doc(ssg_search::artifacts::InputDoc {
79                url: e.url.clone(),
80                title: e.title.clone(),
81                body: format!("{}\n{}", e.title, e.content),
82                excerpt,
83            });
84        }
85
86        let artifacts = builder.build();
87
88        // Materialise the four files under <site>/search/.
89        let dir = ctx.site_dir.join("search");
90        fs::create_dir_all(&dir).with_path(&dir)?;
91
92        let emb_path = dir.join(EMBEDDINGS_FILE);
93        fs::write(&emb_path, &artifacts.embeddings).with_path(&emb_path)?;
94
95        let man_path = dir.join(MANIFEST_FILE);
96        let manifest_json = stamp_embeddings_hash(
97            &artifacts.manifest_json,
98            &artifacts.embeddings,
99        );
100        fs::write(&man_path, manifest_json).with_path(&man_path)?;
101
102        let model_path = dir.join(MODEL_FILE);
103        fs::write(&model_path, &artifacts.model).with_path(&model_path)?;
104
105        let tok_path = dir.join(TOKENIZER_FILE);
106        fs::write(&tok_path, &artifacts.tokenizer).with_path(&tok_path)?;
107
108        log::info!(
109            "[vector-search] Wrote {} docs ({} bytes embeddings) to {}",
110            artifacts.count(),
111            artifacts.embeddings.len(),
112            dir.display()
113        );
114        Ok(())
115    }
116}
117
118/// Adds an `embeddings_sha256` field to the manifest JSON so the
119/// `search_index` audit gate can verify `embeddings.bin` integrity
120/// (the gate reads `manifest.json#embeddings_sha256` and compares it
121/// to the SHA-256 of the binary — see
122/// `src/audit/gates/search_index.rs`). Returns the manifest verbatim
123/// if it fails to parse (defensive; `ssg-search` always emits valid
124/// JSON).
125fn stamp_embeddings_hash(manifest_json: &[u8], embeddings: &[u8]) -> Vec<u8> {
126    use sha2::{Digest, Sha256};
127
128    let Ok(mut manifest) =
129        serde_json::from_slice::<serde_json::Value>(manifest_json)
130    else {
131        return manifest_json.to_vec();
132    };
133    let Some(obj) = manifest.as_object_mut() else {
134        return manifest_json.to_vec();
135    };
136
137    let mut hasher = Sha256::new();
138    hasher.update(embeddings);
139    let digest = hasher.finalize();
140    let mut hex = String::with_capacity(digest.len() * 2);
141    for byte in digest {
142        use std::fmt::Write as _;
143        // Infallible on String; ignore the Result per fmt::Write docs.
144        let _ = write!(hex, "{byte:02x}");
145    }
146
147    let _ = obj.insert(
148        "embeddings_sha256".to_string(),
149        serde_json::Value::String(hex),
150    );
151    serialize_manifest_value(&manifest)
152        .unwrap_or_else(|_| manifest_json.to_vec())
153}
154
155/// Serialize the stamped manifest with a fault-injection hook so tests
156/// can drive the defensive fallback in [`stamp_embeddings_hash`]
157/// (pretty-printing a `Value` parsed from `ssg-search`'s own valid
158/// JSON output, plus one inserted hex-string field, cannot fail in
159/// practice).
160fn serialize_manifest_value(
161    manifest: &serde_json::Value,
162) -> serde_json::Result<Vec<u8>> {
163    fail_point!("search_index::manifest-serialize", |_| Err(
164        <serde_json::Error as serde::ser::Error>::custom(
165            "injected: search_index::manifest-serialize"
166        )
167    ));
168    serde_json::to_vec_pretty(manifest)
169}
170
171/// Returns the first `max_chars` Unicode scalar values of `s`. (Plain
172/// `s[..max]` slicing on bytes would split multibyte sequences.)
173fn truncate_chars(s: &str, max_chars: usize) -> String {
174    if s.chars().count() <= max_chars {
175        return s.to_string();
176    }
177    s.chars().take(max_chars).collect()
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183    use std::path::Path;
184    use tempfile::tempdir;
185
186    fn ctx(p: &Path) -> PluginContext {
187        PluginContext::new(p, p, p, p)
188    }
189
190    fn write_html(dir: &Path, name: &str, body: &str) {
191        let p = dir.join(name);
192        fs::write(
193            &p,
194            format!(
195                "<!DOCTYPE html><html><head><title>{name}</title></head><body>{body}</body></html>"
196            ),
197        )
198        .unwrap();
199    }
200
201    #[test]
202    fn plugin_name() {
203        assert_eq!(VectorSearchPlugin.name(), "vector-search");
204    }
205
206    #[test]
207    fn plugin_is_noop_on_missing_dir() {
208        let tmp = tempdir().unwrap();
209        let missing = tmp.path().join("not-here");
210        let c = ctx(&missing);
211        VectorSearchPlugin.after_compile(&c).unwrap();
212        assert!(!missing.exists());
213    }
214
215    #[test]
216    fn plugin_is_noop_in_dry_run() {
217        let tmp = tempdir().unwrap();
218        write_html(tmp.path(), "page.html", "<p>hello</p>");
219        let c = ctx(tmp.path()).with_dry_run(true);
220        VectorSearchPlugin.after_compile(&c).unwrap();
221        assert!(!tmp.path().join("search").exists());
222    }
223
224    #[test]
225    fn plugin_is_noop_on_empty_corpus() {
226        let tmp = tempdir().unwrap();
227        let c = ctx(tmp.path());
228        VectorSearchPlugin.after_compile(&c).unwrap();
229        // No HTML → SearchIndex empty → no search/ dir is created.
230        assert!(!tmp.path().join("search").exists());
231    }
232
233    #[test]
234    fn plugin_emits_four_artifacts() {
235        let tmp = tempdir().unwrap();
236        write_html(
237            tmp.path(),
238            "a.html",
239            "<p>rust webassembly compiles to portable browser modules</p>",
240        );
241        write_html(
242            tmp.path(),
243            "b.html",
244            "<p>baking sourdough bread starter flour</p>",
245        );
246        let c = ctx(tmp.path());
247        VectorSearchPlugin.after_compile(&c).unwrap();
248
249        let dir = tmp.path().join("search");
250        assert!(dir.exists());
251        assert!(dir.join("embeddings.bin").exists());
252        assert!(dir.join("manifest.json").exists());
253        assert!(dir.join("model.bin").exists());
254        assert!(dir.join("tokenizer.bin").exists());
255    }
256
257    #[test]
258    fn embeddings_size_is_n_x_d_x_4() {
259        let tmp = tempdir().unwrap();
260        write_html(tmp.path(), "a.html", "<p>foo bar baz</p>");
261        write_html(tmp.path(), "b.html", "<p>quux corge grault</p>");
262        let c = ctx(tmp.path());
263        VectorSearchPlugin.after_compile(&c).unwrap();
264        let emb = fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
265        assert_eq!(emb.len(), 2 * 256 * 4);
266    }
267
268    #[test]
269    fn manifest_has_one_entry_per_html() {
270        let tmp = tempdir().unwrap();
271        write_html(tmp.path(), "x.html", "<p>alpha</p>");
272        write_html(tmp.path(), "y.html", "<p>beta</p>");
273        let c = ctx(tmp.path());
274        VectorSearchPlugin.after_compile(&c).unwrap();
275        let json = fs::read_to_string(tmp.path().join("search/manifest.json"))
276            .unwrap();
277        let m: serde_json::Value = serde_json::from_str(&json).unwrap();
278        assert_eq!(m["count"].as_u64().unwrap(), 2);
279        assert_eq!(m["entries"].as_array().unwrap().len(), 2);
280    }
281
282    #[test]
283    fn manifest_carries_embeddings_sha256_matching_bin() {
284        use sha2::{Digest, Sha256};
285        let tmp = tempdir().unwrap();
286        write_html(tmp.path(), "a.html", "<p>hash me</p>");
287        let c = ctx(tmp.path());
288        VectorSearchPlugin.after_compile(&c).unwrap();
289
290        let emb = fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
291        let mut h = Sha256::new();
292        h.update(&emb);
293        let expected =
294            h.finalize()
295                .iter()
296                .fold(String::with_capacity(64), |mut s, b| {
297                    use std::fmt::Write;
298                    let _ = write!(s, "{b:02x}");
299                    s
300                });
301
302        let json = fs::read_to_string(tmp.path().join("search/manifest.json"))
303            .unwrap();
304        let m: serde_json::Value = serde_json::from_str(&json).unwrap();
305        assert_eq!(m["embeddings_sha256"].as_str().unwrap(), expected);
306    }
307
308    #[test]
309    fn stamp_embeddings_hash_passes_through_invalid_json() {
310        let raw = b"not json".to_vec();
311        assert_eq!(stamp_embeddings_hash(&raw, b"x"), raw);
312        let arr = b"[1,2]".to_vec();
313        assert_eq!(stamp_embeddings_hash(&arr, b"x"), arr);
314    }
315
316    #[test]
317    fn truncate_chars_respects_unicode_boundaries() {
318        let s = "résumé café";
319        let t = truncate_chars(s, 6);
320        // 6 chars: "résumé" (no panic on multi-byte char)
321        assert_eq!(t, "résumé");
322    }
323
324    #[test]
325    fn truncate_chars_returns_full_string_when_short() {
326        let s = "short";
327        assert_eq!(truncate_chars(s, 100), "short");
328    }
329
330    #[test]
331    fn embeddings_byte_identical_across_runs() {
332        let tmp = tempdir().unwrap();
333        write_html(
334            tmp.path(),
335            "a.html",
336            "<p>identical content produces identical embeddings</p>",
337        );
338        let c = ctx(tmp.path());
339        VectorSearchPlugin.after_compile(&c).unwrap();
340        let first = fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
341        // Rebuild
342        VectorSearchPlugin.after_compile(&c).unwrap();
343        let second =
344            fs::read(tmp.path().join("search/embeddings.bin")).unwrap();
345        assert_eq!(first, second);
346    }
347
348    // -------------------------------------------------------------------
349    // IO error branches
350    // -------------------------------------------------------------------
351
352    #[test]
353    #[cfg(unix)]
354    fn after_compile_fails_when_site_has_unreadable_subdir() {
355        use std::os::unix::fs::PermissionsExt;
356        let tmp = tempdir().unwrap();
357        write_html(tmp.path(), "a.html", "<p>hello</p>");
358        let locked = tmp.path().join("locked");
359        fs::create_dir_all(&locked).unwrap();
360        fs::set_permissions(&locked, fs::Permissions::from_mode(0o000))
361            .unwrap();
362
363        let res = VectorSearchPlugin.after_compile(&ctx(tmp.path()));
364
365        let _ = fs::set_permissions(&locked, fs::Permissions::from_mode(0o755));
366        // Root CI runners bypass perms; only assert when it errored.
367        if let Err(e) = res {
368            assert!(!format!("{e}").is_empty());
369        }
370    }
371
372    #[test]
373    #[serial_test::parallel]
374    fn after_compile_fails_when_search_dir_squatted_by_file() {
375        let tmp = tempdir().unwrap();
376        write_html(tmp.path(), "a.html", "<p>hello</p>");
377        fs::write(tmp.path().join("search"), "not a dir").unwrap();
378        let err = VectorSearchPlugin
379            .after_compile(&ctx(tmp.path()))
380            .unwrap_err();
381        assert!(!format!("{err}").is_empty());
382    }
383
384    /// Squats `search/<name>` with a directory so the corresponding
385    /// `fs::write` fails.
386    fn assert_write_fails_when_squatted(name: &str) {
387        let tmp = tempdir().unwrap();
388        write_html(tmp.path(), "a.html", "<p>hello</p>");
389        fs::create_dir_all(tmp.path().join("search").join(name)).unwrap();
390        let err = VectorSearchPlugin
391            .after_compile(&ctx(tmp.path()))
392            .unwrap_err();
393        assert!(!format!("{err}").is_empty());
394    }
395
396    #[test]
397    fn after_compile_fails_when_embeddings_squatted_by_dir() {
398        assert_write_fails_when_squatted(EMBEDDINGS_FILE);
399    }
400
401    #[test]
402    fn after_compile_fails_when_manifest_squatted_by_dir() {
403        assert_write_fails_when_squatted(MANIFEST_FILE);
404    }
405
406    #[test]
407    fn after_compile_fails_when_model_squatted_by_dir() {
408        assert_write_fails_when_squatted(MODEL_FILE);
409    }
410
411    #[test]
412    fn after_compile_fails_when_tokenizer_squatted_by_dir() {
413        assert_write_fails_when_squatted(TOKENIZER_FILE);
414    }
415}
416
417#[cfg(all(test, feature = "test-fault-injection"))]
418mod fault_tests {
419    use super::*;
420    use serial_test::serial;
421
422    /// RAII guard that disables a failpoint on drop.
423    struct FailGuard(&'static str);
424
425    impl Drop for FailGuard {
426        fn drop(&mut self) {
427            let _ = fail::cfg(self.0, "off");
428        }
429    }
430
431    #[test]
432    #[serial]
433    fn stamp_embeddings_hash_falls_back_on_injected_serialize_failure() {
434        // Drives the defensive `unwrap_or_else` fallback that otherwise
435        // cannot be reached: re-serializing a manifest `Value` parsed
436        // from `ssg-search`'s own output cannot fail in practice.
437        let _guard = FailGuard("search_index::manifest-serialize");
438        fail::cfg("search_index::manifest-serialize", "return")
439            .expect("activate failpoint");
440
441        let manifest_json = br#"{"count":0,"entries":[]}"#;
442        let out = stamp_embeddings_hash(manifest_json, b"embeddings");
443        assert_eq!(
444            out, manifest_json,
445            "injected failure must fall back to the original manifest bytes"
446        );
447    }
448}
ssg/plugins/search_index.rs

ssg/plugins/
search_index.rs
