fix(codex): extend retrieval timeout

scripts/lib/semantic-retrieval.py

@@ -10,6 +10,8 @@ import sys
 
 NAME = 'my-ai-qwen3-retrieval'
 FILES = ('server.py', 'requirements.txt')
+CUDA_INDEX = 'https://download.pytorch.org/whl/cu126'
+TOOL_TIMEOUT_SECONDS = 1800
 
 
 def digest(path):
@@ -43,11 +45,38 @@ def ensure_runtime(target, dry_run):
     python = python_path(venv)
     if not python.is_file():
         raise ValueError('Semantic retrieval virtual environment is incomplete.')
-    subprocess.run([str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')],
-                   check=True, timeout=900)
+    arguments = [str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')]
+    if shutil.which('nvidia-smi'):
+        try:
+            subprocess.run([*arguments, '--extra-index-url', CUDA_INDEX], check=True, timeout=900)
+        except subprocess.CalledProcessError:
+            print('Semantic retrieval: CUDA runtime installation failed; falling back to CPU.', file=sys.stderr)
+            subprocess.run(arguments, check=True, timeout=900)
+    else:
+        subprocess.run(arguments, check=True, timeout=900)
     return python
 
 
+def set_codex_tool_timeout(path, seconds):
+    lines = path.read_text(encoding='utf-8').splitlines(keepends=True)
+    section = f'[mcp_servers.{NAME}]'
+    start = next((number for number, line in enumerate(lines) if line.strip() == section), None)
+    if start is None:
+        raise ValueError('Codex semantic retrieval configuration is missing.')
+    end = next((number for number in range(start + 1, len(lines)) if lines[number].lstrip().startswith('[')), len(lines))
+    for number in range(start + 1, end):
+        if lines[number].lstrip().startswith('tool_timeout_sec'):
+            newline = '\r\n' if lines[number].endswith('\r\n') else '\n'
+            lines[number] = f'tool_timeout_sec = {seconds}{newline}'
+            break
+    else:
+        while end > start + 1 and not lines[end - 1].strip():
+            end -= 1
+        newline = '\r\n' if lines and lines[0].endswith('\r\n') else '\n'
+        lines.insert(end, f'tool_timeout_sec = {seconds}{newline}')
+    path.write_text(''.join(lines), encoding='utf-8')
+
+
 def load_state(path):
     if not path.exists():
         return {'version': 1, 'clients': [], 'files': {}}
@@ -121,6 +150,8 @@ def sync(root, home, clients, enabled, dry_run, update, summary=False):
         if update and client in existing_clients:
             command(client, [client, 'mcp', 'remove', NAME] + (['--scope', 'user'] if client == 'claude' else []), home, False)
         command(client, command_args, home, False)
+        if client == 'codex':
+            set_codex_tool_timeout(home / '.codex' / 'config.toml', TOOL_TIMEOUT_SECONDS)
         if client not in state['clients']:
             state['clients'].append(client)
             save_state(state_path, state, False)

scripts/test-retrieval-server.py

@@ -27,6 +27,12 @@ def embedding():
 
 
 class RetrievalServerTests(unittest.TestCase):
+    def test_preferred_device_uses_cuda_with_cpu_fallback(self):
+        with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))):
+            self.assertEqual(MODULE.preferred_device(), 'cuda')
+        with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))):
+            self.assertEqual(MODULE.preferred_device(), 'cpu')
+
     @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):
         os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE']
@@ -125,6 +131,26 @@ class RetrievalServerTests(unittest.TestCase):
             finally:
                 index.db.close()
 
+    def test_refresh_batches_embeddings_for_changed_files(self):
+        with TemporaryDirectory() as temporary:
+            root = Path(temporary)
+            (root / 'first.md').write_text('first', encoding='utf-8')
+            (root / 'second.md').write_text('second', encoding='utf-8')
+            index = MODULE.Index(root, root / 'data')
+            encoded = []
+
+            def encode(values, **kwargs):
+                encoded.append(list(values))
+                return [embedding() for _ in values]
+
+            index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
+            try:
+                index.rebuild()
+                self.assertEqual(len(encoded), 1)
+                self.assertEqual(len(encoded[0]), 2)
+            finally:
+                index.db.close()
+
     def test_scan_prunes_ignored_directories_and_stops_at_limit(self):
         with TemporaryDirectory() as temporary:
             root = Path(temporary)

scripts/test-semantic-retrieval.py

@@ -6,6 +6,7 @@ from contextlib import redirect_stdout
 from pathlib import Path
 import tempfile
 import unittest
+from unittest.mock import patch
 
 
 ROOT = Path(__file__).resolve().parent.parent
@@ -15,6 +16,40 @@ SPEC.loader.exec_module(MODULE)
 
 
 class SemanticRetrievalTests(unittest.TestCase):
+    def test_codex_registration_sets_tool_timeout(self):
+        with tempfile.TemporaryDirectory() as directory:
+            config = Path(directory) / 'config.toml'
+            config.write_text('[mcp_servers.my-ai-qwen3-retrieval]\ncommand = "python"\n\n[mcp_servers.my-ai-qwen3-retrieval.env]\nHF_HOME = "cache"\n', encoding='utf-8')
+
+            MODULE.set_codex_tool_timeout(config, 1800)
+
+            self.assertEqual(config.read_text(encoding='utf-8'),
+                             '[mcp_servers.my-ai-qwen3-retrieval]\ncommand = "python"\ntool_timeout_sec = 1800\n\n[mcp_servers.my-ai-qwen3-retrieval.env]\nHF_HOME = "cache"\n')
+
+    def test_runtime_prefers_cuda_packages_for_nvidia_gpu(self):
+        with tempfile.TemporaryDirectory() as directory:
+            target = Path(directory)
+            python = MODULE.python_path(target / '.venv')
+            python.parent.mkdir(parents=True)
+            python.touch()
+            with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run') as run:
+                MODULE.ensure_runtime(target, False)
+
+            self.assertEqual(run.call_args.args[0][-2:], ['--extra-index-url', MODULE.CUDA_INDEX])
+
+    def test_runtime_falls_back_to_cpu_packages_when_cuda_install_fails(self):
+        with tempfile.TemporaryDirectory() as directory:
+            target = Path(directory)
+            python = MODULE.python_path(target / '.venv')
+            python.parent.mkdir(parents=True)
+            python.touch()
+            failure = MODULE.subprocess.CalledProcessError(1, 'pip')
+            with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run', side_effect=[failure, None]) as run:
+                MODULE.ensure_runtime(target, False)
+
+            self.assertEqual(run.call_count, 2)
+            self.assertNotIn('--extra-index-url', run.call_args.args[0])
+
     def test_summary_dry_run_omits_details_without_writes(self):
         with tempfile.TemporaryDirectory() as directory:
             home = Path(directory) / 'home'

shared/retrieval/server.py

@@ -118,6 +118,11 @@ def document_text(path: str, symbol: str, content: str):
     return f"Path: {path}\nSymbol: {symbol}\n\n{content}"
 
 
+def preferred_device():
+    import torch
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
 def reciprocal_rank_fusion(*rankings):
     scores = {}
     for ranking in rankings:
@@ -165,13 +170,15 @@ class Index:
         if self.model is None:
             from sentence_transformers import SentenceTransformer
             self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
-                                             prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
+                                             prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
+                                             device=preferred_device())
         return self.model
 
     def reranker(self):
         if self.reranker_model is None:
             from sentence_transformers import CrossEncoder
-            self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION}, default_prompt_name="query")
+            self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION},
+                                               default_prompt_name="query", device=preferred_device())
         return self.reranker_model
 
     def rerank(self, query: str, rows: list[dict], limit: int):
@@ -193,19 +200,31 @@ class Index:
 
     def refresh(self, force=False):
         previous = dict(self.db.execute("select path, digest from files"))
+        encoder = None
+        changed = []
+        documents = []
+        document_offsets = []
+        for path in self.files():
+            relative = path.relative_to(self.root).as_posix()
+            content = path.read_bytes()
+            digest = hashlib.sha256(content).hexdigest()
+            old_digest = previous.pop(relative, None)
+            if not force and digest == old_digest:
+                continue
+            if encoder is None:
+                encoder = self.encoder()
+            values = list(chunks(content.decode("utf-8", errors="ignore"), relative, encoder.tokenizer))
+            file_documents = [document_text(relative, symbol, value) for symbol, value in values]
+            document_offsets.append(len(documents))
+            documents.extend(file_documents)
+            changed.append((relative, digest, values, file_documents))
+        vectors = encoder.encode(documents, normalize_embeddings=True) if documents else []
         with self.db:
-            for path in self.files():
-                relative = path.relative_to(self.root).as_posix()
-                content = path.read_bytes()
-                digest = hashlib.sha256(content).hexdigest()
-                old_digest = previous.pop(relative, None)
-                if not force and digest == old_digest:
-                    continue
-                values = list(chunks(content.decode("utf-8", errors="ignore"), relative, self.encoder().tokenizer))
-                documents = [document_text(relative, symbol, value) for symbol, value in values]
-                vectors = self.encoder().encode(documents, normalize_embeddings=True) if documents else []
+            for number, (relative, digest, values, file_documents) in enumerate(changed):
+                start = document_offsets[number]
+                file_vectors = vectors[start:start + len(file_documents)]
                 self.delete_path(relative)
-                for (symbol, value), document, vector in zip(values, documents, vectors):
+                for (symbol, value), document, vector in zip(values, file_documents, file_vectors):
                     if len(vector) != EMBEDDING_DIMENSIONS:
                         raise ValueError("Embedding must have 1024 dimensions")
                     cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",