fix: verify complete offline assets before replacing staged runtime

tools/embeddinggemma2/stage_runtime.py

@@ -9,6 +9,20 @@ import shutil
 ROOT = Path(__file__).resolve().parents[2]
 
 
+def verify_bundle(root, manifest):
+    """Verify the complete offline model pair before modifying staged resources."""
+    artifacts = []
+    for spec in manifest["files"]:
+        artifact = root / ".searchi/models/embeddinggemma2" / spec["path"]
+        if artifact.stat().st_size != spec["bytes"]:
+            raise RuntimeError("Model size mismatch")
+        with artifact.open("rb") as stream:
+            if hashlib.file_digest(stream, "sha256").hexdigest() != spec["sha256"]:
+                raise RuntimeError("Model checksum mismatch")
+        artifacts.append(artifact)
+    return artifacts
+
+
 def main():
     parser = argparse.ArgumentParser(description="Stage Searchi runtime resources")
     parser.add_argument("--backend", choices=["cpu", "auto", "metal", "vulkan", "cuda"], default="cpu")
@@ -21,23 +35,23 @@ def main():
     source /= name
     if not source.is_file():
         raise RuntimeError("Build the pinned native worker before staging")
+    license_path = ROOT / ".searchi/references/llama.cpp/LICENSE"
+    if not license_path.is_file():
+        raise RuntimeError("Native worker license is missing")
+    artifacts = []
+    if args.bundle_model:
+        manifest = json.loads(Path(__file__).with_name("model-manifest.json").read_text())
+        artifacts = verify_bundle(ROOT, manifest)
     destination = ROOT / "apps/desktop/src-tauri/resources"
     destination.mkdir(parents=True, exist_ok=True)
     shutil.copy2(source, destination / name)
     # Include the upstream runtime license with the statically linked helper.
-    shutil.copy2(ROOT / ".searchi/references/llama.cpp/LICENSE", destination / "LLAMA-LICENSE.txt")
+    shutil.copy2(license_path, destination / "LLAMA-LICENSE.txt")
     if args.bundle_model:
-        manifest = json.loads(Path(__file__).with_name("model-manifest.json").read_text())
         model_destination = destination / "embeddinggemma2"
         model_destination.mkdir(exist_ok=True)
-        for spec in manifest["files"]:
-            artifact = ROOT / ".searchi/models/embeddinggemma2" / spec["path"]
-            if artifact.stat().st_size != spec["bytes"]:
-                raise RuntimeError("Model size mismatch")
-            with artifact.open("rb") as stream:
-                if hashlib.file_digest(stream, "sha256").hexdigest() != spec["sha256"]:
-                    raise RuntimeError("Model checksum mismatch")
-            shutil.copy2(artifact, model_destination / spec["path"])
+        for artifact in artifacts:
+            shutil.copy2(artifact, model_destination / artifact.name)
     print(f"Staged native worker: {destination / name}")
 
 

tools/embeddinggemma2/test_stage_runtime.py

@@ -0,0 +1,48 @@
+"""Offline installation regressions; no model downloads or inference."""
+import hashlib
+import os
+from pathlib import Path
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import stage_runtime
+
+
+class StagingTests(unittest.TestCase):
+    def test_corrupt_second_model_preserves_existing_staged_worker(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            name = "searchi-gemma-worker.exe" if os.name == "nt" else "searchi-gemma-worker"
+            runtime = root / ".searchi/runtime/cpu"
+            if os.name == "nt":
+                runtime /= "Release"
+            runtime.mkdir(parents=True)
+            (runtime / name).write_bytes(b"new worker")
+            license_path = root / ".searchi/references/llama.cpp/LICENSE"
+            license_path.parent.mkdir(parents=True)
+            license_path.write_text("license")
+            destination = root / "apps/desktop/src-tauri/resources"
+            destination.mkdir(parents=True)
+            (destination / name).write_bytes(b"previous worker")
+            models = root / ".searchi/models/embeddinggemma2"
+            models.mkdir(parents=True)
+            specs = []
+            for filename, content in [("model.gguf", b"verified"), ("projector.gguf", b"original")]:
+                (models / filename).write_bytes(content)
+                specs.append({"path": filename, "bytes": len(content), "sha256": hashlib.sha256(content).hexdigest()})
+            (models / "projector.gguf").write_bytes(b"corrupt!")
+            with patch.object(stage_runtime, "ROOT", root), patch("sys.argv", ["stage_runtime", "--bundle-model"]), patch.object(stage_runtime.json, "loads", return_value={"files": specs}):
+                with self.assertRaisesRegex(RuntimeError, "checksum mismatch"):
+                    stage_runtime.main()
+            self.assertEqual((destination / name).read_bytes(), b"previous worker")
+            self.assertFalse((destination / "embeddinggemma2").exists())
+
+    def test_missing_model_rejects_complete_bundle(self):
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaises(FileNotFoundError):
+                stage_runtime.verify_bundle(Path(directory), {"files": [{"path": "missing.gguf", "bytes": 1, "sha256": "unused"}]})
+
+
+if __name__ == "__main__":
+    unittest.main()