searchi/crates/searchi-service/src/pipeline.rs

//! Durable bounded semantic source tasks. Only completed source groups may
//! supply vectors; tokenizer/inference work stays outside SQLite transactions.
use super::*;
use searchi_core::IndexedPassage;
use store::{content_hash, SourceTask};

pub(super) struct BuildWork {
    pub report: EmbeddingReport,
    pub keys: Vec<String>,
    pub media: Vec<searchi_media::MediaIdentity>,
    documents: HashMap<String, searchi_media::MediaIdentity>,
}

impl BuildWork {
    pub fn track(&mut self, task: &SourceTask) -> Result<()> {
        ensure!(
            self.keys.len() < MAX_SCAN_CHUNKS,
            "Semantic source task budget exceeded"
        );
        self.keys.push(task.key.clone());
        Ok(())
    }
}

pub(super) fn task_key(parts: &[&str]) -> String {
    let mut input = String::new();
    for part in parts {
        input.push_str(&format!("{}:", part.len()));
        input.push_str(part);
    }
    content_hash(&input)
}

pub(super) fn text_task(passage: &IndexedPassage, policy: &str) -> Result<SourceTask> {
    let key = task_key(&[
        "text",
        &passage.path,
        &passage.source.kind,
        &passage.source.number.to_string(),
    ]);
    Ok(SourceTask {
        key,
        path: passage.path.clone(),
        source_type: "text".into(),
        source_hash: content_hash(&passage.text),
        title_hash: content_hash(bounded_title(&passage.filename)),
        policy: policy_base(policy).into(),
        passage_id: passage.id,
    })
}

impl HybridSearch {
    pub(super) fn rebuild_generation(
        &mut self,
        fingerprint: String,
        media_limits: Option<&searchi_media::MediaLimits>,
        build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
    ) -> Result<EmbeddingReport> {
        self.rebuild_generation_guarded(fingerprint, media_limits, false, build)?
            .ok_or_else(|| anyhow::anyhow!("Explicit semantic rebuild was unexpectedly deferred"))
    }

    pub(super) fn rebuild_generation_preserving_media(
        &mut self,
        fingerprint: String,
        build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
    ) -> Result<Option<EmbeddingReport>> {
        self.rebuild_generation_guarded(fingerprint, None, true, build)
    }

    fn rebuild_generation_guarded(
        &mut self,
        fingerprint: String,
        media_limits: Option<&searchi_media::MediaLimits>,
        preserve_full_media: bool,
        build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
    ) -> Result<Option<EmbeddingReport>> {
        let lock = self.store.build_lock()?;
        if preserve_full_media && self.store.has_compatible_full_generation(&fingerprint)? {
            return Ok(None);
        }
        let (generation, resumed) = self.store.start_or_resume(&fingerprint, &lock)?;
        // This revision covers the entire discovery/build/finalization pass.
        // A concurrent lexical update aborts publication; next invocation
        // reconciles and resumes all still-compatible completed groups.
        let revision = self.store.revision()?;
        let previous_fingerprint = std::mem::replace(&mut self.generation_fingerprint, fingerprint);
        let mut work = BuildWork {
            report: EmbeddingReport {
                media_files: 0,
                media_failures: 0,
                failures: Vec::new(),
                passages: 0,
                chunks: 0,
                generation,
                model_fingerprint: self.generation_fingerprint.clone(),
                reused_sources: 0,
                reused_chunks: 0,
                encoded_sources: 0,
                document_inference_calls: 0,
                image_inference_calls: 0,
                audio_inference_calls: 0,
                resumed,
            },
            keys: Vec::new(),
            media: Vec::new(),
            documents: HashMap::new(),
        };
        let result = self
            .discover_tasks(media_limits)
            .and_then(|tasks| {
                self.store
                    .prune_before_inference(generation, &tasks, revision)?;
                build(self, generation, &mut work)
            })
            .and_then(|()| {
                self.store
                    .reconcile_tasks(generation, &work.keys, revision)?;
                self.store
                    .activate_checked_with_validation(generation, revision, || {
                        for identity in work.media.iter().chain(work.documents.values()) {
                            // Full bytes, not just stat metadata: in-place changes can
                            // preserve size and timestamps, including on mounted volumes.
                            let current =
                                searchi_media::MediaIdentity::read(&identity.path, identity.bytes)?;
                            ensure!(
                                current.sha256 == identity.sha256
                                    && current.modified_ns == identity.modified_ns,
                                "Source changed during semantic indexing; retry required"
                            );
                        }
                        Ok(())
                    })?;
                Ok(Some(work.report))
            });
        if result.is_err() {
            self.generation_fingerprint = previous_fingerprint;
        }
        result
    }

    /// Discover source identities in keyset batches before spending the encoder
    /// budget. This extra byte-verification pass enables safe staging pruning;
    /// it is indexing I/O, not a foreground latency optimization.
    fn discover_tasks(
        &self,
        media_limits: Option<&searchi_media::MediaLimits>,
    ) -> Result<Vec<SourceTask>> {
        let mut tasks = Vec::new();
        let mut after = 0;
        loop {
            let passages = self.lexical.passage_batch(after, 64)?;
            if passages.is_empty() {
                break;
            }
            for passage in passages {
                after = passage.id;
                if passage.text.trim().is_empty() {
                    continue;
                }
                ensure!(
                    tasks.len() < MAX_SCAN_CHUNKS,
                    "Semantic source task budget exceeded"
                );
                tasks.push(text_task(&passage, &self.generation_fingerprint)?);
            }
        }
        if let Some(limits) = media_limits {
            let mut after = 0;
            let mut count = 0;
            loop {
                let files = self.store.media_batch(after)?;
                if files.is_empty() {
                    break;
                }
                for (id, path, _, _) in files {
                    after = id;
                    let Some(kind) = searchi_media::classify(Path::new(&path)) else {
                        continue;
                    };
                    count += 1;
                    ensure!(
                        count <= 10000,
                        "Media generation exceeds 10000-file resource budget"
                    );
                    ensure!(
                        tasks.len() < MAX_SCAN_CHUNKS,
                        "Semantic source task budget exceeded"
                    );
                    let identity = searchi_media::MediaIdentity::read(
                        Path::new(&path),
                        if kind == searchi_media::MediaKind::Image {
                            64 * 1024 * 1024
                        } else {
                            limits.max_file_bytes
                        },
                    );
                    tasks.push(SourceTask {
                        key: task_key(&["media", &path]),
                        path,
                        source_type: "media".into(),
                        source_hash: identity.map(|identity| identity.sha256).unwrap_or_default(),
                        title_hash: String::new(),
                        policy: self.generation_fingerprint.clone(),
                        passage_id: -id,
                    });
                }
            }
        }
        Ok(tasks)
    }

    pub(super) fn build(&mut self, generation: i64, work: &mut BuildWork) -> Result<()> {
        let mut after_id = 0;
        loop {
            let passages = self.lexical.passage_batch(after_id, 64)?;
            if passages.is_empty() {
                break;
            }
            for passage in passages {
                after_id = passage.id;
                if passage.text.trim().is_empty() {
                    continue;
                }
                if !work.documents.contains_key(&passage.path) {
                    ensure!(
                        work.documents.len() < MAX_SCAN_CHUNKS,
                        "Document source budget exceeded"
                    );
                    let identity = searchi_media::MediaIdentity::read(
                        Path::new(&passage.path),
                        32 * 1024 * 1024,
                    )?;
                    let extracted = self.lexical.extracted_source_hash(&passage.path)?;
                    ensure!(extracted.as_deref() == Some(identity.sha256.as_str()),
                        "Document bytes differ from extracted passages; reindex the folder before semantic indexing");
                    work.documents.insert(passage.path.clone(), identity);
                }
                let task = text_task(&passage, &self.generation_fingerprint)?;
                work.track(&task)?;
                self.store.stage_task(generation, &task)?;
                let result = self.build_text_source(generation, &task, &passage, work);
                if let Err(error) = result {
                    if let Err(cleanup) =
                        self.store
                            .fail_source(generation, &task, &format!("{error:#}"))
                    {
                        return Err(error.context(format!(
                            "Failed to discard incomplete text source: {cleanup:#}"
                        )));
                    }
                    return Err(error);
                }
            }
        }
        Ok(())
    }

    fn build_text_source(
        &mut self,
        generation: i64,
        task: &SourceTask,
        passage: &IndexedPassage,
        work: &mut BuildWork,
    ) -> Result<()> {
        if let Some(count) = self.store.reuse_text_source(generation, task, passage)? {
            ensure!(
                work.report.chunks + count <= MAX_SCAN_CHUNKS,
                "Corpus exceeds ANN chunk budget"
            );
            work.report.chunks += count;
            work.report.passages += 1;
            work.report.reused_chunks += count;
            work.report.reused_sources += 1;
            return Ok(());
        }
        let title = bounded_title(&passage.filename);
        let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
        ensure!(
            work.report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
            "Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
        );
        let batch_limit = self.embedder.document_batch_limit().clamp(1, MAX_BATCH);
        for (batch_index, batch) in windows.chunks(batch_limit).enumerate() {
            let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
            work.report.document_inference_calls += 1;
            let vectors = self.embedder.embed_documents(&texts, title)?;
            self.store.insert(
                generation,
                passage,
                batch_index * batch_limit,
                batch,
                &vectors,
            )?;
        }
        self.store
            .complete_source(generation, task, windows.len())?;
        work.report.chunks += windows.len();
        work.report.passages += 1;
        work.report.encoded_sources += 1;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Fault-only provider: never returns a vector or claims semantic quality.
    struct RejectInference;
    impl EmbeddingProvider for RejectInference {
        fn fingerprint(&self) -> &str {
            "fault-injection-only"
        }
        fn dimensions(&self) -> usize {
            GEMMA_DIMENSIONS
        }
        fn split_windows(&mut self, text: &str, _: &str) -> Result<Vec<TextWindow>> {
            Ok(vec![TextWindow {
                text: text.into(),
                start_byte: 0,
                end_byte: text.len(),
            }])
        }
        fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
            anyhow::bail!("injected primary inference failure")
        }
        fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
            anyhow::bail!("fault provider never supplies query vectors")
        }
    }

    #[test]
    fn cleanup_failure_keeps_primary_error_reader_policy_and_staging() -> Result<()> {
        let temp = tempfile::tempdir()?;
        let root = temp.path().join("corpus");
        std::fs::create_dir(&root)?;
        std::fs::write(
            root.join("notes.txt"),
            "A real source whose inference is deliberately rejected.",
        )?;
        let db = temp.path().join("index.db");
        SearchIndex::open(&db)?.index_folder(&root)?;
        let mut service = HybridSearch::with_provider(&db, Box::new(RejectInference))?;
        let prior_policy = service.generation_fingerprint.clone();
        let connection = rusqlite::Connection::open(&db)?;
        connection.execute_batch("CREATE TRIGGER reject_task_cleanup BEFORE UPDATE OF state ON semantic_tasks WHEN new.state='failed' BEGIN SELECT RAISE(ABORT,'injected task cleanup failure'); END;")?;
        let error = service.index_embeddings().unwrap_err();
        let message = format!("{error:#}");
        assert!(
            message.contains("injected primary inference failure"),
            "{message}"
        );
        assert!(
            message.contains("injected task cleanup failure"),
            "{message}"
        );
        assert_eq!(service.generation_fingerprint, prior_policy);
        let staging: usize =
            connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
                r.get(0)
            })?;
        assert_eq!(staging, 1);
        let state: String =
            connection.query_row("SELECT state FROM semantic_tasks", [], |r| r.get(0))?;
        assert_eq!(
            state, "running",
            "Failed cleanup transaction must roll back"
        );
        connection.execute_batch("DROP TRIGGER reject_task_cleanup;")?;
        let error = service.index_embeddings().unwrap_err();
        assert!(format!("{error:#}").contains("injected primary inference failure"));
        assert_eq!(service.generation_fingerprint, prior_policy);
        let staging: usize =
            connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
                r.get(0)
            })?;
        assert_eq!(staging, 1, "Resume must retain the same staging generation");
        Ok(())
    }
    #[test]
    fn discovery_and_encoding_use_the_same_normalized_text_component_identity() -> Result<()> {
        let temp = tempfile::tempdir()?;
        let root = temp.path().join("docs");
        std::fs::create_dir(&root)?;
        std::fs::write(
            root.join("notes.txt"),
            "Exact source bytes before a deliberately failed encoder.",
        )?;
        let db = temp.path().join("index.db");
        SearchIndex::open(&db)?.index_folder(&root)?;
        let mut service = HybridSearch::with_provider(&db, Box::new(RejectInference))?;
        let base = generation_fingerprint(service.embedder.fingerprint());
        service.generation_fingerprint = format!("{base}:text-only");
        let discovered = service.discover_tasks(None)?.remove(0);
        assert_eq!(discovered.policy, base);
        let error = service
            .index_text_embeddings_preserving_media()
            .unwrap_err();
        assert!(format!("{error:#}").contains("injected primary inference failure"));
        let connection = rusqlite::Connection::open(&db)?;
        let encoded: (String, String, String, String) = connection.query_row(
            "SELECT task_key,source_hash,title_hash,policy FROM semantic_tasks",
            [],
            |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
        )?;
        assert_eq!(
            encoded,
            (
                discovered.key,
                discovered.source_hash,
                discovered.title_hash,
                discovered.policy
            )
        );
        Ok(())
    }
}