Verify extraction provenance and prune semantic staging before inference

crates/searchi-service/src/store/incremental.rs

@@ -181,6 +181,131 @@ impl Store {
         Ok(())
     }
 
+    /// Reconcile completed staging groups before any encoder call. A bounded
+    /// SQL snapshot preserves renamed donors (including swaps and duplicates)
+    /// while rebuilding only current compatible memberships. Active vectors
+    /// are never touched. Full range/vector validation still happens on reuse.
+    pub(crate) fn prune_before_inference(
+        &mut self,
+        generation: i64,
+        tasks: &[SourceTask],
+        revision: i64,
+    ) -> Result<()> {
+        ensure!(
+            tasks.len() <= MAX_SCAN_CHUNKS,
+            "Generation exceeds source task budget"
+        );
+        for task in tasks {
+            validate_task(task)?;
+        }
+        // The bounded snapshot can still contain 100K x 768d vectors. Keep it
+        // SQLite-backed on disk rather than turning that into a Rust/heap copy.
+        self.connection.pragma_update(None, "temp_store", "FILE")?;
+        let tx = self
+            .connection
+            .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
+        let current: i64 = tx.query_row(
+            "SELECT corpus_revision FROM semantic_meta WHERE singleton=1",
+            [],
+            |r| r.get(0),
+        )?;
+        ensure!(
+            current == revision,
+            "Indexed corpus changed during semantic discovery; retry required"
+        );
+        let active: bool = tx.query_row(
+            "SELECT active_generation=?1 FROM semantic_meta WHERE singleton=1",
+            [generation],
+            |r| r.get::<_, Option<bool>>(0).map(|v| v.unwrap_or(false)),
+        )?;
+        ensure!(!active, "Cannot prune the active semantic generation");
+        let existing: usize = tx.query_row(
+            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+            [generation],
+            |r| r.get(0),
+        )?;
+        ensure!(
+            existing <= MAX_SCAN_CHUNKS,
+            "Staging generation exceeds semantic chunk budget"
+        );
+        tx.execute_batch("CREATE TEMP TABLE IF NOT EXISTS semantic_discovery_tasks(task_key TEXT PRIMARY KEY,path TEXT NOT NULL,source_type TEXT NOT NULL,source_hash TEXT NOT NULL,title_hash TEXT NOT NULL,policy TEXT NOT NULL,passage_id INTEGER NOT NULL);
+            DELETE FROM semantic_discovery_tasks;
+            CREATE TEMP TABLE IF NOT EXISTS semantic_prune_tasks AS SELECT * FROM semantic_tasks WHERE 0;
+            CREATE TEMP TABLE IF NOT EXISTS semantic_prune_chunks AS SELECT * FROM semantic_chunks WHERE 0;
+            DELETE FROM semantic_prune_tasks; DELETE FROM semantic_prune_chunks;")?;
+        for task in tasks {
+            tx.execute(
+                "INSERT INTO semantic_discovery_tasks VALUES(?1,?2,?3,?4,?5,?6,?7)",
+                params![
+                    task.key,
+                    task.path,
+                    task.source_type,
+                    task.source_hash,
+                    task.title_hash,
+                    task.policy,
+                    task.passage_id
+                ],
+            )?;
+        }
+        tx.execute("INSERT INTO semantic_prune_tasks SELECT * FROM semantic_tasks t WHERE generation=?1 AND state='done' AND source_hash!='' AND expected_chunks IS NOT NULL AND expected_chunks=(SELECT count(*) FROM semantic_chunks c WHERE c.generation=t.generation AND c.task_key=t.task_key)",[generation])?;
+        tx.execute(
+            "INSERT INTO semantic_prune_chunks SELECT * FROM semantic_chunks WHERE generation=?1",
+            [generation],
+        )?;
+        tx.execute_batch("CREATE UNIQUE INDEX IF NOT EXISTS temp.semantic_prune_task_key ON semantic_prune_tasks(task_key);
+            CREATE INDEX IF NOT EXISTS temp.semantic_prune_reuse ON semantic_prune_tasks(source_hash,title_hash,policy,source_type);
+            CREATE INDEX IF NOT EXISTS temp.semantic_prune_membership ON semantic_prune_chunks(task_key);")?;
+        let mut mappings = Vec::new();
+        {
+            let mut statement = tx.prepare("SELECT d.task_key,COALESCE((SELECT old.task_key FROM semantic_prune_tasks old WHERE old.task_key=d.task_key AND old.source_hash=d.source_hash AND old.title_hash=d.title_hash AND old.policy=d.policy AND old.source_type=d.source_type),(SELECT old.task_key FROM semantic_prune_tasks old WHERE old.source_hash=d.source_hash AND old.title_hash=d.title_hash AND old.policy=d.policy AND old.source_type=d.source_type ORDER BY old.task_key LIMIT 1)) FROM semantic_discovery_tasks d")?;
+            let rows = statement.query_map([], |r| {
+                Ok((r.get::<_, String>(0)?, r.get::<_, Option<String>>(1)?))
+            })?;
+            for row in rows {
+                let (key, donor) = row?;
+                if let Some(donor) = donor {
+                    mappings.push((key, donor));
+                }
+            }
+        }
+        let mut total = 0usize;
+        for (_, donor) in &mappings {
+            let count: usize = tx.query_row(
+                "SELECT expected_chunks FROM semantic_prune_tasks WHERE task_key=?1",
+                [donor],
+                |r| r.get(0),
+            )?;
+            total = total
+                .checked_add(count)
+                .context("Semantic chunk count overflow")?;
+            ensure!(
+                total <= MAX_SCAN_CHUNKS,
+                "Current completed sources exceed semantic chunk budget"
+            );
+        }
+        tx.execute(
+            "DELETE FROM semantic_chunks WHERE generation=?1",
+            [generation],
+        )?;
+        tx.execute(
+            "DELETE FROM semantic_tasks WHERE generation=?1",
+            [generation],
+        )?;
+        for (key, donor) in mappings {
+            tx.execute("INSERT INTO semantic_tasks(generation,task_key,path,source_type,source_hash,title_hash,policy,passage_id,state,attempts,expected_chunks,message)
+                SELECT ?1,d.task_key,d.path,d.source_type,d.source_hash,d.title_hash,d.policy,d.passage_id,'done',old.attempts,old.expected_chunks,NULL FROM semantic_discovery_tasks d JOIN semantic_prune_tasks old ON old.task_key=?3 WHERE d.task_key=?2",params![generation,key,donor])?;
+            tx.execute("INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector,source_type,file_size,file_mtime,task_key)
+                SELECT ?1,d.passage_id,c.ordinal,d.path,c.kind,c.number,c.content_hash,c.start_byte,c.end_byte,c.snippet,c.vector,c.source_type,c.file_size,c.file_mtime,d.task_key FROM semantic_prune_chunks c JOIN semantic_discovery_tasks d ON d.task_key=?2 WHERE c.task_key=?3",params![generation,key,donor])?;
+        }
+        tx.execute(
+            "DELETE FROM semantic_media_errors WHERE generation=?1",
+            [generation],
+        )?;
+        tx.execute_batch("DELETE FROM semantic_prune_tasks; DELETE FROM semantic_prune_chunks; DELETE FROM semantic_discovery_tasks;")?;
+        tx.commit()?;
+        Ok(())
+    }
+
     pub(crate) fn reconcile_tasks(
         &mut self,
         generation: i64,
@@ -221,9 +346,21 @@ impl Store {
         Ok(())
     }
 
+    #[cfg(test)]
     pub(crate) fn activate_checked(&mut self, generation: i64, revision: i64) -> Result<()> {
+        self.activate_checked_with_validation(generation, revision, || Ok(()))
+    }
+
+    /// Validate live sources after graph preparation, outside the SQLite write
+    /// transaction. Rejection preserves both the active reader and staging.
+    pub(crate) fn activate_checked_with_validation(
+        &mut self,
+        generation: i64,
+        revision: i64,
+        validate_sources: impl FnOnce() -> Result<()>,
+    ) -> Result<()> {
         validate_completion(&self.connection, generation)?;
-        self.activate_revision(generation, Some(revision))
+        self.activate_revision(generation, Some(revision), validate_sources)
     }
 }
 
@@ -468,6 +605,107 @@ mod tests {
         store.complete_source(g, t, 1)
     }
 
+    #[test]
+    fn pre_inference_pruning_preserves_renamed_donors_at_the_chunk_ceiling() -> Result<()> {
+        // Actual 100K membership boundary; vectors are storage-only fixtures.
+        let dir = tempfile::tempdir()?;
+        let root = dir.path().join("docs");
+        std::fs::create_dir(&root)?;
+        std::fs::write(root.join("notes.txt"), "complete source")?;
+        let db = dir.path().join("index.db");
+        let mut lexical = SearchIndex::open(&db)?;
+        lexical.index_folder(&root)?;
+        let p = lexical.passage_batch(0, 1)?.remove(0);
+        let mut store = Store::open(&db, 2)?;
+        let lock = store.build_lock()?;
+        let active = store.start_or_resume("prune-policy", &lock)?.0;
+        insert_complete(&mut store, active, &p, &task(&p))?;
+        store.activate_checked(active, store.revision()?)?;
+        let staging = store.start_or_resume("prune-policy", &lock)?.0;
+        let original = task(&p);
+        insert_complete(&mut store, staging, &p, &original)?;
+        store.connection.execute("WITH RECURSIVE ord(n) AS (SELECT 1 UNION ALL SELECT n+1 FROM ord WHERE n+1<?2)
+            INSERT INTO semantic_chunks(generation,passage_id,ordinal,path,kind,number,content_hash,start_byte,end_byte,snippet,vector,source_type,file_size,file_mtime,task_key)
+            SELECT generation,passage_id,n,path,kind,number,content_hash,start_byte,end_byte,snippet,vector,source_type,file_size,file_mtime,task_key FROM semantic_chunks CROSS JOIN ord WHERE generation=?1 AND ordinal=0",params![staging,MAX_SCAN_CHUNKS as i64])?;
+        store.complete_source(staging, &original, MAX_SCAN_CHUNKS)?;
+        let moved_dir = root.join("moved");
+        std::fs::create_dir(&moved_dir)?;
+        std::fs::rename(root.join("notes.txt"), moved_dir.join("notes.txt"))?;
+        lexical.index_folder(&root)?;
+        let moved = lexical.passage_batch(0, 1)?.remove(0);
+        let moved_task = task(&moved);
+        let revision = store.revision()?;
+        assert!(store
+            .prune_before_inference(staging, std::slice::from_ref(&moved_task), revision - 1)
+            .is_err());
+        let old_count: usize = store.connection.query_row(
+            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+            [staging],
+            |r| r.get(0),
+        )?;
+        assert_eq!(
+            old_count, MAX_SCAN_CHUNKS,
+            "Revision failure must preserve staging"
+        );
+        store.prune_before_inference(staging, std::slice::from_ref(&moved_task), revision)?;
+        store.stage_task(staging, &moved_task)?;
+        assert_eq!(
+            store.reuse_text_source(staging, &moved_task, &moved)?,
+            Some(MAX_SCAN_CHUNKS)
+        );
+        let (count, old_paths): (usize, usize) = store.connection.query_row(
+            "SELECT count(*),sum(path!=?2) FROM semantic_chunks WHERE generation=?1",
+            params![staging, moved.path],
+            |r| Ok((r.get(0)?, r.get(1)?)),
+        )?;
+        assert_eq!((count, old_paths), (MAX_SCAN_CHUNKS, 0));
+        let mut duplicate = moved_task.clone();
+        duplicate.key.push_str("-another-path");
+        duplicate.path.push_str("-another-path");
+        assert!(
+            store
+                .prune_before_inference(staging, &[moved_task.clone(), duplicate], revision)
+                .is_err(),
+            "True current membership over budget must reject atomically"
+        );
+        let count: usize = store.connection.query_row(
+            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+            [staging],
+            |r| r.get(0),
+        )?;
+        assert_eq!(count, MAX_SCAN_CHUNKS);
+        std::fs::write(
+            moved_dir.join("notes.txt"),
+            "replacement source with new content",
+        )?;
+        lexical.index_folder(&root)?;
+        let replacement = lexical.passage_batch(0, 1)?.remove(0);
+        let replacement_task = task(&replacement);
+        store.prune_before_inference(
+            staging,
+            std::slice::from_ref(&replacement_task),
+            store.revision()?,
+        )?;
+        let count: usize = store.connection.query_row(
+            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+            [staging],
+            |r| r.get(0),
+        )?;
+        assert_eq!(
+            count, 0,
+            "Obsolete ceiling-sized group must be removed before replacement inference"
+        );
+        insert_complete(&mut store, staging, &replacement, &replacement_task)?;
+        assert_eq!(store.active_fingerprint()?.as_deref(), Some("prune-policy"));
+        let active_count: usize = store.connection.query_row(
+            "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+            [active],
+            |r| r.get(0),
+        )?;
+        assert_eq!(active_count, 1, "Pruning must never remove active vectors");
+        Ok(())
+    }
+
     #[test]
     fn resume_retains_completed_groups_and_discards_unfinished_vectors() -> Result<()> {
         let dir = tempfile::tempdir()?;
@@ -583,6 +821,59 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn late_source_validation_rejects_publication_without_holding_write_lock() -> Result<()> {
+        let dir = tempfile::tempdir()?;
+        let root = dir.path().join("docs");
+        std::fs::create_dir(&root)?;
+        std::fs::write(root.join("notes.txt"), "source identity before graph build")?;
+        let db = dir.path().join("index.db");
+        let mut lexical = SearchIndex::open(&db)?;
+        lexical.index_folder(&root)?;
+        let passage = lexical.passage_batch(0, 1)?.remove(0);
+        let source = task(&passage);
+        let mut store = Store::open(&db, 2)?;
+        let lock = store.build_lock()?;
+        let (active, _) = store.start_or_resume("old-policy", &lock)?;
+        insert_complete(&mut store, active, &passage, &source)?;
+        store.activate_checked(active, store.revision()?)?;
+        let (staging, _) = store.start_or_resume("new-policy", &lock)?;
+        insert_complete(&mut store, staging, &passage, &source)?;
+        let revision = store.revision()?;
+        let mut called = false;
+        let error = store
+            .activate_checked_with_validation(staging, revision, || {
+                called = true;
+                // An independent write succeeds only when validation runs outside
+                // the publishing connection's Immediate transaction.
+                let other = rusqlite::Connection::open(&db)?;
+                other.busy_timeout(std::time::Duration::ZERO)?;
+                other.execute_batch("BEGIN IMMEDIATE; COMMIT;")?;
+                std::fs::write(root.join("notes.txt"), "direct filesystem mutation")?;
+                anyhow::bail!("source changed after graph preparation")
+            })
+            .unwrap_err();
+        assert!(called);
+        assert!(error
+            .to_string()
+            .contains("source changed after graph preparation"));
+        assert_eq!(store.active_fingerprint()?.as_deref(), Some("old-policy"));
+        assert_eq!(
+            store.ann_cache.as_ref().map(|(generation, _)| *generation),
+            Some(active)
+        );
+        assert_eq!(store.start_or_resume("new-policy", &lock)?, (staging, true));
+        assert_eq!(
+            store.connection.query_row(
+                "SELECT count(*) FROM semantic_chunks WHERE generation=?1",
+                [staging],
+                |row| row.get::<_, usize>(0),
+            )?,
+            1
+        );
+        Ok(())
+    }
+
     #[test]
     fn revision_race_cannot_publish_and_failed_activation_keeps_staging() -> Result<()> {
         let dir = tempfile::tempdir()?;