fix: stop the re-embed migration from destroying the index it rebuilds

Five defects from a review of the embeddings work, four of them silent.

- Write local files atomically. `LocalStorage.save_file` streamed straight onto
  the destination, so an interrupted write left a truncated file. `reembed`
  rewrites every index it touches, and a half-written `index.faiss` loads at
  neither the old width nor the new one -- the source was unrecoverable, with
  no backup and no temp file left behind. Bytes now land beside the destination
  and move into place with `os.replace`. S3 was already safe (single PUT).

- Read pgvector chunks a page at a time. `reembed_pgvector` materialised every
  `(id, text)` row for a source before embedding -- ~1.6 GB at 200k chunks and
  several times that for non-Latin scripts, with the `PGresult` held alongside
  until the cursor closed. Inside the shipped 4Gi limit, while also holding the
  model, that is an OOMKill -- which is exactly the SIGKILL the point above
  turned into a destroyed index. It now walks the source by keyset.

- Bound the first wave of delegated embeds. The failure cooldown is only latched
  once the first `get()` returns, so every request already in flight paid the
  full EMBEDDINGS_DELEGATE_TIMEOUT: measured 64 threads all timing out together,
  and at the shipped 60s across a 96-thread WSGI pool that is an API serving
  nothing at all, health checks included. One caller now probes while the rest
  fail fast; after a single success the gate leaves the path entirely.

- Ship EMBEDDINGS_NAME commented in .env-template. The comment directly above it
  says to leave it commented when upgrading, and the line shipped set. Any value
  reaching `.env` lands in `model_fields_set`, which makes `resolve_embeddings_pin`
  bail -- so a template-derived `.env` disabled the legacy pin outright and
  repointed a populated index at a different 768-dim model, where no width check
  fires. The pin already picks granite for a fresh install and mpnet for an
  existing one, so nothing needs to be set by hand.

- Stamp `sources.model` on wiki sources. They were created with the column NULL
  and then embedded like any other source, and the boot check reads NULL as
  "pre-dates the column, therefore the legacy model" -- reporting a correctly
  embedded source as stale on every startup of every process. Stamped at
  creation, and again on each page re-embed so existing rows heal.

The two docs that promised the FAISS index survives a failed run said so of the
embed only; both now describe the write, and upgrading.mdx says to stop ingest
for the duration.

.env-template

@@ -11,16 +11,18 @@ INTERNAL_KEY=<internal key for worker-to-backend authentication>
 # NOVITA_API_KEY=<your-novita-api-key>
 # OPEN_ROUTER_API_KEY=<your-openrouter-api-key>
 
-# Embedding model. New installs should use granite: multilingual, 32k context,
-# and the same 768 dimensions as the legacy model.
+# Embedding model. Leave it commented out and DocsGPT picks one for you: a
+# fresh install is pinned to granite (multilingual, 32k context, the same 768
+# dimensions as the legacy model), and an install that already has sources
+# keeps the model its index was built with.
 #
-# UPGRADING AN EXISTING INSTALL? Leave this commented out. Vectors already in
-# your index were produced by the old model, and changing this without
-# re-embedding gives you retrieval against a different vector space -- which
-# fails silently, because both models are 768-dimensional. To switch, set this
-# and then run:
+# Setting it here overrides that pin, so only set it deliberately. On an index
+# that already has vectors, changing it without re-embedding leaves queries
+# searching a different vector space than the stored vectors -- which fails
+# silently, because both models are 768-dimensional. To switch, set it and then
+# run:
 #     python -m application.scripts.reembed
-EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2
+# EMBEDDINGS_NAME=ibm-granite/granite-embedding-311m-multilingual-r2
 
 # Remote Embeddings (Optional - for using a remote embeddings API instead of
 # running the model in-process). When set, the app calls the remote API and

application/api/user/sources/routes.py

@@ -800,6 +800,12 @@ class CreateWikiSource(Resource):
                     config={"kind": "wiki"},
                     directory_structure={},
                     tokens=0,
+                    # Wiki pages are embedded like any other source, so record
+                    # which model did it. Left unset, the column reads as NULL,
+                    # which the boot mismatch check takes to mean "pre-dates the
+                    # column, therefore the legacy model" -- and reports a
+                    # correctly-embedded source as stale on every startup.
+                    model=settings.EMBEDDINGS_NAME,
                 )
                 if initial_content:
                     WikiPagesRepository(conn).upsert(

application/scripts/reembed.py

@@ -197,15 +197,89 @@ def reembed_graph_nodes(store, conn, source_id: str, batch_size: int, dry_run: b
     return written
 
 
+def _count_source_chunks(conn, table: str, source_id: str) -> int:
+    """How many chunk rows one source holds."""
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            sql.SQL("SELECT count(*) FROM {table} WHERE source_id = %s").format(
+                table=sql.Identifier(table)
+            ),
+            (source_id,),
+        )
+        return int(cursor.fetchone()[0])
+    finally:
+        cursor.close()
+
+
+def _fetch_chunk_page(
+    conn, table: str, source_id: str, after_id: Optional[Any], limit: int
+) -> List[Tuple[Any, str]]:
+    """One page of ``(id, text)`` ordered by id, starting after ``after_id``.
+
+    Keyed off ``id`` rather than ``OFFSET`` so each page costs an index seek,
+    and the page boundaries stay stable across the updates the caller makes
+    between pages -- only the vector column changes, never ``id``.
+    """
+    cursor = conn.cursor()
+    try:
+        if after_id is None:
+            cursor.execute(
+                sql.SQL(
+                    "SELECT id, text FROM {table} WHERE source_id = %s "
+                    "ORDER BY id LIMIT %s"
+                ).format(table=sql.Identifier(table)),
+                (source_id, limit),
+            )
+        else:
+            cursor.execute(
+                sql.SQL(
+                    "SELECT id, text FROM {table} WHERE source_id = %s AND id > %s "
+                    "ORDER BY id LIMIT %s"
+                ).format(table=sql.Identifier(table)),
+                (source_id, after_id, limit),
+            )
+        return cursor.fetchall()
+    finally:
+        cursor.close()
+
+
+def _write_chunk_page(conn, table: str, vector_column: str, store, page) -> None:
+    """Embed one page and commit its vectors."""
+    vectors = store._embedding.embed_documents([row[1] or "" for row in page])
+    cursor = conn.cursor()
+    try:
+        cursor.executemany(
+            sql.SQL("UPDATE {table} SET {column} = %s::vector WHERE id = %s").format(
+                table=sql.Identifier(table),
+                column=sql.Identifier(vector_column),
+            ),
+            [(str(list(vector)), row[0]) for vector, row in zip(vectors, page)],
+        )
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        cursor.close()
+
+
 def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
     """Rewrite one source's vectors in place.
 
     Rows are updated, never deleted and re-inserted, so an interrupted run
     leaves every chunk present and simply re-does its batch next time.
 
+    Chunks are read a page at a time rather than all at once. A source's text is
+    the whole corpus: at the default chunk size 200k chunks materialise to about
+    1.6 GB of Python strings, several times that for non-Latin scripts, and the
+    result set is held alongside them until the cursor closes. That is enough to
+    be OOM-killed inside a 4 GiB container while also holding the model -- and a
+    kill here is what a torn index costs the most.
+
     Args:
         source_id: Source whose chunks to re-embed.
-        batch_size: Chunks per embed call and per transaction.
+        batch_size: Chunks per page, per embed call and per transaction.
         dry_run: When true, count the work and change nothing.
 
     Returns:
@@ -215,44 +289,22 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in
     table, vector_column = store._table_name, store._vector_column
     conn = store._get_connection()
 
-    cursor = conn.cursor()
+    seen = written = 0
     try:
-        cursor.execute(
-            sql.SQL("SELECT id, text FROM {table} WHERE source_id = %s ORDER BY id").format(
-                table=sql.Identifier(table)
-            ),
-            (source_id,),
-        )
-        rows = cursor.fetchall()
-    finally:
-        cursor.close()
-
-    written = 0
-    try:
-        if not dry_run:
-            for batch in _batched(rows, batch_size):
-                vectors = store._embedding.embed_documents([row[1] or "" for row in batch])
-                cursor = conn.cursor()
-                try:
-                    cursor.executemany(
-                        sql.SQL(
-                            "UPDATE {table} SET {column} = %s::vector WHERE id = %s"
-                        ).format(
-                            table=sql.Identifier(table),
-                            column=sql.Identifier(vector_column),
-                        ),
-                        [
-                            (str(list(vector)), row[0])
-                            for vector, row in zip(vectors, batch)
-                        ],
-                    )
-                    conn.commit()
-                except Exception:
-                    conn.rollback()
-                    raise
-                finally:
-                    cursor.close()
-                written += len(batch)
+        if dry_run:
+            seen = _count_source_chunks(conn, table, source_id)
+        else:
+            after_id: Optional[Any] = None
+            while True:
+                page = _fetch_chunk_page(
+                    conn, table, source_id, after_id, batch_size
+                )
+                if not page:
+                    break
+                _write_chunk_page(conn, table, vector_column, store, page)
+                after_id = page[-1][0]
+                seen += len(page)
+                written += len(page)
 
         # The graph seeds every traversal from its own vectors, so leaving them
         # in the old model's space is the same silent mismatch this script
@@ -268,7 +320,7 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in
                 )
     finally:
         store.close()
-    return len(rows), written
+    return seen, written
 
 
 def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
@@ -276,7 +328,9 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
 
     A FAISS index is a flat array, so this rebuilds rather than updates. Every
     new vector is computed *before* the old index is touched, so an interrupt
-    during embedding leaves the existing index intact.
+    during embedding leaves the existing index intact -- and ``save_local``
+    writes each file to a temporary path and moves it into place, so an
+    interrupt during the write does too.
 
     Args:
         source_id: Source whose index to rebuild.

application/storage/local.py

@@ -1,6 +1,7 @@
 """Local file system implementation."""
 import os
 import shutil
+import tempfile
 from typing import BinaryIO, List, Callable
 
 from application.storage.base import BaseStorage
@@ -36,16 +37,49 @@ class LocalStorage(BaseStorage):
         return resolved
 
     def save_file(self, file_data: BinaryIO, path: str, **kwargs) -> dict:
-        """Save a file to local storage."""
+        """Save a file, replacing any existing one atomically.
+
+        The bytes land on a temporary file beside the destination and are moved
+        into place with ``os.replace``, so an interrupted write leaves the
+        previous file intact instead of a truncated one. Streaming straight onto
+        the destination is unrecoverable for a file that is rewritten in place:
+        a half-written ``index.faiss`` loads at neither the old width nor the
+        new one, and ``application.scripts.reembed`` rewrites every index it
+        touches.
+        """
         full_path = self._get_full_path(path)
+        directory = os.path.dirname(full_path)
+        os.makedirs(directory, exist_ok=True)
 
-        os.makedirs(os.path.dirname(full_path), exist_ok=True)
-
-        if hasattr(file_data, 'save'):
-            file_data.save(full_path)
-        else:
-            with open(full_path, 'wb') as f:
-                shutil.copyfileobj(file_data, f)
+        # Same directory as the destination, so the replace below is a rename
+        # within one filesystem rather than a copy.
+        fd, temp_path = tempfile.mkstemp(
+            dir=directory, prefix=f".{os.path.basename(full_path)}.", suffix=".tmp"
+        )
+        try:
+            with os.fdopen(fd, "wb") as f:
+                if hasattr(file_data, "save"):
+                    file_data.save(f)
+                else:
+                    shutil.copyfileobj(file_data, f)
+                f.flush()
+                os.fsync(f.fileno())
+            # mkstemp is 0600; keep whatever the file already had, or fall back
+            # to what a plain open() would have produced.
+            try:
+                mode = os.stat(full_path).st_mode & 0o777
+            except FileNotFoundError:
+                mode = 0o644
+            os.chmod(temp_path, mode)
+            os.replace(temp_path, full_path)
+        except BaseException:
+            # A successful replace consumes the temp file; every other exit
+            # leaves it behind.
+            try:
+                os.unlink(temp_path)
+            except OSError:
+                pass
+            raise
 
         return {
             'storage_type': 'local'

application/vectorstore/embeddings_delegated.py

@@ -22,6 +22,7 @@ network hop rather than a broker round trip.
 from __future__ import annotations
 
 import logging
+import threading
 import time
 from typing import Any, List, Optional
 
@@ -40,6 +41,13 @@ EMBED_TASK = "application.vectorstore.embeddings_tasks.embed_texts"
 #: single retrieval into one timeout rather than one per source.
 _FAILURE_COOLDOWN = 30.0
 
+#: How long a caller waits for the outcome of the dispatch already in flight
+#: before giving up on its own. Only applies while the worker is unproven --
+#: once one dispatch has succeeded, every caller goes straight to the broker.
+#: Comfortably above a healthy round trip (~60 ms on a prefork worker) and far
+#: below ``EMBEDDINGS_DELEGATE_TIMEOUT``, which is the point.
+_PROBE_WAIT = 2.0
+
 _NO_WORKER_HINT = (
     "Start a worker consuming it, point EMBEDDINGS_BASE_URL at an embedding "
     "service, or set EMBEDDINGS_DELEGATE_TO_WORKER=false to load the model in "
@@ -89,12 +97,21 @@ class DelegatedEmbeddings:
         self._local: Any = None
         self._dimension: Optional[int] = dimension_for(embeddings_name)
         self._failed_at: Optional[float] = None
+        # A dispatch has completed successfully, so the worker is known to be
+        # consuming the queue and callers need not take turns proving it.
+        self._verified = False
+        self._probing = False
+        self._state_lock = threading.Lock()
+        self._probe_done = threading.Event()
 
     def _cooldown_remaining(self) -> float:
         """Seconds left of the fail-fast window after a failed dispatch."""
-        if self._failed_at is None:
+        # One load: a concurrent success clearing the latch between two reads
+        # would otherwise subtract from None.
+        failed_at = self._failed_at
+        if failed_at is None:
             return 0.0
-        return max(0.0, _FAILURE_COOLDOWN - (time.monotonic() - self._failed_at))
+        return max(0.0, _FAILURE_COOLDOWN - (time.monotonic() - failed_at))
 
     def _local_embeddings(self):
         """The in-process model, built once, for use inside a worker task."""
@@ -104,26 +121,10 @@ class DelegatedEmbeddings:
             self._local = build_local_embeddings(self.embeddings_name, self.embeddings_key)
         return self._local
 
-    def _dispatch(self, texts: List[str]) -> List[List[float]]:
-        """Run the embed task on the worker and wait for its vectors."""
+    def _send(self, texts: List[str], queue: str, timeout: int) -> List[List[float]]:
+        """Publish the embed task and wait for its vectors."""
         from application.celery_init import celery
 
-        queue = getattr(settings, "EMBEDDINGS_QUEUE", "embeddings")
-        timeout = getattr(settings, "EMBEDDINGS_DELEGATE_TIMEOUT", 60)
-
-        # A missing worker is a property of the deployment, not of this call,
-        # so once one dispatch has timed out the next is not worth another full
-        # timeout. Without this latch a single retrieval pays the timeout twice
-        # -- once in ``fanout.embed_questions``, then again per source when it
-        # falls back to letting each store embed its own query.
-        remaining = self._cooldown_remaining()
-        if remaining > 0:
-            raise RuntimeError(
-                f"Skipping the embed dispatch: a previous request to the {queue!r} "
-                f"queue failed and the {_FAILURE_COOLDOWN}s cooldown has "
-                f"{remaining:.0f}s left. {_NO_WORKER_HINT}"
-            )
-
         result = celery.send_task(EMBED_TASK, args=[texts, self.embeddings_name], queue=queue)
         try:
             vectors = result.get(timeout=timeout)
@@ -137,8 +138,76 @@ class DelegatedEmbeddings:
         finally:
             _forget(result)
         self._failed_at = None
+        self._verified = True
         return vectors
 
+    def _cooldown_error(self, queue: str, remaining: float) -> RuntimeError:
+        return RuntimeError(
+            f"Skipping the embed dispatch: a previous request to the {queue!r} "
+            f"queue failed and the {_FAILURE_COOLDOWN}s cooldown has "
+            f"{remaining:.0f}s left. {_NO_WORKER_HINT}"
+        )
+
+    def _dispatch(self, texts: List[str]) -> List[List[float]]:
+        """Run the embed task on the worker and wait for its vectors.
+
+        A missing worker is a property of the deployment, not of this call, so
+        at most one caller waits out ``EMBEDDINGS_DELEGATE_TIMEOUT`` to discover
+        it. Two guards do that:
+
+        The cooldown latch covers requests arriving *after* a failure -- without
+        it a single retrieval pays the timeout twice, once in
+        ``fanout.embed_questions`` and again per source when it falls back to
+        letting each store embed its own query.
+
+        The probe covers requests already in flight *alongside* the first one,
+        which the latch cannot: nothing is latched until that first ``get()``
+        returns, so every thread in the opening wave would otherwise block for
+        the full timeout at once -- at the shipped 60s and 96 WSGI threads, an
+        API that serves nothing at all, health checks included.
+        """
+        queue = getattr(settings, "EMBEDDINGS_QUEUE", "embeddings")
+        timeout = getattr(settings, "EMBEDDINGS_DELEGATE_TIMEOUT", 60)
+
+        remaining = self._cooldown_remaining()
+        if remaining > 0:
+            raise self._cooldown_error(queue, remaining)
+
+        if self._verified:
+            return self._send(texts, queue, timeout)
+
+        with self._state_lock:
+            # "send" -- proven while we waited for the lock, just go.
+            # "wait" -- another caller is already finding out; don't pay a
+            #           second full timeout to learn the same thing.
+            # "probe" -- nobody is; this call is the one that finds out.
+            role = "send" if self._verified else "wait" if self._probing else "probe"
+            if role == "probe":
+                self._probing = True
+                self._probe_done.clear()
+
+        if role == "probe":
+            try:
+                return self._send(texts, queue, timeout)
+            finally:
+                with self._state_lock:
+                    self._probing = False
+                self._probe_done.set()
+
+        if role == "wait":
+            self._probe_done.wait(_PROBE_WAIT)
+            remaining = self._cooldown_remaining()
+            if remaining > 0:
+                raise self._cooldown_error(queue, remaining)
+            if not self._verified:
+                raise RuntimeError(
+                    f"Skipping the embed dispatch: an earlier request to the "
+                    f"{queue!r} queue is still unanswered after {_PROBE_WAIT}s, so "
+                    f"no worker appears to be consuming it. {_NO_WORKER_HINT}"
+                )
+
+        return self._send(texts, queue, timeout)
+
     def embed_documents(self, documents: List[str]) -> List[List[float]]:
         """Embed a list of texts, preserving order."""
         if not documents:

application/vectorstore/faiss.py

@@ -325,7 +325,13 @@ class FaissStore(BaseVectorStore):
             f.write(dump_pickle_sidecar(self.documents, self.index_to_docstore_id))
 
     def _save_to_storage(self) -> bool:
-        """Persist the index through the configured storage backend."""
+        """Persist the index through the configured storage backend.
+
+        Each file is replaced atomically by the backend, so an interrupted save
+        leaves the previous one readable. The three are still written in
+        sequence: a crash between them pairs a new index with an older sidecar,
+        which stays consistent because the ids and row order are preserved.
+        """
         with tempfile.TemporaryDirectory() as temp_dir:
             self._write_index_files(temp_dir)
             storage_path = get_vectorstore(self.source_id)

application/worker.py

@@ -2523,6 +2523,14 @@ def reembed_wiki_page_worker(self, source_id, path, content_hash, user):
 
         with db_session() as conn:
             WikiPagesRepository(conn).set_embed_status(source_id, path, "embedded")
+            # These chunks were just embedded with the configured model, so the
+            # source now names it. Wiki sources created before this was recorded
+            # carry NULL, which the boot mismatch check reads as the legacy
+            # model and reports as stale; stamping here heals them on the next
+            # page edit.
+            SourcesRepository(conn).update(
+                source_id, user, {"model": settings.EMBEDDINGS_NAME}
+            )
     except Exception:
         with db_session() as conn:
             WikiPagesRepository(conn).set_embed_status(source_id, path, "failed")

docs/content/Deploying/DocsGPT-Settings.mdx

@@ -59,7 +59,7 @@ Here are some of the most fundamental settings you'll likely want to configure:
 
 - **`EMBEDDINGS_NAME`**: This setting defines which embedding model DocsGPT will use to generate vector embeddings for your documents. Embeddings are numerical representations of text that allow DocsGPT to understand the semantic meaning of your documents for efficient search and retrieval.
 
-  - **Default value:** `huggingface_sentence-transformers/all-mpnet-base-v2`, kept so an existing index stays readable. New installs are pointed at `ibm-granite/granite-embedding-311m-multilingual-r2` (multilingual, 32k context, same 768 dimensions) by `.env-template` and `setup.sh`.
+  - **Default value:** leave it unset and DocsGPT picks for you at first boot, recording the choice so it never changes underneath you: a fresh install is pinned to `ibm-granite/granite-embedding-311m-multilingual-r2` (multilingual, 32k context, same 768 dimensions), and an install that already has sources is pinned to `huggingface_sentence-transformers/all-mpnet-base-v2` so its index stays readable. Setting it here overrides that pin.
   - **Other options:** Any FastEmbed built-in model, or any Hugging Face repository shipping an ONNX export. See [Embeddings](/Models/embeddings).
   - **Changing it on an existing index requires re-embedding** — same-width models swap without any error and silently degrade retrieval. Run `python -m application.scripts.reembed`.
 

docs/content/upgrading.mdx

@@ -55,7 +55,11 @@ Useful flags:
 | `--sources a,b` | Only these source ids — also how you retry a failed source |
 | `--batch-size N` | Chunks per embed call (default 64) |
 
-The script processes sources independently: one failing source is reported and skipped rather than aborting the run, and the exit code is non-zero if any failed. For `pgvector` it updates rows in place, so an interrupted run simply re-does its last batch. For `faiss` it builds the replacement index in memory and only writes once the whole rebuild succeeds, leaving the existing index untouched on failure.
+The script processes sources independently: one failing source is reported and skipped rather than aborting the run, and the exit code is non-zero if any failed. For `pgvector` it reads a page of chunks at a time and updates rows in place, so memory stays flat on a large index and an interrupted run simply re-does its last batch. For `faiss` it builds the replacement index in memory, writes each file to a temporary path, and moves it into place — so an interrupt during either the rebuild or the write leaves the existing index intact rather than truncated.
+
+<Callout type="warning">
+  Stop ingest before you run this. It reads each source's chunks and writes the vectors back; anything ingested while it runs can be overwritten by the rebuild (`faiss`) or missed by it (`pgvector`).
+</Callout>
 
 <Callout type="info">
   Running [GraphRAG](/Sources/GraphRAG)? The script also rewrites `graph_nodes.name_embedding`, which seeds every graph traversal. Those vectors are written once at extraction time and share the chunk vectors' width, so leaving them in the old model's space degrades graph retrieval just as silently as the chunk vectors would — and needs no LLM re-extraction to fix.

tests/api/user/sources/test_wiki_routes.py

@@ -95,6 +95,12 @@ class TestCreateWikiSource:
         assert row["type"] == "wiki"
         assert row["config"]["kind"] == "wiki"
         assert int(row["tokens"]) == 0
+        # Wiki pages get embedded like any other source, so the row has to name
+        # the model that did it. NULL reads as "the legacy model" to the boot
+        # mismatch check, which then reports the source as stale forever.
+        from application.core.settings import settings
+
+        assert row["model"] == settings.EMBEDDINGS_NAME
         # No seed content → no re-embed, and never any ingest/reingest task.
         mock_reembed.assert_not_called()
         mock_ingest.assert_not_called()

tests/scripts/test_reembed.py

@@ -7,6 +7,46 @@ import pytest
 from application.scripts import reembed
 
 
+def paginating_cursor(chunk_rows, *, graph_rows=(), graph_table=("graph_nodes",)):
+    """A cursor answering the reads ``reembed_pgvector`` issues, from memory.
+
+    The chunk read is paginated by keyset, so a cursor that returns the same
+    page for every ``fetchall`` never terminates. Returns the cursor and the
+    mutable table backing it.
+
+    Args:
+        chunk_rows: ``(id, text)`` rows for the source's chunk table.
+        graph_rows: ``(id, name)`` rows for ``graph_nodes``.
+        graph_table: What ``to_regclass`` reports; ``(None,)`` for absent.
+
+    Returns:
+        ``(cursor, tables)``, where ``tables["chunks"]`` and
+        ``tables["graph"]`` can be reassigned to change what is served.
+    """
+    tables = {"chunks": list(chunk_rows), "graph": list(graph_rows)}
+    pending = {"rows": []}
+    cursor = MagicMock()
+
+    def execute(query, params=None):
+        text = str(query)
+        if "to_regclass" in text:
+            pending["rows"] = [graph_table]
+        elif "FROM graph_nodes" in text:
+            pending["rows"] = list(tables["graph"])
+        elif "count(*)" in text:
+            pending["rows"] = [(len(tables["chunks"]),)]
+        elif "id > %s" in text:
+            _, after_id, limit = params
+            pending["rows"] = [r for r in tables["chunks"] if r[0] > after_id][:limit]
+        else:
+            pending["rows"] = list(tables["chunks"])[: params[1]]
+
+    cursor.execute.side_effect = execute
+    cursor.fetchall.side_effect = lambda: pending["rows"]
+    cursor.fetchone.side_effect = lambda: pending["rows"][0] if pending["rows"] else None
+    return cursor, tables
+
+
 class TestCLI:
     def test_defaults(self):
         args = reembed.build_parser().parse_args([])
@@ -172,11 +212,18 @@ class TestPgvectorWithoutTheExtension:
 
     @pytest.fixture
     def store(self):
+        """A store whose cursor serves the paginated reads, not a fixed page.
+
+        ``reembed_pgvector`` walks the source by keyset, so a cursor that
+        returns the same rows for every ``fetchall`` never terminates. The fake
+        answers each of the three queries the function issues from one in-memory
+        table, which tests mutate through ``rows``.
+        """
+        cursor, table = paginating_cursor([(1, "alpha"), (2, "beta"), (3, "gamma")])
+
         store = MagicMock()
         store._table_name = "documents"
         store._vector_column = "embedding"
-        cursor = MagicMock()
-        cursor.fetchall.return_value = [(1, "alpha"), (2, "beta"), (3, "gamma")]
         conn = MagicMock()
         conn.cursor.return_value = cursor
         store._get_connection.return_value = conn
@@ -186,60 +233,85 @@ class TestPgvectorWithoutTheExtension:
         with patch.object(
             reembed.VectorCreator, "create_vectorstore", return_value=store
         ), patch.object(reembed.settings, "GRAPHRAG_ENABLED", False):
-            yield store, conn, cursor
+            yield store, conn, cursor, table
 
     def test_reads_and_rewrites_every_chunk(self, store):
-        _, conn, cursor = store
+        _, conn, cursor, _ = store
         seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
         assert (seen, written) == (3, 3)
         cursor.executemany.assert_called_once()
         conn.commit.assert_called()
 
-    def test_dry_run_neither_embeds_nor_writes(self, store):
-        fake_store, conn, cursor = store
+    def test_pages_are_bounded_by_batch_size(self, store):
+        """The read is paginated, so a huge source never lands in memory at once."""
+        _, _, cursor, table = store
+        table["chunks"] = [(i, f"chunk-{i}") for i in range(1, 11)]
+        seen, written = reembed.reembed_pgvector("s1", batch_size=3, dry_run=False)
+        assert (seen, written) == (10, 10)
+        selects = [
+            call for call in cursor.execute.call_args_list
+            if "SELECT id, text" in str(call.args[0])
+        ]
+        # Four pages of at most 3, then one empty page to end the walk.
+        assert len(selects) == 5
+        assert all(call.args[1][-1] == 3 for call in selects)
+
+    def test_dry_run_counts_without_reading_the_text(self, store):
+        fake_store, _, cursor, _ = store
         seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=True)
         assert (seen, written) == (3, 0)
         fake_store._embedding.embed_documents.assert_not_called()
         cursor.executemany.assert_not_called()
+        assert any(
+            "count(*)" in str(call.args[0]) for call in cursor.execute.call_args_list
+        )
 
     def test_batches_commit_separately(self, store):
-        _, conn, cursor = store
+        _, conn, cursor, _ = store
         reembed.reembed_pgvector("s1", batch_size=2, dry_run=False)
         # 3 rows at batch 2 is two write transactions.
         assert cursor.executemany.call_count == 2
         assert conn.commit.call_count == 2
 
     def test_failed_batch_rolls_back_and_raises(self, store):
-        fake_store, conn, cursor = store
+        _, conn, cursor, _ = store
         cursor.executemany.side_effect = RuntimeError("write failed")
         with pytest.raises(RuntimeError):
             reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
         conn.rollback.assert_called_once()
 
     def test_connection_is_returned_even_on_failure(self, store):
-        fake_store, _, cursor = store
+        fake_store, _, cursor, _ = store
         cursor.executemany.side_effect = RuntimeError("boom")
         with pytest.raises(RuntimeError):
             reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
         fake_store.close.assert_called_once()
 
     def test_null_text_does_not_crash_the_embed_call(self, store):
-        fake_store, _, cursor = store
-        cursor.fetchall.return_value = [(1, None), (2, "beta")]
+        fake_store, _, _, table = store
+        table["chunks"] = [(1, None), (2, "beta")]
         seen, written = reembed.reembed_pgvector("s1", batch_size=64, dry_run=False)
         assert (seen, written) == (2, 2)
         assert fake_store._embedding.embed_documents.call_args.args[0] == ["", "beta"]
 
     def test_empty_source_is_a_no_op(self, store):
-        fake_store, _, cursor = store
-        cursor.fetchall.return_value = []
+        fake_store, _, _, table = store
+        table["chunks"] = []
         assert reembed.reembed_pgvector("s1", batch_size=64, dry_run=False) == (0, 0)
         fake_store._embedding.embed_documents.assert_not_called()
 
-    def test_source_discovery_returns_sorted_ids(self, store):
-        _, _, cursor = store
+    def test_source_discovery_returns_sorted_ids(self):
+        store = MagicMock()
+        store._table_name = "documents"
+        cursor = MagicMock()
         cursor.fetchall.return_value = [("b",), ("a",)]
-        assert reembed.list_source_ids("pgvector") == ["b", "a"]
+        conn = MagicMock()
+        conn.cursor.return_value = cursor
+        store._get_connection.return_value = conn
+        with patch.object(
+            reembed.VectorCreator, "create_vectorstore", return_value=store
+        ):
+            assert reembed.list_source_ids("pgvector") == ["b", "a"]
 
 
 class TestFaissSourceDiscovery:
@@ -353,12 +425,13 @@ class TestGraphNodeReembedding:
 
     @staticmethod
     def _pgvector_mocks():
+        cursor, _ = paginating_cursor(
+            [(1, "alpha"), (2, "beta")],
+            graph_rows=[("n1", "Alpha"), ("n2", "Beta")],
+        )
         store = MagicMock()
         store._table_name = "documents"
         store._vector_column = "embedding"
-        cursor = MagicMock()
-        cursor.fetchone.return_value = ("graph_nodes",)
-        cursor.fetchall.return_value = [(1, "alpha"), (2, "beta")]
         conn = MagicMock()
         conn.cursor.return_value = cursor
         store._get_connection.return_value = conn

tests/storage/test_local_storage.py

@@ -16,6 +16,13 @@ def local_storage(temp_base_dir):
     return LocalStorage(base_dir=temp_base_dir)
 
 
+@pytest.fixture
+def real_storage(tmp_path):
+    """Storage over a real directory, for the write paths worth exercising."""
+    base = os.path.realpath(str(tmp_path))
+    return LocalStorage(base_dir=base), base
+
+
 @pytest.mark.unit
 class TestLocalStorageInitialization:
 
@@ -37,50 +44,87 @@ class TestLocalStorageInitialization:
         with pytest.raises(ValueError, match="Path traversal detected"):
             local_storage._get_full_path("/absolute/path/test.txt")
 
-    @patch("os.makedirs")
-    @patch("builtins.open", new_callable=mock_open)
-    @patch("shutil.copyfileobj")
-    def test_save_file_creates_directory_and_saves(
-        self, mock_copyfileobj, mock_file, mock_makedirs, local_storage
-    ):
-        file_data = io.BytesIO(b"test content")
-        path = "documents/test.txt"
+    def test_save_file_creates_directory_and_saves(self, real_storage):
+        storage, base = real_storage
+        result = storage.save_file(io.BytesIO(b"test content"), "documents/test.txt")
 
-        result = local_storage.save_file(file_data, path)
-
-        expected_dir = os.path.join(os.path.realpath("/tmp/test_storage"), "documents")
-        expected_file = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
-
-        assert mock_makedirs.call_count == 1
-        assert os.path.normpath(mock_makedirs.call_args[0][0]) == os.path.normpath(
-            expected_dir
-        )
-        assert mock_makedirs.call_args[1] == {"exist_ok": True}
-
-        assert mock_file.call_count == 1
-        assert os.path.normpath(mock_file.call_args[0][0]) == os.path.normpath(
-            expected_file
-        )
-        assert mock_file.call_args[0][1] == "wb"
-
-        mock_copyfileobj.assert_called_once_with(file_data, mock_file())
+        written = os.path.join(base, "documents/test.txt")
+        assert os.path.isfile(written)
+        with open(written, "rb") as f:
+            assert f.read() == b"test content"
         assert result == {"storage_type": "local"}
 
-    @patch("os.makedirs")
-    def test_save_file_with_save_method(self, mock_makedirs, local_storage):
-        file_data = MagicMock()
-        file_data.save = MagicMock()
-        path = "documents/test.txt"
+    def test_save_file_with_save_method(self, real_storage):
+        """Werkzeug's ``FileStorage.save`` accepts the destination handle."""
+        storage, base = real_storage
 
-        result = local_storage.save_file(file_data, path)
+        class _Uploaded:
+            def save(self, dst):
+                dst.write(b"from save()")
 
-        expected_file = os.path.join(os.path.realpath("/tmp/test_storage"), "documents/test.txt")
-        assert file_data.save.call_count == 1
-        assert os.path.normpath(file_data.save.call_args[0][0]) == os.path.normpath(
-            expected_file
-        )
+        result = storage.save_file(_Uploaded(), "documents/test.txt")
+
+        with open(os.path.join(base, "documents/test.txt"), "rb") as f:
+            assert f.read() == b"from save()"
         assert result == {"storage_type": "local"}
 
+    def test_save_file_replaces_existing_content(self, real_storage):
+        storage, base = real_storage
+        storage.save_file(io.BytesIO(b"first"), "a/b.bin")
+        storage.save_file(io.BytesIO(b"second"), "a/b.bin")
+
+        with open(os.path.join(base, "a/b.bin"), "rb") as f:
+            assert f.read() == b"second"
+
+    def test_failed_write_leaves_the_previous_file_intact(self, real_storage):
+        """The whole point of writing through a temp file.
+
+        ``reembed`` rewrites an index in place; a half-written ``index.faiss``
+        loads at neither the old width nor the new one, so an interrupted write
+        must leave the previous bytes untouched rather than truncate them.
+        """
+        storage, base = real_storage
+        storage.save_file(io.BytesIO(b"the good index"), "indexes/s1/index.faiss")
+
+        class _DiesHalfway:
+            def __init__(self):
+                self._served = False
+
+            def read(self, size=-1):
+                if self._served:
+                    raise OSError("connection reset")
+                self._served = True
+                return b"corrupt"
+
+        with pytest.raises(OSError, match="connection reset"):
+            storage.save_file(_DiesHalfway(), "indexes/s1/index.faiss")
+
+        with open(os.path.join(base, "indexes/s1/index.faiss"), "rb") as f:
+            assert f.read() == b"the good index"
+
+    def test_failed_write_leaves_no_temp_file_behind(self, real_storage):
+        storage, base = real_storage
+
+        class _Explodes:
+            def read(self, size=-1):
+                raise OSError("boom")
+
+        with pytest.raises(OSError):
+            storage.save_file(_Explodes(), "indexes/s1/index.faiss")
+
+        assert os.listdir(os.path.join(base, "indexes/s1")) == []
+
+    def test_save_file_keeps_the_existing_permissions(self, real_storage):
+        """``mkstemp`` is 0600; the replacement must not silently narrow access."""
+        storage, base = real_storage
+        storage.save_file(io.BytesIO(b"one"), "a/b.bin")
+        target = os.path.join(base, "a/b.bin")
+        os.chmod(target, 0o640)
+
+        storage.save_file(io.BytesIO(b"two"), "a/b.bin")
+
+        assert os.stat(target).st_mode & 0o777 == 0o640
+
     def test_save_file_with_absolute_path_outside_base_raises(self, local_storage):
         file_data = io.BytesIO(b"test content")
         path = "/absolute/path/test.txt"

tests/vectorstore/test_embeddings_delegated.py

@@ -1,5 +1,7 @@
 """Query embedding runs on the worker so the API holds no model."""
 
+import threading
+import time
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -187,6 +189,106 @@ class TestFailureCooldown:
         assert celery.send_task.call_count == 3
 
 
+class TestTheConcurrentFirstWave:
+    """The latch cannot cover requests already in flight beside the first one.
+
+    Nothing is latched until that first ``get()`` returns, so every thread in
+    the opening wave would otherwise block for the full
+    ``EMBEDDINGS_DELEGATE_TIMEOUT`` at once -- at the shipped 60s across a 96
+    thread WSGI pool, an API that serves nothing at all.
+    """
+
+    @staticmethod
+    def _blocking_celery(release, outcome):
+        """A worker whose ``get`` blocks until ``release`` is set."""
+        def get(timeout=None):
+            release.wait(5)
+            if isinstance(outcome, Exception):
+                raise outcome
+            return outcome
+
+        result = MagicMock()
+        result.get.side_effect = get
+        celery = MagicMock()
+        celery.send_task.return_value = result
+        return celery
+
+    def _race(self, celery, embeddings, release, threads=8):
+        """Start ``threads`` embeds, let them pile up, then unblock the prober."""
+        errors, values = [], []
+        started = threading.Barrier(threads + 1)
+
+        def call():
+            started.wait(5)
+            try:
+                values.append(embeddings.embed_query("q"))
+            except Exception as exc:  # noqa: BLE001 -- recorded for the assertions
+                errors.append(exc)
+
+        with patch.dict("sys.modules", {"application.celery_init": MagicMock(celery=celery)}):
+            workers = [threading.Thread(target=call) for _ in range(threads)]
+            for worker in workers:
+                worker.start()
+            started.wait(5)
+            time.sleep(0.1)  # let the followers reach the probe gate
+            release.set()
+            for worker in workers:
+                worker.join(10)
+        return values, errors
+
+    def test_only_one_caller_waits_on_an_unproven_worker(self, not_in_worker):
+        release = threading.Event()
+        celery = self._blocking_celery(release, TimeoutError("no worker"))
+        embeddings = DelegatedEmbeddings("granite-311m")
+        with patch(
+            "application.vectorstore.embeddings_delegated._PROBE_WAIT", 0.05
+        ):
+            values, errors = self._race(celery, embeddings, release)
+
+        assert values == []
+        assert len(errors) == 8
+        # One probe published; the rest gave up without their own round trip.
+        assert celery.send_task.call_count == 1
+        assert sum("still unanswered" in str(e) for e in errors) == 7
+
+    def test_the_fast_failure_still_names_the_remedy(self, not_in_worker):
+        release = threading.Event()
+        celery = self._blocking_celery(release, TimeoutError("no worker"))
+        embeddings = DelegatedEmbeddings("granite-311m")
+        with patch("application.vectorstore.embeddings_delegated._PROBE_WAIT", 0.05):
+            _, errors = self._race(celery, embeddings, release, threads=3)
+        assert all("EMBEDDINGS_DELEGATE_TO_WORKER=false" in str(e) for e in errors)
+
+    def test_a_healthy_worker_serves_the_whole_wave(self, not_in_worker):
+        release = threading.Event()
+        celery = self._blocking_celery(release, [[0.1, 0.2]])
+        embeddings = DelegatedEmbeddings("granite-311m")
+        values, errors = self._race(celery, embeddings, release)
+
+        assert errors == []
+        assert values == [[0.1, 0.2]] * 8
+        # The probe proves the worker, then every follower dispatches for real.
+        assert celery.send_task.call_count == 8
+        assert embeddings._verified is True
+
+    def test_a_proven_worker_adds_no_gate(self, not_in_worker):
+        """After one success the probe is out of the path entirely."""
+        release = threading.Event()
+        release.set()
+        celery = self._blocking_celery(release, [[0.3]])
+        embeddings = DelegatedEmbeddings("granite-311m")
+        with patch.dict("sys.modules", {"application.celery_init": MagicMock(celery=celery)}):
+            embeddings.embed_query("warm")
+        assert embeddings._verified is True
+
+        with patch.object(embeddings, "_state_lock") as lock:
+            with patch.dict(
+                "sys.modules", {"application.celery_init": MagicMock(celery=celery)}
+            ):
+                embeddings.embed_query("q")
+            lock.__enter__.assert_not_called()
+
+
 class TestTheResultIsForgotten:
     """A query vector must not outlive the query that asked for it.
 

tests/worker/test_reembed_wiki_page.py

@@ -100,6 +100,46 @@ class TestReembedWikiPageWorker:
         )
         assert result == {"status": "embedded", "added": 2, "deleted": 3}
 
+    def test_a_reembed_stamps_the_model_on_the_source(
+        self, pg_conn, patch_worker_db, task_self, monkeypatch
+    ):
+        """Wiki sources are created with ``model`` NULL, which the boot mismatch
+        check reads as the legacy model and reports as stale on every startup.
+        Stamping here heals a source created before it was recorded.
+        """
+        from application import worker
+        from application.core.settings import settings
+
+        source_id = _seed_source(pg_conn)
+        assert SourcesRepository(pg_conn).get_any(source_id, "alice")["model"] is None
+
+        _patch_store(monkeypatch, MagicMock(name="vector_store"))
+        _patch_repo(monkeypatch, {"content": "body", "title": "T"})
+        _patch_chunker(monkeypatch, [Document(text="c1")])
+
+        worker.reembed_wiki_page_worker(
+            task_self, source_id, "p.md", "hash-1", "alice"
+        )
+
+        row = SourcesRepository(pg_conn).get_any(source_id, "alice")
+        assert row["model"] == settings.EMBEDDINGS_NAME
+
+    def test_a_purge_leaves_the_model_alone(
+        self, pg_conn, patch_worker_db, task_self, monkeypatch
+    ):
+        """Deleting a page embeds nothing, so it claims nothing about the model."""
+        from application import worker
+
+        source_id = _seed_source(pg_conn)
+        _patch_store(monkeypatch, MagicMock(name="vector_store"))
+        _patch_repo(monkeypatch, None)
+
+        worker.reembed_wiki_page_worker(
+            task_self, source_id, "gone.md", "hash-x", "alice"
+        )
+
+        assert SourcesRepository(pg_conn).get_any(source_id, "alice")["model"] is None
+
     def test_page_missing_purges(
         self, pg_conn, patch_worker_db, task_self, monkeypatch
     ):