feat(retrieval): preload Qwen models during setup

adapters/plugins.tsv

@@ -1,7 +1,7 @@
 name	claude_marketplace	claude_plugin	codex_marketplace	codex_plugin	codex_method	codex_source	codex_skill	mcporter_name	mcporter_url
 Ponytail	DietrichGebert/ponytail	ponytail@ponytail	DietrichGebert/ponytail	ponytail@ponytail	plugin	-	-
 i-have-adhd	ayghri/i-have-adhd	i-have-adhd@i-have-adhd	ayghri/i-have-adhd	i-have-adhd@i-have-adhd	plugin	-	-
-Superpowers	-	superpowers@claude-plugins-official	openai-curated-remote	superpowers@openai-curated-remote	plugin	-	-
+Superpowers	anthropics/claude-plugins-official	superpowers@claude-plugins-official	openai-curated-remote	superpowers@openai-curated-remote	plugin	-	-
 Context7	upstash/context7	context7@context7-marketplace	upstash/context7	context7@context7-marketplace	plugin	-	-	context7	https://mcp.context7.com/mcp
 Caveman	JuliusBrussee/caveman	caveman@caveman	thinkhome-org/caveman-codex	caveman@thinkhome-caveman	plugin	-	-
 QMD	tobi/qmd	qmd@qmd	-	-	qmd	@tobilu/qmd	qmd

scripts/lib/semantic-retrieval.py

@@ -10,6 +10,7 @@ import sys
 
 FILES = ('server.py', 'benchmark.py', 'requirements.txt')
 CUDA_INDEX = 'https://download.pytorch.org/whl/cu126'
+MODEL_PRELOAD_TIMEOUT = 900
 
 
 def digest(path):
@@ -41,6 +42,11 @@ def ensure_runtime(target, dry_run):
     return python
 
 
+def ensure_models(python, target):
+    subprocess.run([str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')],
+                   check=True, timeout=MODEL_PRELOAD_TIMEOUT)
+
+
 def load_state(path):
     if not path.exists():
         return {'version': 1, 'clients': [], 'files': {}}
@@ -96,7 +102,9 @@ def sync(root, home, clients, enabled, dry_run, update, summary=False):
     for name in FILES:
         shutil.copy2(source / name, target / name)
         state['files'][name] = digest(target / name)
+    save_state(state_path, state, False)
     python = ensure_runtime(target, False)
+    ensure_models(python, target)
     state['clients'] = list(dict.fromkeys(state['clients']))
     save_state(state_path, state, False)
     for client in clients:

scripts/tests/test-managed-extensions.py

@@ -1,4 +1,5 @@
 import contextlib
+import csv
 import importlib.util
 import io
 import json
@@ -11,7 +12,8 @@ import unittest
 from unittest.mock import patch
 
 
-spec = importlib.util.spec_from_file_location('extensions', Path(__file__).parents[2] / 'scripts/lib/managed-extensions.py')
+root = Path(__file__).parents[2]
+spec = importlib.util.spec_from_file_location('extensions', root / 'scripts/lib/managed-extensions.py')
 extensions = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(extensions)
 
@@ -184,11 +186,12 @@ class ManagedExtensionTests(unittest.TestCase):
             ['claude', 'plugin', 'enable', 'example@test'],
         ], self.commands)
 
-    def test_claude_refreshes_marketplace_before_install(self):
-        entry = plugin('Superpowers')
-        entry['claude_plugin'] = 'superpowers@claude-plugins-official'
+    def test_superpowers_adds_official_claude_marketplace_before_install(self):
+        with (root / 'adapters/plugins.tsv').open(encoding='utf-8-sig', newline='') as stream:
+            entry = next(entry for entry in csv.DictReader(stream, delimiter='\t') if entry['name'] == 'Superpowers')
         self.manager().sync([entry], ['claude'], {'Superpowers'})
         self.assertEqual([
+            ['claude', 'plugin', 'marketplace', 'add', 'anthropics/claude-plugins-official'],
             ['claude', 'plugin', 'marketplace', 'update', 'claude-plugins-official'],
             ['claude', 'plugin', 'install', 'superpowers@claude-plugins-official', '--scope', 'user'],
         ], self.commands)

scripts/tests/test-retrieval-server.py

@@ -82,6 +82,25 @@ class RetrievalServerTests(unittest.TestCase):
         with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))):
             self.assertEqual(MODULE.preferred_device(), 'cpu')
 
+    def test_preload_models_warms_embedding_and_reranking_from_model_cache(self):
+        with TemporaryDirectory() as temporary:
+            cache = Path(temporary) / 'model-cache'
+            calls = []
+            encoder = SimpleNamespace(encode=lambda values, **kwargs: calls.append(('encode', values, kwargs)) or [embedding()])
+            reranker = SimpleNamespace(predict=lambda pairs, **kwargs: calls.append(('predict', pairs, kwargs)) or [0.0])
+
+            with patch.object(MODULE, 'shared_encoder', return_value=encoder), \
+                 patch.object(MODULE, 'shared_reranker', return_value=reranker), \
+                 patch.dict(os.environ, {'HF_HOME': 'other-cache'}):
+                getattr(MODULE, 'preload_models', lambda _: None)(cache)
+                self.assertEqual(str(cache.resolve()), os.environ['HF_HOME'])
+
+            self.assertTrue(cache.is_dir())
+            self.assertEqual([
+                ('encode', ['Verify semantic retrieval models.'], {'prompt_name': 'query'}),
+                ('predict', [('Verify semantic retrieval models.', 'Verify semantic retrieval models.')], {'show_progress_bar': False}),
+            ], calls)
+
     @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.')
     def test_real_models(self):
         if not huggingface_available():

scripts/tests/test-semantic-retrieval.py

@@ -48,6 +48,19 @@ class SemanticRetrievalTests(unittest.TestCase):
             self.assertEqual(run.call_count, 2)
             self.assertNotIn('--extra-index-url', run.call_args.args[0])
 
+    def test_sync_preloads_models_after_installing_runtime(self):
+        with tempfile.TemporaryDirectory() as directory:
+            home = Path(directory) / 'home'
+            target = home / '.my-ai-configuration/semantic-retrieval'
+            python = MODULE.python_path(target / '.venv')
+            with patch.object(MODULE, 'ensure_runtime', return_value=python), \
+                 patch.object(MODULE.subprocess, 'run') as run:
+                MODULE.sync(ROOT, home, ['codex'], True, False, False)
+
+            self.assertEqual([[
+                str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')
+            ]], [call.args[0] for call in run.call_args_list])
+
     def test_summary_dry_run_omits_details_without_writes(self):
         with tempfile.TemporaryDirectory() as directory:
             home = Path(directory) / 'home'

shared/retrieval/server.py

@@ -201,6 +201,18 @@ def shared_reranker():
         return _shared_reranker
 
 
+def preload_models(model_cache: Path):
+    model_cache.mkdir(parents=True, exist_ok=True)
+    os.environ["HF_HOME"] = str(model_cache.resolve())
+    query = "Verify semantic retrieval models."
+    embeddings = shared_encoder().encode([query], prompt_name="query")
+    if len(embeddings) != 1 or len(embeddings[0]) != EMBEDDING_DIMENSIONS:
+        raise ValueError("Embedding model returned unexpected dimensions.")
+    scores = shared_reranker().predict([(query, query)], show_progress_bar=False)
+    if len(scores) != 1:
+        raise ValueError("Reranking model returned unexpected scores.")
+
+
 def reciprocal_rank_fusion(*rankings):
     scores = {}
     for ranking in rankings:
@@ -489,15 +501,20 @@ def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("action", choices=("search", "rebuild", "daemon"))
+    parser.add_argument("action", choices=("search", "rebuild", "daemon", "preload"))
     parser.add_argument("--query")
     parser.add_argument("--top-k", type=int, default=FINAL_RESULTS)
     parser.add_argument("--root", default=os.getcwd(), type=Path)
-    parser.add_argument("--data-dir", required=True, type=Path)
+    parser.add_argument("--data-dir", type=Path)
     parser.add_argument("--model-cache", required=True, type=Path)
     args = parser.parse_args()
-    data_dir = args.data_dir.resolve()
     model_cache = args.model_cache.resolve()
+    if args.action == "preload":
+        preload_models(model_cache)
+        return
+    if args.data_dir is None:
+        parser.error(f"{args.action} requires --data-dir")
+    data_dir = args.data_dir.resolve()
     if args.action == "daemon":
         data_dir.mkdir(parents=True, exist_ok=True)
         run_daemon(data_dir, model_cache)