feat: parser improvements and fixes

application/agents/base.py

@@ -1139,6 +1139,53 @@ class BaseAgent(ABC):
 
     # ---- LLM generation ----
 
+    def _llm_supports_tools(self) -> bool:
+        """Whether the LLM accepts a ``tools`` payload.
+
+        ``_supports_tools`` is a method on every real provider (the old
+        truthiness test on the *bound method* never gated anything), but
+        test doubles sometimes set it as a plain bool, so both shapes are
+        honored. A provider that never implemented the check — ``BaseLLM``
+        raises — keeps the historical permissive behavior; the provider
+        drops the tools itself.
+
+        Returns:
+            True when the ``tools`` kwarg should be sent.
+        """
+        supports = getattr(self.llm, "_supports_tools", None)
+        if supports is None:
+            return False
+        if not callable(supports):
+            return bool(supports)
+        try:
+            return bool(supports())
+        except Exception:
+            logger.debug(
+                "Tool-support check failed for %s; sending tools anyway",
+                type(self.llm).__name__,
+                exc_info=True,
+            )
+            return True
+
+    def _structured_output_kwarg(self) -> Optional[str]:
+        """Name of the gen kwarg this LLM takes structured output on.
+
+        Read off the LLM *class*, not ``llm_name``: every OpenAI-wire
+        provider (``openai_compatible``, ``groq``, ``novita``,
+        ``openrouter``, ``docsgpt``) dispatches through an ``OpenAILLM``
+        subclass, so a string comparison against "openai" silently
+        dropped guided decoding for all of them. ``BaseLLM`` owns the
+        declaration so the cross-provider fallback adapter can read it too.
+
+        Returns:
+            ``"response_format"``, ``"response_schema"``, or None when the
+            provider has no structured-output kwarg.
+        """
+        # ``type(self.llm)`` — an instance attribute on a test double would
+        # otherwise leak a truthy Mock into the gen kwargs.
+        kwarg = getattr(type(self.llm), "structured_output_kwarg", None)
+        return kwarg if kwarg in ("response_format", "response_schema") else None
+
     def _llm_gen(
         self,
         messages: List[Dict],
@@ -1162,11 +1209,7 @@ class BaseAgent(ABC):
         if self.attachments:
             gen_kwargs["_usage_attachments"] = self.attachments
 
-        if (
-            hasattr(self.llm, "_supports_tools")
-            and self.llm._supports_tools
-            and self.tools
-        ):
+        if self.tools and self._llm_supports_tools():
             gen_kwargs["tools"] = self.tools
         if (
             self.json_schema
@@ -1176,14 +1219,12 @@ class BaseAgent(ABC):
             structured_format = self.llm.prepare_structured_output_format(
                 self.json_schema, strict=getattr(self, "json_schema_strict", True)
             )
-            if structured_format:
-                if self.llm_name == "openai":
-                    gen_kwargs["response_format"] = structured_format
-                elif self.llm_name == "google":
-                    gen_kwargs["response_schema"] = structured_format
+            structured_kwarg = self._structured_output_kwarg()
+            if structured_format and structured_kwarg:
+                gen_kwargs[structured_kwarg] = structured_format
         elif (
             getattr(self, "json_object", False)
-            and self.llm_name == "openai"
+            and self._structured_output_kwarg() == "response_format"
             and hasattr(self.llm, "_supports_structured_output")
             and self.llm._supports_structured_output()
         ):

application/agents/tools/read_document.py

@@ -2,7 +2,8 @@
 
 The ``read_document`` action resolves a parent-scoped input artifact, enqueues a
 ``parse_document`` task on the dedicated ``parsing`` Celery queue, and awaits the
-result with a timeout. The run-scoped authz gate is enforced TWICE — here before
+result with a timeout (inside a worker it parses inline under the same bound — see
+``_dispatch``). The run-scoped authz gate is enforced TWICE — here before
 enqueue (reject cross-tenant) and again in the worker (re-resolve, never trusting a
 raw path). When a ``json_schema`` is supplied the structured payload is validated
 through the existing jsonschema path; the full result may also be persisted as a
@@ -12,7 +13,11 @@ through the existing jsonschema path; the full result may also be persisted as a
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List, Optional
+import signal
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from typing import Any, Callable, Dict, List, Optional
 
 from celery import current_task
 
@@ -39,6 +44,15 @@ except Exception:  # pragma: no cover - jsonschema is a declared dependency
     jsonschema = None  # type: ignore[assignment]
 
 
+class _InlineParseTimeout(BaseException):
+    """Raised when an inline (in-worker) parse outlives its window.
+
+    Derives from ``BaseException`` on purpose: the SIGALRM interrupt lands deep inside the
+    parser, whose blanket ``except Exception`` would otherwise swallow it and report a
+    generic parse failure instead of the shared timed-out result.
+    """
+
+
 class ReadDocumentTool(Tool):
     """Read Document
     Parse a document (PDF, Word, PowerPoint, ...) to text, markdown, or structured data.
@@ -58,6 +72,8 @@ class ReadDocumentTool(Tool):
         self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
         self.message_id: Optional[str] = self.config.get("message_id")
         self._last_artifact_id: Optional[str] = None
+        # Byte size of the resolved input, used to scale the awaited parse window.
+        self._input_size: Optional[int] = None
 
     # ------------------------------------------------------------------
     # Tool ABC
@@ -157,6 +173,7 @@ class ReadDocumentTool(Tool):
         """Resolve the input run-scoped (reject cross-tenant before enqueue), enqueue+await, validate."""
         input_id = kwargs.get("input")
         json_schema = kwargs.get("json_schema")
+        self._input_size = None
         if not isinstance(input_id, str) or not input_id.strip():
             return {"status": "error", "error": "input artifact id is required."}
         if json_schema is not None:
@@ -199,8 +216,17 @@ class ReadDocumentTool(Tool):
         and blocking on ``get()`` would self-deadlock: concurrent agent tasks each hold a pool
         slot blocked in ``get()`` so ``parse_document`` never gets a free slot. So parse INLINE
         in-process inside a worker; only dispatch+await (degrading on timeout/failure) from web.
+
+        Both branches are bounded by the SAME size-scaled window: the inline parse carries no
+        Celery time limit of its own, so an unbounded one would pin the agent's worker slot
+        until the OUTER task's limit (webhook runs have none) kills the whole agent run.
         """
         parent = self._parent()
+        from application.api.user.tasks import parse_timeout_for_size
+
+        # OCR cost scales with pages, so the parse window grows with the document's size
+        # (floored at DOCUMENT_PARSE_TIMEOUT).
+        timeout = parse_timeout_for_size(self._input_size)
 
         # ``current_task`` is a Celery proxy: truthy only while this runs inside a worker task,
         # falsy in the web process (the bare proxy is NOT identity-None, so test truthiness).
@@ -208,7 +234,12 @@ class ReadDocumentTool(Tool):
             from application.worker import run_parse_document
 
             try:
-                result = run_parse_document(artifact_id, parent, self.user_id, options)
+                result = self._run_inline_bounded(
+                    lambda: run_parse_document(artifact_id, parent, self.user_id, options), timeout
+                )
+            except _InlineParseTimeout:
+                logger.warning("read_document: inline parse timed out after %ss; abandoning it", int(timeout))
+                return {"status": "error", "error": f"document parsing timed out after {int(timeout)}s."}
             except Exception as exc:
                 logger.exception("read_document: inline parse failed")
                 return {"status": "error", "error": f"document parsing failed: {type(exc).__name__}: {exc}"}
@@ -218,12 +249,18 @@ class ReadDocumentTool(Tool):
 
         from celery.exceptions import TimeoutError as CeleryTimeoutError
 
-        from application.api.user.tasks import parse_document
+        from application.api.user.tasks import parse_document, parse_task_time_limits
 
-        timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
+        # The task's per-call time limits are raised to match the awaited window: bound to
+        # the base timeout at import, the worker would otherwise self-terminate a large
+        # parse long before this await gives up.
         queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
         try:
-            async_result = parse_document.apply_async(args=[artifact_id, parent, self.user_id, options], queue=queue)
+            async_result = parse_document.apply_async(
+                args=[artifact_id, parent, self.user_id, options],
+                queue=queue,
+                **parse_task_time_limits(timeout),
+            )
             # The web process (not a worker) awaits here; ``disable_sync_subtasks=False`` keeps
             # the call correct if invoked from a non-prefork (eventlet/gevent) worker where the
             # inline branch above still ran but the blanket guard would otherwise raise.
@@ -237,6 +274,87 @@ class ReadDocumentTool(Tool):
             return {"status": "error", "error": "document parsing produced an unexpected result."}
         return result
 
+    # ------------------------------------------------------------------
+    # Inline parse bound
+    # ------------------------------------------------------------------
+    @classmethod
+    def _run_inline_bounded(cls, fn: Callable[[], Any], timeout: Optional[float]) -> Any:
+        """Run ``fn`` under a wall-clock bound, raising ``_InlineParseTimeout`` on expiry.
+
+        Prefers a SIGALRM interval timer (POSIX, main thread — prefork/solo Celery pools run
+        the task on the child's main thread), which actually interrupts the parse between
+        bytecodes and frees the slot. Otherwise falls back to a helper thread, which can only
+        stop WAITING: the orphaned parse keeps running until the parser finishes.
+
+        Args:
+            fn: Zero-argument callable performing the parse.
+            timeout: Window in seconds; ``None`` or non-positive runs unbounded.
+
+        Returns:
+            Whatever ``fn`` returns.
+        """
+        if not timeout or timeout <= 0:
+            return fn()
+        if cls._sigalrm_usable():
+            return cls._run_with_sigalrm(fn, timeout)
+        return cls._run_in_thread(fn, timeout)
+
+    @staticmethod
+    def _sigalrm_usable() -> bool:
+        """True when this thread may arm SIGALRM without clobbering another user of the timer."""
+        if not hasattr(signal, "SIGALRM") or not hasattr(signal, "setitimer"):
+            return False
+        if threading.current_thread() is not threading.main_thread():
+            return False
+        try:
+            previous = signal.getsignal(signal.SIGALRM)
+            armed = signal.getitimer(signal.ITIMER_REAL)
+        except (ValueError, OSError):  # pragma: no cover - non-main thread / no timer
+            return False
+        if previous not in (None, signal.SIG_DFL):
+            return False
+        return not (armed[0] or armed[1])
+
+    @staticmethod
+    def _run_with_sigalrm(fn: Callable[[], Any], timeout: float) -> Any:
+        """Run ``fn`` with a SIGALRM deadline, always restoring the handler and cancelling the timer."""
+
+        def _on_alarm(signum: int, frame: Any) -> None:
+            raise _InlineParseTimeout()
+
+        previous = signal.signal(signal.SIGALRM, _on_alarm)
+        signal.setitimer(signal.ITIMER_REAL, timeout)
+        try:
+            return fn()
+        finally:
+            # Nested so the handler is restored even if the alarm lands in this very
+            # window (fn returned just as the timer expired).
+            try:
+                signal.setitimer(signal.ITIMER_REAL, 0)
+            finally:
+                signal.signal(signal.SIGALRM, previous)
+
+    @staticmethod
+    def _run_in_thread(fn: Callable[[], Any], timeout: float) -> Any:
+        """Run ``fn`` in a helper thread and stop waiting at ``timeout`` (the parse is NOT killed)."""
+        executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="read-document-inline")
+        try:
+            future = executor.submit(fn)
+            try:
+                return future.result(timeout=timeout)
+            except FutureTimeoutError:
+                if future.done():  # ``fn`` itself raised a TimeoutError; surface it as a failure
+                    raise
+                logger.warning(
+                    "read_document: inline parse exceeded %.0fs off the main thread; the parse thread "
+                    "keeps running until the parser finishes (it cannot be interrupted)",
+                    timeout,
+                )
+                raise _InlineParseTimeout() from None
+        finally:
+            # Never block the agent run on the orphaned parse.
+            executor.shutdown(wait=False)
+
     # ------------------------------------------------------------------
     # Input resolution (run-scoped gate, before enqueue)
     # ------------------------------------------------------------------
@@ -260,6 +378,8 @@ class ReadDocumentTool(Tool):
                     if artifact_id is not None
                     else None
                 )
+                if artifact is not None:
+                    self._input_size = self._version_size(repo, artifact_id, artifact)
         except Exception:
             logger.exception("read_document: failed to resolve input artifact")
             return {"status": "error", "error": f"failed to resolve input artifact {raw_id}."}
@@ -274,6 +394,17 @@ class ReadDocumentTool(Tool):
             return {"status": "error", "error": f"input artifact {raw_id} not found in this conversation/run."}
         return str(artifact_id)
 
+    @staticmethod
+    def _version_size(repo: Any, artifact_id: Any, artifact: Dict[str, Any]) -> Optional[int]:
+        """Best-effort byte size of the artifact's current version (scales the parse window)."""
+        try:
+            version = repo.get_version(artifact_id, artifact.get("current_version"))
+        except Exception:
+            logger.debug("read_document: could not read input size; using the base parse window")
+            return None
+        size = (version or {}).get("size")
+        return int(size) if isinstance(size, (int, float)) else None
+
     def _bridge_chat_attachment(self, raw_id: str) -> Any:
         """Bridge a referenced chat attachment to a conversation artifact id; None on miss, error dict on failure."""
         if not self.conversation_id or not self.user_id:
@@ -281,6 +412,10 @@ class ReadDocumentTool(Tool):
         attachment = match_attachment(self.config.get("attachments"), raw_id, self.user_id)
         if attachment is None:
             return None
+        # The attachment row carries the authoritative byte size; the bridged artifact
+        # is a copy of the same bytes, so it sizes the parse window just as well.
+        size = attachment.get("size")
+        self._input_size = int(size) if isinstance(size, (int, float)) else None
         try:
             return bridge_attachment(attachment, user_id=self.user_id, conversation_id=self.conversation_id)
         except AttachmentBridgeError as exc:

application/agents/workflow_agent.py

@@ -387,8 +387,43 @@ class WorkflowAgent(BaseAgent):
                     "mime_type": ref["mime_type"],
                 }
             )
+            self._record_preextracted_text(ref["artifact_id"], attachment)
         return refs, dropped
 
+    def _record_preextracted_text(self, artifact_id: str, attachment: Dict[str, Any]) -> None:
+        """Stage an attachment's upload-time text on the engine so nodes skip re-parsing it."""
+        if self._engine is None:
+            return
+        text = self._usable_attachment_text(attachment)
+        if text:
+            self._engine.preextracted_text[str(artifact_id)] = text
+
+    @staticmethod
+    def _usable_attachment_text(attachment: Dict[str, Any]) -> Optional[str]:
+        """Return an attachment's stored text when it is safe to reuse verbatim, else None.
+
+        The attachment was already parsed (and possibly OCR'd) by ``store_attachment``;
+        reusing that text is what keeps a scanned document from being OCR'd once per
+        node. Reuse requires the stored text to be the WHOLE document: a truncated row
+        was cut head-first at ``ATTACHMENT_MAX_TOKENS`` and would silently lose the tail
+        that the workflow's own head+tail bounding keeps, so it falls back to a parse.
+        A row with no ``extraction`` metadata at all (legacy, pre-dating that field) is
+        trusted when its ``content`` is non-empty: nothing records otherwise, and the
+        non-workflow chat path already feeds that same text to the model.
+        """
+        content = attachment.get("content")
+        if not isinstance(content, str) or not content.strip():
+            return None
+        metadata = attachment.get("metadata")
+        extraction = metadata.get("extraction") if isinstance(metadata, dict) else None
+        if extraction is None:
+            return content
+        if not isinstance(extraction, dict):
+            return None
+        if extraction.get("status") != "ok" or extraction.get("truncated"):
+            return None
+        return content
+
     @staticmethod
     def _read_attachment_bytes(storage: Any, upload_path: str, max_bytes: int) -> Optional[bytes]:
         """Read an attachment with a bounded read and a guaranteed handle close; None on failure."""

application/agents/workflows/workflow_engine.py

@@ -76,6 +76,13 @@ class WorkflowEngine:
         # embedded draft), so run-scoped artifacts must NOT be persisted as orphans.
         # Defaults True; WorkflowAgent flips it off on the draft path.
         self.run_persisted: bool = True
+        # Transient, in-memory ``artifact_id -> text`` for input documents whose text
+        # was already extracted when the attachment was stored, so a node reuses it
+        # instead of re-parsing (re-OCRing) the same bytes. Deliberately NOT part of
+        # ``state``: state is snapshotted into the run row, and the text is large.
+        # Populated once per run by ``WorkflowAgent._bridge_attachments``; lifetime is
+        # this engine instance, i.e. exactly the run.
+        self.preextracted_text: Dict[str, str] = {}
         self.state: WorkflowState = {}
         self.execution_log: List[Dict[str, Any]] = []
         self._condition_result: Optional[str] = None
@@ -805,16 +812,22 @@ class WorkflowEngine:
                 attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
                 native_count += 1
             else:
-                # Inline-text mimes are read directly (cheap); other mimes route
-                # through the parsing worker -- a blocking, ~120s per-document call.
-                # Cap how many documents a single node sends down that path so a
-                # node referencing many non-native documents (e.g. the ``*`` token)
-                # can't serialize dozens of parses. Inline text is not capped.
-                needs_parse = not self._is_inline_text_mime(mime_type)
+                # Inline-text mimes are read directly (cheap), and a document whose
+                # text was already extracted when its attachment was stored is reused
+                # as-is (no parse at all); every other mime routes through the parsing
+                # worker -- a blocking, size-scaled per-document call. Cap how many
+                # documents a single node sends down THAT path so a node referencing
+                # many non-native documents (e.g. the ``*`` token) can't serialize
+                # dozens of parses. Inline text and reused text are not capped: they
+                # cost no blocking parse, so charging them would starve the budget.
+                needs_parse = (
+                    artifact_id not in self.preextracted_text
+                    and not self._is_inline_text_mime(mime_type)
+                )
                 if needs_parse:
                     # Count (and gate on) the parse ATTEMPT, not the success: a
-                    # timed-out/failed parse is the ~120s worst case we must bound,
-                    # so it has to consume cap budget too. Otherwise a degraded
+                    # timed-out/failed parse is the full-window worst case we must
+                    # bound, so it has to consume cap budget too. Otherwise a degraded
                     # parsing backend (every parse fails) never advances the count
                     # and the node keeps issuing blocking parses without limit.
                     if extract_count >= extract_max:
@@ -822,7 +835,7 @@ class WorkflowEngine:
                         continue
                     extract_count += 1
                 content = self._extract_attachment_text(
-                    artifact_id, storage_path, mime_type, filename, max_bytes
+                    artifact_id, storage_path, mime_type, filename, max_bytes, size=size
                 )
                 if content is None:
                     logger.warning(
@@ -873,12 +886,30 @@ class WorkflowEngine:
         return mime_type == "application/pdf" and supports_images
 
     def _extract_attachment_text(
-        self, artifact_id: str, storage_path: str, mime_type: str, filename: str, max_bytes: int
+        self,
+        artifact_id: str,
+        storage_path: str,
+        mime_type: str,
+        filename: str,
+        max_bytes: int,
+        size: Optional[int] = None,
     ) -> Optional[str]:
-        """Get an attachment's text: inline already-text formats, else parse via the parsing worker; None on failure."""
+        """Get an attachment's text: reuse upload-time extraction, else inline text mimes, else parse."""
         from application.parser.document_reader import truncate_text_head_tail
         from application.storage.storage_creator import StorageCreator
 
+        # An uploaded chat attachment was already parsed (and possibly OCR'd) when it
+        # was stored, so re-parsing it here would repeat the dominant cost of the run
+        # for every node that references it. Bounded with the same head+tail window the
+        # inline-text path uses.
+        preextracted = self.preextracted_text.get(artifact_id)
+        if preextracted:
+            logger.info(
+                "Workflow node: reusing upload-time extraction for %s (%s); skipping re-parse",
+                filename,
+                artifact_id,
+            )
+            return truncate_text_head_tail(preextracted)
         if self._is_inline_text_mime(mime_type):
             try:
                 data = StorageCreator.get_storage().get_file(storage_path).read()
@@ -902,13 +933,17 @@ class WorkflowEngine:
             return truncate_text_head_tail(text)
         # Non-text mimes parse via the dedicated parsing queue (works on any backend,
         # no sandbox): the worker re-resolves the artifact run-scoped and reads its bytes.
-        return self._parse_document_text(artifact_id)
+        return self._parse_document_text(artifact_id, size=size)
 
-    def _parse_document_text(self, artifact_id: str) -> Optional[str]:
-        """Enqueue ``parse_document`` for this run and await the bounded markdown; None on failure."""
+    def _parse_document_text(self, artifact_id: str, size: Optional[int] = None) -> Optional[str]:
+        """Enqueue ``parse_document`` for this run and await the size-scaled markdown; None on failure."""
         from celery.exceptions import TimeoutError as CeleryTimeoutError
 
-        from application.api.user.tasks import parse_document
+        from application.api.user.tasks import (
+            parse_document,
+            parse_task_time_limits,
+            parse_timeout_for_size,
+        )
         from application.core.settings import settings
 
         user_id = self._resolve_user_id()
@@ -916,11 +951,15 @@ class WorkflowEngine:
             return None
         options = {"output": "markdown", "include_tables": False, "persist": False}
         queue = getattr(settings, "DOCUMENT_PARSE_QUEUE", "parsing")
-        timeout = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120))
+        # OCR cost scales with pages, so the window grows with the document's size
+        # (floored at DOCUMENT_PARSE_TIMEOUT); the task's per-call time limits are
+        # raised to match, else the worker would self-terminate mid-parse.
+        timeout = parse_timeout_for_size(size)
         try:
             async_result = parse_document.apply_async(
                 args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options],
                 queue=queue,
+                **parse_task_time_limits(timeout),
             )
             # A workflow can run inside a Celery worker (scheduled runs / webhooks).
             # In a prefork worker ``task_join_will_block()`` is process-wide, so the

application/api/user/tasks.py

@@ -1,5 +1,6 @@
 import logging
 from datetime import timedelta
+from typing import Dict, Optional
 
 from sqlalchemy.exc import DataError
 
@@ -270,6 +271,47 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None):
     return resp
 
 
+# Seconds the hard (SIGKILL) limit trails the soft limit, giving the soft handler
+# room to unwind and return a terminal result before the worker is killed.
+_PARSE_HARD_LIMIT_GRACE = 30
+
+
+def parse_timeout_for_size(size_bytes: Optional[int]) -> float:
+    """Seconds to allow one ``parse_document``: floored at the base timeout, scaled by size, capped.
+
+    Args:
+        size_bytes: Byte size of the document, or None when unknown (base timeout).
+
+    Returns:
+        The parse window in seconds.
+    """
+    from application.core.settings import settings
+
+    base = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120) or 120)
+    per_mib = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 0) or 0)
+    ceiling = float(getattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", base) or base)
+    size = float(size_bytes) if isinstance(size_bytes, (int, float)) else 0.0
+    scaled = base + per_mib * max(size, 0.0) / (1024 * 1024)
+    return min(ceiling, max(base, scaled))
+
+
+def parse_task_time_limits(timeout: float) -> Dict[str, int]:
+    """Per-call Celery limits for a ``parse_document`` the caller awaits for ``timeout`` seconds.
+
+    The task's import-time ``soft_time_limit`` is bound to the BASE timeout, so a caller
+    awaiting a longer, size-scaled window must raise the per-call limits to match or the
+    worker self-terminates the parse first.
+
+    Args:
+        timeout: The awaited parse window in seconds.
+
+    Returns:
+        ``apply_async`` kwargs carrying the soft and hard time limits.
+    """
+    soft = max(1, int(timeout))
+    return {"soft_time_limit": soft, "time_limit": soft + _PARSE_HARD_LIMIT_GRACE}
+
+
 # Not DURABLE: the read_document tool awaits this synchronously with a timeout, so a
 # blind autoretry would double-parse and the caller would already have degraded. The
 # task is routed to the dedicated ``parsing`` queue (celeryconfig task_routes) so a
@@ -286,18 +328,23 @@ def parse_document(self, artifact_id, parent, user_id, options=None):
         # A pathological/malicious document must not pin a parsing-worker slot past the
         # window the caller already abandoned. Return the worker's clean error shape so
         # the slot frees and the Redis result backend still gets a terminal result.
-        limit = getattr(self, "soft_time_limit", None)
+        # ``request.timelimit`` is (hard, soft) and carries the caller's PER-CALL limit,
+        # so prefer it over the import-time default when reporting the window.
+        limit = (getattr(self.request, "timelimit", None) or (None, None))[1]
+        if limit is None:
+            limit = getattr(self, "soft_time_limit", None)
         suffix = f" after {int(limit)}s" if limit else ""
         return {"status": "error", "error": f"document parsing timed out{suffix}."}
 
 
-# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the same window read_document awaits) so
+# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the floor of the window callers await) so
 # the prefork worker self-terminates a runaway parse instead of pinning the slot; the hard
-# limit is the SIGKILL backstop if the soft handler can't unwind in time.
+# limit is the SIGKILL backstop if the soft handler can't unwind in time. Callers awaiting a
+# size-scaled window override both per call via ``parse_task_time_limits``.
 try:
     from application.core.settings import settings as _parse_settings
     parse_document.soft_time_limit = int(_parse_settings.DOCUMENT_PARSE_TIMEOUT)
-    parse_document.time_limit = parse_document.soft_time_limit + 30
+    parse_document.time_limit = parse_document.soft_time_limit + _PARSE_HARD_LIMIT_GRACE
 except Exception:
     pass
 

application/app.py

@@ -31,7 +31,10 @@ from application.api.v1 import v1_bp  # noqa: E402
 from application.celery_init import celery  # noqa: E402
 from application.core.secret_key import resolve_jwt_secret_key  # noqa: E402
 from application.core.settings import settings  # noqa: E402
-from application.storage.db.bootstrap import ensure_database_ready  # noqa: E402
+from application.storage.db.bootstrap import (  # noqa: E402
+    ensure_database_ready,
+    ensure_vector_schema,
+)
 from application.stt.upload_limits import (  # noqa: E402
     build_stt_file_size_limit_message,
     should_reject_stt_request,
@@ -59,6 +62,11 @@ ensure_database_ready(
     logger=logging.getLogger("application.app"),
 )
 
+# Own the vector DB's schema here too, so the retrieval hot path is pure reads
+# instead of re-running DDL for every source of every request.
+if settings.AUTO_VECTOR_SCHEMA:
+    ensure_vector_schema(logger=logging.getLogger("application.app"))
+
 from application.agents.default_tools import (  # noqa: E402
     validate_default_chat_tools,
 )

application/cache.py

@@ -26,6 +26,73 @@ def _cache_default(value):
         return f"<bytes:sha256:{hashlib.sha256(bytes(value)).hexdigest()}>"
     return repr(value)
 
+
+# Generation kwargs that never reach the provider: usage-accounting side
+# channels only. Everything else the caller passes (``response_format``,
+# ``response_schema``, ``tool_choice``, ``reasoning_effort``, sampling
+# params, ...) is part of the request and therefore part of the key —
+# otherwise a workflow node that changed its JSON schema replays the old
+# schema's cached answer for the whole TTL.
+_CACHE_KEY_IGNORED_KWARGS = frozenset({"_usage_attachments", "attachments"})
+
+# Kwargs that make the answer depend on provider-held state no key can
+# capture. ``previous_response_id`` chains a Responses API turn server
+# side, and a cache hit would also skip the ``_last_response_id``
+# bookkeeping the next turn needs — so skip the cache entirely.
+_CACHE_BYPASS_KWARGS = ("previous_response_id",)
+
+
+def _gen_kwargs_fingerprint(extra: dict | None) -> str:
+    """Stable fingerprint of the generation-affecting kwargs.
+
+    Args:
+        extra: Keyword arguments forwarded to the generation call.
+
+    Returns:
+        A sorted JSON dump of the semantic kwargs, or "" when there are none.
+
+    Raises:
+        ValueError: If the kwargs cannot be serialized (callers treat this
+            as "do not cache").
+    """
+    if not extra:
+        return ""
+    filtered = {
+        key: value
+        for key, value in extra.items()
+        if key not in _CACHE_KEY_IGNORED_KWARGS and value is not None
+    }
+    if not filtered:
+        return ""
+    try:
+        return json.dumps(filtered, sort_keys=True, default=_cache_default)
+    except (TypeError, ValueError) as e:
+        raise ValueError(f"Unserializable generation kwargs: {e}") from e
+
+
+def _bypasses_cache(extra: dict | None) -> bool:
+    """Whether a kwarg ties the call to provider-side conversation state."""
+    if not extra:
+        return False
+    return any(extra.get(key) for key in _CACHE_BYPASS_KWARGS)
+
+
+def _is_stream_payload(raw: str) -> bool:
+    """Whether a cached value is a ``stream_cache`` chunk envelope.
+
+    ``gen_cache`` and ``stream_cache`` share one key space, so a
+    non-streaming read can land on a stream entry; without this guard the
+    serialized chunk list would be handed back as the model's answer.
+    """
+    if not raw.startswith("{"):
+        return False
+    try:
+        decoded = json.loads(raw)
+    except ValueError:
+        return False
+    return isinstance(decoded, dict) and isinstance(decoded.get("chunks"), list)
+
+
 _redis_instance = None
 _redis_creation_failed = False
 _instance_lock = Lock()
@@ -116,23 +183,44 @@ def get_pubsub_redis_instance():
     return _pubsub_redis_instance
 
 
-def gen_cache_key(messages, model="docgpt", tools=None):
+def gen_cache_key(messages, model="docgpt", tools=None, extra=None):
+    """Build the Redis key for one generation call.
+
+    Args:
+        messages: Chat messages for the call.
+        model: Model identifier.
+        tools: Tool schemas, when the call carries any.
+        extra: Remaining generation kwargs (``response_format``,
+            ``response_schema``, sampling params, ...). Non-semantic keys
+            are dropped before hashing; the suffix is omitted entirely when
+            nothing semantic remains, so keys for plain calls are unchanged.
+
+    Returns:
+        Hex digest used as the cache key.
+
+    Raises:
+        ValueError: If ``messages`` holds a non-dict entry, or ``extra``
+            cannot be serialized.
+    """
     if not all(isinstance(msg, dict) for msg in messages):
         raise ValueError("All messages must be dictionaries.")
     messages_str = json.dumps(messages, default=_cache_default)
     tools_str = json.dumps(str(tools)) if tools else ""
     combined = f"{model}_{messages_str}_{tools_str}"
+    extra_str = _gen_kwargs_fingerprint(extra)
+    if extra_str:
+        combined = f"{combined}_{extra_str}"
     cache_key = get_hash(combined)
     return cache_key
 
 
 def gen_cache(func):
     def wrapper(self, model, messages, stream, tools=None, *args, **kwargs):
-        if tools is not None:
+        if tools is not None or _bypasses_cache(kwargs):
             return func(self, model, messages, stream, tools, *args, **kwargs)
-        
+
         try:
-            cache_key = gen_cache_key(messages, model, tools)
+            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
         except ValueError as e:
             logger.error(f"Cache key generation failed: {e}")
             return func(self, model, messages, stream, tools, *args, **kwargs)
@@ -142,7 +230,9 @@ def gen_cache(func):
             try:
                 cached_response = redis_client.get(cache_key)
                 if cached_response:
-                    return cached_response.decode("utf-8")
+                    decoded = cached_response.decode("utf-8")
+                    if not _is_stream_payload(decoded):
+                        return decoded
             except Exception as e:
                 logger.error(f"Error getting cached response: {e}", exc_info=True)
 
@@ -160,12 +250,12 @@ def gen_cache(func):
 
 def stream_cache(func):
     def wrapper(self, model, messages, stream, tools=None, *args, **kwargs):
-        if tools is not None:
+        if tools is not None or _bypasses_cache(kwargs):
             yield from func(self, model, messages, stream, tools, *args, **kwargs)
             return
-        
+
         try:
-            cache_key = gen_cache_key(messages, model, tools)
+            cache_key = gen_cache_key(messages, model, tools, extra=kwargs)
         except ValueError as e:
             logger.error(f"Cache key generation failed: {e}")
             yield from func(self, model, messages, stream, tools, *args, **kwargs)

application/core/settings.py

@@ -75,6 +75,9 @@ class Settings(BaseSettings):
     AUTO_MIGRATE: bool = True
     # On app startup, create the target Postgres database if it's missing (requires CREATEDB privilege). Dev-friendly default.
     AUTO_CREATE_DB: bool = True
+    # On app startup, create the pgvector/graph tables and verify the embedding dimension. Set False to manage them
+    # out-of-band — there is no Alembic migration for the vector DB because it may be a separate cluster.
+    AUTO_VECTOR_SCHEMA: bool = True
     LLM_PATH: str = os.path.join(current_dir, "models/docsgpt-7b-f16.gguf")
     DEFAULT_MAX_HISTORY: int = 150
     DEFAULT_LLM_TOKEN_LIMIT: int = 128000  # Fallback when model not found in registry
@@ -108,6 +111,12 @@ class Settings(BaseSettings):
     DOCLING_COMPILE_TORCH_MODELS: bool = False
     DOCLING_TABULAR_MAX_BYTES: int = 2_000_000
     DOCLING_MARKUP_MAX_BYTES: int = 8_000_000
+    # Chars-per-page floor below which an OCR'd PDF/image parse is treated as a docling
+    # pipeline dropout (long-running workers were observed returning zero characters for
+    # every scanned page after a long scanned PDF, with no error) rather than as content.
+    # Such a parse is retried once on a fresh full-page-OCR converter and then fails
+    # loudly instead of indexing an empty document. 0 disables the guard.
+    DOCLING_OCR_MIN_CHARS_PER_PAGE: int = 20
     # Read PDF *attachments* via their embedded text layer (pypdfium2) instead
     # of docling, falling back to docling when there is no text layer to read.
     # Attachments go into a prompt, so docling's structural markdown earns far
@@ -127,6 +136,9 @@ class Settings(BaseSettings):
     # ``RetrieverCreator.retrievers`` registry keys (``classic`` / ``default``),
     # NOT the legacy ``classic_rag`` label which never matched the registry.
     RETRIEVERS_ENABLED: list = ["classic", "default"]
+    # Concurrent per-source vector searches within one retrieval (multi-source chats);
+    # the query is embedded once and shared across sources.
+    RETRIEVAL_MAX_PARALLEL_SOURCES: int = 4
     # Kill-switch for per-source retrieval dispatch. When False the retrieval
     # path collapses to today's single-retriever behavior (consumed by the
     # Dispatcher in a later change; defined here so the flag exists up front).
@@ -225,6 +237,9 @@ class Settings(BaseSettings):
     # dialect form (``postgresql+psycopg://``) are all accepted and
     # normalized internally for ``psycopg.connect()``.
     PGVECTOR_CONNECTION_STRING: Optional[str] = None
+    # Per-process psycopg connection pool for the vector store; 0 = one direct
+    # connection per store instance, the legacy behaviour.
+    PGVECTOR_POOL_MAX_SIZE: int = 8
     # IVFFlat probes for vector search. ``None`` derives sqrt(lists) from the
     # index itself; set an integer to pin it. Higher = better recall, more scan.
     PGVECTOR_IVFFLAT_PROBES: Optional[int] = None
@@ -428,6 +443,12 @@ class Settings(BaseSettings):
     # ``read_document`` parsing on a dedicated Celery ``parsing`` queue (backend parser).
     DOCUMENT_PARSE_QUEUE: str = "parsing"  # queue the parse_document task is routed to
     DOCUMENT_PARSE_TIMEOUT: int = 120  # seconds the tool awaits the enqueued parse before degrading
+    # The base timeout is a FLOOR: the awaited window (and the task's per-call Celery
+    # time limits) grow with the document's size, because OCR cost scales with pages
+    # -- a 30-page scan needs ~60s, a 100-page scan several minutes. Without this a
+    # large scan is silently dropped from the node/tool at the base window.
+    DOCUMENT_PARSE_TIMEOUT_PER_MB: int = 60  # extra seconds of parse window per MiB of input
+    DOCUMENT_PARSE_TIMEOUT_MAX: int = 900  # absolute ceiling on the size-scaled parse window
     DOCUMENT_PARSE_MAX_BYTES: int = 0  # cap on a parsed document's bytes (0 = reuse SANDBOX_MAX_INPUT_BYTES)
     DOCUMENT_MAX_DECOMPRESSED_BYTES: int = 300 * 1024 * 1024
     DOCUMENT_MAX_ARCHIVE_ENTRIES: int = 10000

application/graphrag/extraction.py

@@ -26,7 +26,9 @@ from typing import Any, Callable, Dict, List, Optional
 from application.core.settings import settings
 from application.llm.llm_creator import LLMCreator
 from application.storage.db.source_config import SourceConfig
-from application.vectorstore.base import EmbeddingsSingleton
+# ``EmbeddingsSingleton`` is re-exported here so callers and tests can reach the
+# shared instance cache from this module.
+from application.vectorstore.base import EmbeddingsSingleton, get_embeddings  # noqa: F401
 
 logger = logging.getLogger(__name__)
 
@@ -193,9 +195,7 @@ def extract_graph_for_source(
     skipped_over_cap = max(0, len(pending) - cap)
     to_process = pending[:cap]
 
-    embedding = EmbeddingsSingleton.get_instance(
-        settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-    )
+    embedding = get_embeddings()
 
     llm = _build_extraction_llm(
         _resolve_extraction_model(config), user, request_id

application/graphrag/store.py

@@ -1,12 +1,14 @@
 """Per-source knowledge-graph store co-located with the pgvector ``documents`` table.
 
-GraphRAG is pgvector-only: the graph tables live in the same DB as the
-pgvector store and are created on-demand (``CREATE TABLE IF NOT EXISTS`` +
-``CREATE EXTENSION IF NOT EXISTS vector``), mirroring
-``PGVectorStore._ensure_table_exists`` rather than going through app-DB Alembic.
-That DB may be a separate cluster (e.g. Neon) from the app DB where ``sources``
-lives, so ``source_id`` is a plain indexed UUID column with no cross-DB FK and
-all ids are generated in Python.
+GraphRAG is pgvector-only: the graph tables live in the same DB as the pgvector
+store and are created at boot by ``ensure_vector_schema`` (``CREATE TABLE IF NOT
+EXISTS`` + ``CREATE EXTENSION IF NOT EXISTS vector``), mirroring
+``PGVectorStore.create_schema`` rather than going through app-DB Alembic.
+Constructing a store runs no DDL and opens no connection; the write methods
+re-check the schema once per instance as a safety net for a process that never
+ran the boot hook. That DB may be a separate cluster (e.g. Neon) from the app DB
+where ``sources`` lives, so ``source_id`` is a plain indexed UUID column with no
+cross-DB FK and all ids are generated in Python.
 """
 
 from __future__ import annotations
@@ -18,9 +20,15 @@ from typing import Any, Dict, List, Optional
 from psycopg.types.json import Jsonb
 
 from application.core.settings import settings
+from application.vectorstore import pgconn
 
 DEFAULT_NAME_EMBEDDING_DIM = 768
 
+# Mirrors ``PGVectorStore.DEFAULT_POOL_MAX_SIZE``; duplicated rather than
+# imported so the graph store never pulls the vector store (and its embeddings
+# stack) in at import time. 0 disables pooling.
+DEFAULT_POOL_MAX_SIZE = 8
+
 MAX_SUBGRAPH_NODES = 500
 MAX_SUBGRAPH_EDGES = 2000
 
@@ -30,6 +38,14 @@ GRAPH_OVERVIEW_MAX_LIMIT = 250
 PGVECTOR_SOURCE_COLUMN = "source_id"
 
 
+def _resolve_pool_max_size() -> int:
+    """Pool size from settings, defensively — 0 means one direct connection."""
+    value = getattr(settings, "PGVECTOR_POOL_MAX_SIZE", DEFAULT_POOL_MAX_SIZE)
+    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+        return value
+    return DEFAULT_POOL_MAX_SIZE
+
+
 def _safe_identifier(name: str) -> str:
     """Return ``name`` if it is a bare SQL identifier, else raise.
 
@@ -97,35 +113,74 @@ class GraphStore:
         self._psycopg = psycopg
         self._register_vector = register_vector
         self._connection = None
-        self._ensure_tables()
+        self._pooled = False
+        self._pool_max_size = _resolve_pool_max_size()
+        # No DDL here: the graph tables are created at boot alongside the
+        # pgvector table, and every write re-checks once (see _ensure_tables_once).
+        self._tables_ensured = False
 
     def _get_connection(self):
+        """Get or create this store's connection, pooled unless pooling is off.
+
+        Shares :mod:`application.vectorstore.pgconn`'s per-DSN pool with
+        ``PGVectorStore``, so a retrieval that touches both pays one checkout
+        each instead of two fresh connect handshakes.
+        """
         if self._connection is None or self._connection.closed:
-            self._connection = self._psycopg.connect(self._connection_string)
-            self._register_vector(self._connection)
+            if self._pool_max_size > 0:
+                self._connection = pgconn.pool_for(
+                    self._connection_string, self._pool_max_size
+                ).getconn()
+                self._pooled = True
+                # No _register_pgvector_types here: the pool's ``configure``
+                # hook already ran it on this physical connection, and repeating
+                # it would cost a catalog lookup on every checkout. The DDL path
+                # (_ensure_tables) still re-registers after CREATE EXTENSION.
+            else:
+                self._connection = self._psycopg.connect(self._connection_string)
+                self._register_pgvector_types(self._connection)
+                self._pooled = False
         return self._connection
 
+    def _register_pgvector_types(self, conn) -> None:
+        """Register pgvector's adapters, tolerating a not-yet-created extension.
+
+        ``register_vector`` looks the ``vector`` type up in the catalog and
+        raises when it is absent — which is exactly the state of a brand-new
+        database, before ``CREATE EXTENSION`` has run. Swallow that so the
+        schema bootstrap can proceed; it re-registers once the type exists.
+        """
+        try:
+            self._register_vector(conn)
+        except Exception as e:
+            logging.debug("pgvector types not registered yet: %s", e)
+
     def _embedding_dim(self) -> int:
         """Dimension of the configured embeddings model, matching ``PGVectorStore``.
 
         Falls back to ``DEFAULT_NAME_EMBEDDING_DIM`` so the graph table and the
         pgvector ``documents`` table always agree on the configured model.
         """
-        from application.vectorstore.base import EmbeddingsSingleton
+        from application.vectorstore.base import get_embeddings
 
-        embedding = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-        )
+        embedding = get_embeddings()
         return getattr(embedding, "dimension", DEFAULT_NAME_EMBEDDING_DIM)
 
-    def _ensure_tables(self):
-        conn = self._get_connection()
+    @staticmethod
+    def create_schema(conn, *, dimension: int = DEFAULT_NAME_EMBEDDING_DIM) -> None:
+        """Create the graph tables and indexes on ``conn`` without committing.
+
+        Shared by the boot hook (``ensure_vector_schema``) and the store's own
+        write-path safety net; the caller owns the transaction.
+
+        Args:
+            conn: Open psycopg connection to the pgvector database.
+            dimension: Width of the node name-embedding vectors.
+        """
         cursor = conn.cursor()
         try:
             cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
 
-            embedding_dim = self._embedding_dim()
-
             cursor.execute(
                 f"""
                 CREATE TABLE IF NOT EXISTS graph_nodes (
@@ -137,7 +192,7 @@ class GraphStore:
                     description TEXT,
                     degree INT DEFAULT 0,
                     doc_freq INT DEFAULT 0,
-                    name_embedding vector({embedding_dim}),
+                    name_embedding vector({dimension}),
                     UNIQUE (source_id, normalized_name)
                 );
                 """
@@ -206,14 +261,44 @@ class GraphStore:
             # post-filter in search_nodes_by_embedding that silently returns
             # zero nodes, and graph_rag does not fall back when the source has
             # nodes. Add an index deliberately once a graph is large enough.
+        finally:
+            cursor.close()
 
+    def _ensure_tables(self):
+        """Create the graph schema under an advisory lock, then commit."""
+        # Same key as the pgvector store and the boot hook: one lock guards all
+        # DDL in this database, so concurrent workers never race each other.
+        from application.vectorstore.pgvector import SCHEMA_LOCK_KEY
+
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+                )
+            finally:
+                cursor.close()
+            self.create_schema(conn, dimension=self._embedding_dim())
             conn.commit()
+            # The extension may have just been created; pick up its adapters.
+            self._register_pgvector_types(conn)
         except Exception as e:
             conn.rollback()
             logging.error(f"Error creating graph tables: {e}")
             raise
-        finally:
-            cursor.close()
+
+    def _ensure_tables_once(self) -> None:
+        """Create the graph schema on this instance's first write, at most once.
+
+        Readers never create tables — boot owns the schema. This is the safety
+        net for a process that never ran the boot hook (scripts, tests, the
+        first extraction on a brand-new deployment).
+        """
+        if getattr(self, "_tables_ensured", False):
+            return
+        self._ensure_tables()
+        self._tables_ensured = True
 
     def _upsert_node(
         self,
@@ -287,6 +372,7 @@ class GraphStore:
         incremented, the type is refreshed if previously empty, and the
         embedding is refreshed when provided. Returns the node id either way.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -351,6 +437,7 @@ class GraphStore:
         source_chunk_ids: Optional[List[str]] = None,
     ) -> str:
         """Insert an edge and bump the degree of both endpoints. Returns its id."""
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -384,6 +471,7 @@ class GraphStore:
         )
 
     def link_node_chunk(self, source_id: str, node_id: str, chunk_id: str):
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -415,6 +503,7 @@ class GraphStore:
         are not bumped here — the caller runs ``set_node_degrees`` once at the
         end. Returns ``(nodes_upserted, edges_added)``.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         node_ids: Dict[str, str] = {}
@@ -540,6 +629,53 @@ class GraphStore:
             cursor.close()
             conn.rollback()
 
+    def count_nodes_many(self, source_ids: List[str]) -> Dict[str, int]:
+        """Node counts for several sources in one round trip.
+
+        Replaces the retriever's per-source ``count_nodes`` fan-out: N sources
+        used to cost N queries (each on its own fresh connection). Ids with no
+        rows are filled in with 0 in Python, so the caller always gets an entry
+        for every id it asked about.
+
+        Args:
+            source_ids: Source ids to count; empty/falsy entries are ignored.
+
+        Returns:
+            dict: ``{source_id: node_count}``, zero-filled. All zeros when the
+            query fails, which drives the ClassicRAG fallback exactly as a
+            failing ``count_nodes`` does.
+        """
+        ids = [str(s) for s in source_ids if s]
+        if not ids:
+            return {}
+        counts: Dict[str, int] = {source_id: 0 for source_id in ids}
+        # ``source_id`` is a UUID column, so Postgres hands back the canonical
+        # lowercase text. Map it to the exact string the caller passed, or a
+        # differently-cased id would land under a second key and read as 0.
+        by_canonical = {source_id.lower(): source_id for source_id in ids}
+        conn = self._get_connection()
+        cursor = conn.cursor()
+        try:
+            cursor.execute(
+                """
+                SELECT source_id, count(*)
+                FROM graph_nodes
+                WHERE source_id = ANY(%s)
+                GROUP BY source_id;
+                """,
+                (ids,),
+            )
+            for row in cursor.fetchall():
+                returned = str(row[0])
+                counts[by_canonical.get(returned.lower(), returned)] = int(row[1])
+            return counts
+        except Exception as e:
+            logging.error(f"Error counting nodes: {e}")
+            return counts
+        finally:
+            cursor.close()
+            conn.rollback()
+
     def search_nodes_by_embedding(
         self, source_id: str, query_embedding: List[float], k: int = 10
     ) -> List[Dict[str, Any]]:
@@ -860,6 +996,7 @@ class GraphStore:
         (``WHERE id IN (src, dst)`` bumps the endpoint a single time when
         ``src == dst``). ``UNION`` deduplicates the two endpoints of each edge.
         """
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -891,6 +1028,7 @@ class GraphStore:
             cursor.close()
 
     def mark_chunk(self, source_id: str, chunk_id: str, status: str):
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -952,6 +1090,7 @@ class GraphStore:
 
     def delete_by_source(self, source_id: str):
         """Remove every graph row for a source (no FK cascade across clusters)."""
+        self._ensure_tables_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         try:
@@ -972,10 +1111,26 @@ class GraphStore:
         finally:
             cursor.close()
 
+    def close(self) -> None:
+        """Release this store's connection: back to the pool, or closed outright.
+
+        A pooled connection is rolled back first when it is still in a
+        transaction, so the next borrower gets a clean session.
+        """
+        conn = getattr(self, "_connection", None)
+        if conn is None:
+            return
+        self._connection = None
+        pgconn.release(
+            getattr(self, "_connection_string", ""),
+            conn,
+            getattr(self, "_pooled", False),
+        )
+
     def __del__(self):
-        if (
-            hasattr(self, "_connection")
-            and self._connection
-            and not self._connection.closed
-        ):
-            self._connection.close()
+        """Release the connection when the object is destroyed. Never raises."""
+        try:
+            self.close()
+        except Exception:
+            # Interpreter teardown can null out module globals; never raise here.
+            pass

application/llm/base.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import ClassVar
+from typing import ClassVar, Dict, Optional, Tuple
 
 import httpx
 import openai
@@ -38,6 +38,24 @@ class BaseLLM(ABC):
     # calls by vendor. Subclasses override.
     provider_name: ClassVar[str] = "unknown"
 
+    # Name of the gen kwarg this provider takes structured output on
+    # ("response_format" for OpenAI-wire classes, "response_schema" for
+    # Google); None = the provider has no structured-output kwarg.
+    structured_output_kwarg: ClassVar[Optional[str]] = None
+
+    # (json_schema, strict) last passed to ``prepare_structured_output_format``;
+    # lets a cross-provider fallback re-prepare the schema in the backup's own
+    # format instead of forwarding a kwarg the backup cannot read.
+    _structured_output_source: Optional[Tuple[Dict, bool]] = None
+
+    # Gen kwargs that carry structured output, across providers. Kept here
+    # (not on the provider subclasses) because the adapter has to recognize
+    # every provider's kwarg, not just its own.
+    _STRUCTURED_OUTPUT_KWARGS: ClassVar[Tuple[str, ...]] = (
+        "response_format",
+        "response_schema",
+    )
+
     def __init__(
         self,
         decoded_token=None,
@@ -300,6 +318,132 @@ class BaseLLM(ABC):
                 prepared.append({**message, "content": parts})
         return prepared
 
+    @staticmethod
+    def _fallback_enforces_structured_output(fallback) -> bool:
+        """Whether the fallback's capabilities still allow schema enforcement."""
+        supports = getattr(fallback, "_supports_structured_output", None)
+        if supports is None:
+            return True
+        if not callable(supports):
+            return bool(supports)
+        try:
+            return bool(supports())
+        except Exception:
+            logger.debug(
+                "Structured-output capability check failed for %s; assuming supported",
+                type(fallback).__name__,
+                exc_info=True,
+            )
+            return True
+
+    @staticmethod
+    def _recover_structured_output_source(present: Dict) -> Optional[Tuple[Dict, bool]]:
+        """Recover a raw (schema, strict) pair from already-prepared kwargs.
+
+        Callers that hand-build an OpenAI ``response_format`` never went through
+        ``prepare_structured_output_format``, so nothing was recorded; the raw
+        schema is still readable out of the envelope. A Google
+        ``response_schema`` is a lossy, type-mapped conversion — not reversible —
+        so it yields nothing.
+        """
+        response_format = present.get("response_format")
+        if not isinstance(response_format, dict):
+            return None
+        json_schema = response_format.get("json_schema")
+        if not isinstance(json_schema, dict):
+            return None
+        schema = json_schema.get("schema")
+        if not isinstance(schema, dict) or not schema:
+            return None
+        return schema, bool(json_schema.get("strict", True))
+
+    def _adapt_structured_output_kwargs(self, fallback, kwargs: Dict) -> Dict:
+        """Re-express the primary's structured-output kwargs for ``fallback``.
+
+        Structured output is provider-specific: OpenAI-wire classes take
+        ``response_format``, Google takes ``response_schema``. Forwarding the
+        primary's kwarg verbatim to a different-family backup either loses
+        enforcement silently (Google swallows ``response_format`` in
+        ``**kwargs``) or raises ``TypeError`` inside the OpenAI SDK
+        (``response_schema`` is not a Chat-Completions param), which turns the
+        fallback into no fallback at all.
+
+        Args:
+            fallback: Backup LLM the request is about to be re-sent to.
+            kwargs: The primary's generation kwargs.
+
+        Returns:
+            A copy of ``kwargs`` carrying the backup's own structured-output
+            kwarg, or neither when the backup cannot enforce a schema.
+        """
+        adapted = dict(kwargs)
+        present = {
+            name: adapted.pop(name)
+            for name in self._STRUCTURED_OUTPUT_KWARGS
+            if name in adapted
+        }
+        if not present:
+            return adapted
+
+        target = getattr(fallback, "structured_output_kwarg", None)
+        if not target or not self._fallback_enforces_structured_output(fallback):
+            logger.warning(
+                f"Fallback {fallback.model_id} cannot enforce structured "
+                f"output; continuing unstructured"
+            )
+            return adapted
+
+        response_format = present.get("response_format")
+        if (
+            isinstance(response_format, dict)
+            and response_format.get("type") == "json_object"
+        ):
+            # json_object mode is an OpenAI-wire shape with no wired equivalent
+            # elsewhere; keep it only within the same family.
+            if target == "response_format":
+                adapted[target] = response_format
+            else:
+                logger.warning(
+                    f"Fallback {fallback.model_id} has no json_object mode; "
+                    f"continuing unstructured"
+                )
+            return adapted
+
+        if target in present:
+            # Same wire family (OpenAI -> openai_compatible, Google -> Google):
+            # the prepared value is already in the backup's format.
+            adapted[target] = present[target]
+            return adapted
+
+        source = self._structured_output_source
+        if not source:
+            source = self._recover_structured_output_source(present)
+        if not source:
+            logger.warning(
+                f"Fallback {fallback.model_id} needs {target} but the source "
+                f"schema is unavailable; continuing unstructured"
+            )
+            return adapted
+
+        schema, strict = source
+        try:
+            prepared = fallback.prepare_structured_output_format(schema, strict=strict)
+        except Exception:
+            logger.warning(
+                f"Failed to prepare structured output for fallback "
+                f"{fallback.model_id}; continuing unstructured",
+                exc_info=True,
+            )
+            prepared = None
+        if prepared:
+            adapted[target] = prepared
+        else:
+            logger.warning(
+                f"Fallback {fallback.model_id} cannot enforce structured "
+                f"output; continuing unstructured"
+            )
+        return adapted
+
     def _execute_with_fallback(
         self, method_name: str, decorators: list, *args, **kwargs
     ):
@@ -371,6 +515,9 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            fallback_kwargs = self._adapt_structured_output_kwargs(
+                fallback, fallback_kwargs
+            )
             if fallback_kwargs.get("messages"):
                 fallback_kwargs["messages"] = self._prepare_fallback_messages(
                     fallback,
@@ -502,6 +649,9 @@ class BaseLLM(ABC):
             for decorator in decorators:
                 fallback_method = decorator(fallback_method)
             fallback_kwargs = {**kwargs, "model": fallback.model_id}
+            fallback_kwargs = self._adapt_structured_output_kwargs(
+                fallback, fallback_kwargs
+            )
             if fallback_kwargs.get("messages"):
                 fallback_kwargs["messages"] = self._prepare_fallback_messages(
                     fallback,
@@ -673,9 +823,13 @@ class BaseLLM(ABC):
     def _supports_structured_output(self):
         return False
 
-    def prepare_structured_output_format(self, json_schema):
-        """Prepare structured output format specific to the LLM provider"""
-        _ = json_schema
+    def prepare_structured_output_format(self, json_schema, strict=True):
+        """Prepare structured output format specific to the LLM provider.
+
+        Overrides must record ``self._structured_output_source`` so a
+        cross-provider fallback can re-prepare the schema in its own format.
+        """
+        _ = (json_schema, strict)
         return None
 
     def get_supported_attachment_types(self):

application/llm/google_ai.py

@@ -12,6 +12,7 @@ from application.storage.storage_creator import StorageCreator
 
 class GoogleLLM(BaseLLM):
     provider_name = "google"
+    structured_output_kwarg = "response_schema"
 
     def __init__(
         self, api_key=None, user_api_key=None, decoded_token=None, *args, **kwargs
@@ -674,6 +675,8 @@ class GoogleLLM(BaseLLM):
         ``strict`` is accepted for signature parity with the OpenAI provider;
         Google enforces the schema natively via ``response_schema``.
         """
+        # Recorded so a cross-provider fallback can re-prepare the raw schema.
+        self._structured_output_source = (json_schema, strict) if json_schema else None
         if not json_schema:
             return None
         type_map = {

application/llm/openai.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os.path
 
-from openai import OpenAI
+from openai import BadRequestError, OpenAI
 
 from application.core.settings import settings
 from application.llm.base import BaseLLM
@@ -14,6 +14,50 @@ from application.storage.storage_creator import StorageCreator
 # Placeholder sent to OpenAI-compatible backends that require no credentials.
 NO_API_KEY = "sk-no-key"
 
+# Request keys that only make sense alongside ``tools``; all dropped together
+# when an endpoint turns out not to serve tool calling.
+_TOOL_REQUEST_KEYS = ("tools", "tool_choice", "parallel_tool_calls")
+
+# Substrings an OpenAI-compatible server puts in a 400 when it understands the
+# request but cannot serve tool calling: vLLM without --enable-auto-tool-choice
+# / --tool-call-parser, llama.cpp, TGI, Ollama and friends. Matched case
+# insensitively, and only for requests that actually carried tools — so a false
+# positive degrades to a tool-less answer rather than a failed chat.
+_TOOLS_UNSUPPORTED_MARKERS = (
+    "enable-auto-tool-choice",
+    "tool-call-parser",
+    "tool_choice",
+    "tool choice",
+    "tools are not supported",
+    "tools is not supported",
+    "does not support tools",
+    "doesn't support tools",
+    "tool calling",
+    "tool call is not supported",
+    "function calling",
+    "functions are not supported",
+    "does not support function",
+    "unsupported parameter: 'tools'",
+    "unknown field: tools",
+)
+
+
+def _is_tools_unsupported_error(error: Exception) -> bool:
+    """Whether a 400 says the endpoint cannot serve tool calling.
+
+    Args:
+        error: The exception raised by the provider SDK.
+
+    Returns:
+        True when the server's message matches a tool-calling marker.
+    """
+    parts = [str(error)]
+    body = getattr(error, "body", None)
+    if isinstance(body, dict):
+        parts.append(str(body.get("message") or ""))
+    haystack = " ".join(parts).lower()
+    return any(marker in haystack for marker in _TOOLS_UNSUPPORTED_MARKERS)
+
 
 def _truncate_base64_for_logging(messages):
     """
@@ -115,6 +159,12 @@ class _RespChoice:
 
 class OpenAILLM(BaseLLM):
     provider_name = "openai"
+    structured_output_kwarg = "response_format"
+
+    # Flipped once this endpoint answered 400 on a request carrying tools.
+    # Instance-scoped (the assignment shadows this class default), so an
+    # agent loop pays for the discovery round trip at most once.
+    _tools_rejected = False
 
     def __init__(
         self,
@@ -562,6 +612,53 @@ class OpenAILLM(BaseLLM):
                 return normalized
         return ""
 
+    def _note_tools_rejected(self, model, error) -> None:
+        """Record that this endpoint refuses tool calling, warning once."""
+        if not self._tools_rejected:
+            logging.warning(
+                "Endpoint rejected tool calling (model=%s, base_url=%s): %s. "
+                "Retrying this request without tools and skipping tools for "
+                "the rest of this session — set supports_tools: false in the "
+                "model catalog for %s to avoid the discovery round trip.",
+                model,
+                self._effective_base_url,
+                error,
+                model,
+            )
+        self._tools_rejected = True
+
+    @staticmethod
+    def _params_without_tools(params: dict) -> dict:
+        """Copy of ``params`` with every tool-calling key removed."""
+        return {
+            key: value for key, value in params.items()
+            if key not in _TOOL_REQUEST_KEYS
+        }
+
+    def _create_with_tool_fallback(self, create, params: dict, model):
+        """Send a request, retrying once without tools on a tools-unsupported 400.
+
+        Args:
+            create: The provider ``create`` callable for this endpoint.
+            params: Request kwargs, possibly including ``tools``.
+            model: Model id, used for the warning.
+
+        Returns:
+            The provider response — from the first call, or from the
+            tool-less retry when the endpoint rejected tool calling.
+
+        Raises:
+            openai.BadRequestError: Any 400 unrelated to tool calling, and
+                any 400 on a request that carried no tools.
+        """
+        try:
+            return create(**params)
+        except BadRequestError as error:
+            if not params.get("tools") or not _is_tools_unsupported_error(error):
+                raise
+            self._note_tools_rejected(model, error)
+            return create(**self._params_without_tools(params))
+
     def _raw_gen(
         self,
         baseself,
@@ -620,7 +717,9 @@ class OpenAILLM(BaseLLM):
             request_params["response_format"] = response_format
         self._last_usage = None
         self._stream_reached_finish = False
-        response = self.client.chat.completions.create(**request_params)
+        response = self._create_with_tool_fallback(
+            self.client.chat.completions.create, request_params, model
+        )
         logging.debug("OpenAI request completed")
         self._record_chat_usage(getattr(response, "usage", None))
         if tools:
@@ -694,7 +793,12 @@ class OpenAILLM(BaseLLM):
         self._last_usage = None
         self._stream_reached_finish = False
         self._last_finish_reason = None
-        response = self.client.chat.completions.create(**request_params)
+        # The request is issued before the loop, so a tools-unsupported 400
+        # lands here with nothing yielded yet — the retry inside
+        # ``_create_with_tool_fallback`` can never duplicate delivered output.
+        response = self._create_with_tool_fallback(
+            self.client.chat.completions.create, request_params, model
+        )
 
         try:
             for line in response:
@@ -1298,7 +1402,9 @@ class OpenAILLM(BaseLLM):
             stream=False,
             kwargs=kwargs,
         )
-        response = self.client.responses.create(**params)
+        response = self._create_with_tool_fallback(
+            self.client.responses.create, params, model
+        )
         if response is None:
             raise RuntimeError("Responses API returned no response object")
         logging.debug(
@@ -1371,7 +1477,11 @@ class OpenAILLM(BaseLLM):
             stream=True,
             kwargs=kwargs,
         )
-        response = self.client.responses.create(**params)
+        # Issued before the event loop below, so the tool-less retry cannot
+        # duplicate output already yielded.
+        response = self._create_with_tool_fallback(
+            self.client.responses.create, params, model
+        )
 
         func_calls = {}
         reasoning_items = []
@@ -1485,6 +1595,12 @@ class OpenAILLM(BaseLLM):
                 response.close()
 
     def _supports_tools(self):
+        # A 400 already told us this endpoint can't serve tool calling
+        # (vLLM without --enable-auto-tool-choice, and friends). Answer
+        # False from then on so later turns skip tools outright instead of
+        # paying for another rejected round trip.
+        if self._tools_rejected:
+            return False
         # When the LLM was constructed via LLMCreator with a registered
         # AvailableModel, ``self.capabilities`` is the per-model record.
         # BYOM users can disable tool support; respect that. Otherwise
@@ -1515,6 +1631,8 @@ class OpenAILLM(BaseLLM):
             kwargs["reasoning_effort"] = effort
 
     def prepare_structured_output_format(self, json_schema, strict=True):
+        # Recorded so a cross-provider fallback can re-prepare the raw schema.
+        self._structured_output_source = (json_schema, strict) if json_schema else None
         if not json_schema:
             return None
         try:

application/parser/chunking_strategies.py

@@ -288,12 +288,9 @@ class SemanticChunker(_BaseStrategyChunker):
         sentences = self._split_sentences(text)
         if len(sentences) < 2:
             raise ValueError("too few sentences for semantic chunking")
-        from application.vectorstore.base import EmbeddingsSingleton
-        from application.core.settings import settings
+        from application.vectorstore.base import get_embeddings
 
-        embeddings = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME
-        ).embed_documents(sentences)
+        embeddings = get_embeddings().embed_documents(sentences)
         breakpoints = self._breakpoints(embeddings)
         groups = self._group(sentences, breakpoints)
         return [g for g in self._enforce_tokens(groups) if g.strip()]

application/parser/file/docling_parser.py

@@ -9,10 +9,11 @@ images (PNG, JPEG, TIFF, BMP, WEBP), WebVTT, and specialized XML formats.
 import importlib.util
 import logging
 import os
+import re
 import tempfile
 import zipfile
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 from application.parser.file.base_parser import BaseParser, DocumentParseError
 from application.utils import truncate_to_line_boundary
@@ -213,6 +214,124 @@ def _parse_markup_bounded(
             pass
 
 
+# Suffixes whose text comes (wholly or partly) from OCR, and are therefore
+# covered by the near-empty-output dropout guard. Everything else docling
+# handles (docx/xlsx/html/vtt/...) has a native text layer and legitimately
+# short files, so the guard would only produce false alarms there.
+_IMAGE_SUFFIXES = frozenset(
+    {".png", ".jpg", ".jpeg", ".tiff", ".tif", ".bmp", ".webp"}
+)
+_OCR_GUARDED_SUFFIXES = _IMAGE_SUFFIXES | {".pdf"}
+
+# Chars-per-page floor below which an OCR parse is treated as a pipeline
+# dropout rather than as the document's real content. Scanned pages that
+# carry any text at all clear this comfortably; the observed failure mode
+# returns literally zero characters per page.
+_DEFAULT_OCR_MIN_CHARS_PER_PAGE = 20
+
+# Docling emits this placeholder for a picture it did not read; it is markup,
+# not document text, so it must not count towards the chars-per-page ratio.
+_IMAGE_PLACEHOLDER_RE = re.compile(r"<!--\s*image\s*-->")
+
+
+def _ocr_min_chars_per_page() -> int:
+    """Chars-per-page floor for the OCR dropout guard; 0 disables it."""
+    from application.core.settings import settings
+
+    try:
+        return int(
+            getattr(
+                settings,
+                "DOCLING_OCR_MIN_CHARS_PER_PAGE",
+                _DEFAULT_OCR_MIN_CHARS_PER_PAGE,
+            )
+        )
+    except (TypeError, ValueError):
+        return _DEFAULT_OCR_MIN_CHARS_PER_PAGE
+
+
+def _text_char_count(content: Optional[str]) -> int:
+    """Count real text characters in exported content, ignoring placeholders."""
+    if not content:
+        return 0
+    return len(_IMAGE_PLACEHOLDER_RE.sub("", content).strip())
+
+
+def _positive_len(value: object) -> int:
+    """``len(value)`` when it is a positive int, 0 when it is neither."""
+    try:
+        length = len(value)
+    except TypeError:
+        return 0
+    return length if isinstance(length, int) and length > 0 else 0
+
+
+def _result_page_count(result: object, file: Path) -> int:
+    """Page count of a conversion result, never below 1.
+
+    Args:
+        result: Docling ``ConversionResult``.
+        file: Path of the file that was converted.
+
+    Returns:
+        Number of pages; 1 for images and whenever docling reports none.
+    """
+    if Path(file).suffix.lower() in _IMAGE_SUFFIXES:
+        return 1
+    document = getattr(result, "document", None)
+    pages = _positive_len(getattr(document, "pages", None))
+    if pages:
+        return pages
+    num_pages = getattr(document, "num_pages", None)
+    if callable(num_pages):
+        try:
+            reported = num_pages()
+        except Exception:  # pragma: no cover - defensive, docling-version drift
+            reported = None
+        if isinstance(reported, int) and reported > 0:
+            return reported
+    return _positive_len(getattr(result, "pages", None)) or 1
+
+
+def _pdf_text_layer_probe(file: Path) -> Tuple[int, int]:
+    """Measure a PDF's embedded text layer with pypdfium2.
+
+    Used only to explain a suspected OCR dropout: a near-empty parse of a PDF
+    that *does* carry a text layer means docling failed to read text it never
+    needed OCR for, while an absent text layer points at the OCR stage itself.
+
+    Args:
+        file: Path to the PDF.
+
+    Returns:
+        Tuple of (page count, characters in the text layer); (0, 0) when the
+        file cannot be probed.
+    """
+    try:
+        import pypdfium2 as pdfium
+    except ImportError:
+        return 0, 0
+    try:
+        pdf = pdfium.PdfDocument(str(file))
+    except Exception:
+        return 0, 0
+    try:
+        total = 0
+        for index in range(len(pdf)):
+            page = pdf[index]
+            textpage = page.get_textpage()
+            try:
+                total += max(0, textpage.count_chars())
+            finally:
+                textpage.close()
+                page.close()
+        return len(pdf), total
+    except Exception:
+        return 0, 0
+    finally:
+        pdf.close()
+
+
 class DoclingParser(BaseParser):
     """Parser using docling for advanced document processing.
 
@@ -291,6 +410,13 @@ class DoclingParser(BaseParser):
             ocr_options = self._get_ocr_options()
             if ocr_options is not None:
                 pipeline_options.ocr_options = ocr_options
+            # Docling's *default* OCR options carry their own flag, so without
+            # this the setting was silently dropped whenever `_get_ocr_options`
+            # returned None (use_rapidocr=False) — including the dropout retry,
+            # whose whole point is forcing full-page OCR.
+            active_ocr_options = getattr(pipeline_options, "ocr_options", None)
+            if hasattr(active_ocr_options, "force_full_page_ocr"):
+                active_ocr_options.force_full_page_ocr = self.force_full_page_ocr
 
         return DocumentConverter(
             format_options={
@@ -383,18 +509,113 @@ class DoclingParser(BaseParser):
 
         return content
 
+    def _ocr_guard_applies(self, file: Path) -> bool:
+        """Whether the near-empty-output dropout guard covers this parse."""
+        return (
+            bool(self.ocr_enabled)
+            and _ocr_min_chars_per_page() > 0
+            and Path(file).suffix.lower() in _OCR_GUARDED_SUFFIXES
+        )
+
+    def _recover_from_ocr_dropout(
+        self, file: Path, first_pass_chars: int, pages: int
+    ) -> str:
+        """Retry a near-empty OCR parse once on a fresh full-page-OCR converter.
+
+        Docling caches its pipeline (and the threaded page queue behind it) on
+        the ``DocumentConverter`` instance, so a converter that has degraded
+        mid-worker keeps returning empty pages while a fresh one starts clean —
+        hence a brand new converter rather than a second `convert` call.
+
+        ``self._converter`` is dropped either way instead of being replaced by
+        the retry converter: the degraded instance must not survive, but the
+        retry one is configured for full-page OCR, and keeping it would impose
+        that cost (and its worse results on text PDFs) on every later file in
+        the worker. The next parse lazily rebuilds one with normal options.
+
+        Args:
+            file: Path to the file being parsed.
+            first_pass_chars: Text characters the first pass produced.
+            pages: Page count reported for the first pass.
+
+        Returns:
+            The retry's content, when it clears the chars-per-page floor.
+
+        Raises:
+            DocumentParseError: If the retry is near-empty too, so the caller
+                fails loudly instead of indexing an empty document.
+        """
+        name = Path(file).name
+        logger.warning(
+            "OCR output for %s is near-empty (%d chars over %d page(s)); "
+            "retrying once on a fresh converter with full-page OCR",
+            name,
+            first_pass_chars,
+            pages,
+        )
+        original_force_full_page_ocr = self.force_full_page_ocr
+        try:
+            self.force_full_page_ocr = True
+            result = self._create_converter().convert(str(file))
+            content = self._export_content(result.document)
+        finally:
+            self.force_full_page_ocr = original_force_full_page_ocr
+            self._converter = None
+
+        retry_chars = _text_char_count(content)
+        retry_pages = max(pages, _result_page_count(result, file))
+        if retry_chars >= _ocr_min_chars_per_page() * retry_pages:
+            logger.warning(
+                "Recovered %s on retry: %d chars over %d page(s) after a "
+                "near-empty first pass (%d chars); the degraded converter has "
+                "been discarded",
+                name,
+                retry_chars,
+                retry_pages,
+                first_pass_chars,
+            )
+            return content
+
+        detail = ""
+        if Path(file).suffix.lower() == ".pdf":
+            layer_pages, layer_chars = _pdf_text_layer_probe(file)
+            if layer_chars > 0:
+                detail = (
+                    f" The PDF carries a {layer_chars}-char text layer over "
+                    f"{layer_pages} page(s) that docling should have read "
+                    "without OCR at all."
+                )
+            elif layer_pages > 0:
+                detail = (
+                    f" The PDF has no text layer over {layer_pages} page(s), "
+                    "so OCR was the only possible source."
+                )
+        raise DocumentParseError(
+            f"OCR produced {retry_chars} chars over {retry_pages} pages for "
+            f"{name}; likely OCR pipeline dropout — document not indexed."
+            + detail
+        )
+
     def parse_file(self, file: Path, errors: str = "ignore") -> Union[str, List[str]]:
         """Parse file using docling with hybrid OCR.
 
         Uses smart OCR approach where the layout model detects text vs bitmap
         regions. Text is extracted directly, bitmaps are OCR'd only when needed.
 
+        When OCR is enabled for a PDF or image, a near-empty result is treated
+        as a pipeline dropout rather than as the document's content: it is
+        retried once on a fresh converter, and failing that it raises.
+
         Args:
             file: Path to the file to parse
             errors: Error handling mode (ignored, docling handles internally)
 
         Returns:
             Parsed document content as markdown string
+
+        Raises:
+            DocumentParseError: If docling fails, or if an OCR parse stays
+                near-empty across both passes.
         """
         logger.info(f"parse_file called for: {file}")
 
@@ -405,10 +626,29 @@ class DoclingParser(BaseParser):
             logger.info(f"Converting file with hybrid OCR: {file}")
             result = self._converter.convert(str(file))
             content = self._export_content(result.document)
-            logger.info(f"Parse complete, content length: {len(content)} chars")
+            pages = _result_page_count(result, file)
+            chars = _text_char_count(content)
+            logger.info(
+                "Parse complete for %s: %d chars (%d text) over %d page(s), "
+                "%.1f chars/page",
+                Path(file).name,
+                len(content),
+                chars,
+                pages,
+                chars / pages,
+            )
+
+            if self._ocr_guard_applies(file) and chars < (
+                _ocr_min_chars_per_page() * pages
+            ):
+                return self._recover_from_ocr_dropout(file, chars, pages)
 
             return content
 
+        except DocumentParseError:
+            # Already the loud, actionable failure — do not re-wrap it into a
+            # generic "Failed to parse ..." and lose the dropout diagnosis.
+            raise
         except Exception as e:
             logger.error(f"Error parsing file with docling: {e}", exc_info=True)
             # ``errors`` governs *decoding* leniency, not whether a total

application/retriever/classic_rag.py

@@ -1,8 +1,17 @@
 import logging
+from typing import Any, Dict, List, Optional, Tuple
 
 from application.core.settings import settings
 from application.llm.llm_creator import LLMCreator
 from application.retriever.base import BaseRetriever
+from application.retriever.fanout import (  # noqa: F401  (re-exported)
+    DEFAULT_MAX_PARALLEL_SOURCES,
+    EMBEDDING_ATTRS as _EMBEDDING_ATTRS,
+    embed_questions,
+    max_parallel_sources,
+    run_source_jobs,
+    store_embeddings,
+)
 from application.retriever.labels import labels_from_metadata
 from application.utils import num_tokens_from_string
 from application.vectorstore.vector_creator import VectorCreator
@@ -10,6 +19,11 @@ from application.vectorstore.vector_creator import VectorCreator
 logger = logging.getLogger(__name__)
 
 
+def _max_parallel_sources(n_sources: int) -> int:
+    """Worker count for the per-source fan-out, bounded by the source count."""
+    return max_parallel_sources(n_sources, settings)
+
+
 class ClassicRAG(BaseRetriever):
     # The group's real top-k, set by the Dispatcher when it inflates ``chunks``
     # for a prescreen fetch. None → ``chunks`` is already the top-k.
@@ -152,13 +166,25 @@ class ClassicRAG(BaseRetriever):
             logger.error(f"Error rephrasing query: {e}", exc_info=True)
             return self.original_question
 
-    def _fetch_candidates(self, docsearch, question, src_k, score_threshold):
+    def _fetch_candidates(
+        self,
+        docsearch,
+        question: str,
+        src_k: int,
+        score_threshold: Optional[float],
+        query_vector: Optional[List[float]] = None,
+    ):
         """Fetch candidate hits for one vector store (vector search).
 
         Returns plain hits, or ``(hit, score)`` pairs when ``include_scores`` is
         set. Subclasses override this to change candidate sourcing (e.g. RRF
         fusion) while inheriting the surrounding per-source resolution and
         budgeting.
+
+        Args:
+            query_vector: Query embedding computed once for the whole
+                retrieval. Forwarded so the store skips embedding the query
+                again; stores that don't support it ignore the kwarg.
         """
         # ``score_threshold`` is honoured by pgvector/mongodb and safely ignored
         # by stores whose ``search`` swallows kwargs. The candidate count is
@@ -167,6 +193,8 @@ class ClassicRAG(BaseRetriever):
         search_kwargs = {"k": k}
         if score_threshold is not None:
             search_kwargs["score_threshold"] = score_threshold
+        if query_vector is not None:
+            search_kwargs["query_vector"] = query_vector
         if self.include_scores:
             return docsearch.search_with_scores(question, **search_kwargs)
         return docsearch.search(question, **search_kwargs)
@@ -175,6 +203,152 @@ class ClassicRAG(BaseRetriever):
         """Label for the scores ``_fetch_candidates`` attaches (None if unscored)."""
         return getattr(docsearch, "score_kind", None)
 
+    def _resolve_source(
+        self, vectorstore_id: str, chunks_per_source: int
+    ) -> Dict[str, Any]:
+        """Resolve one source's fetch parameters (top-k, threshold, query).
+
+        Per-source overrides come from the Dispatcher; absent, the source gets
+        the global behaviour — byte-identical to the pre-override path.
+        """
+        src_cfg = self.per_source_retrieval.get(vectorstore_id)
+        if src_cfg is None:
+            # No per-source override → the effective rephrase_query defaults to
+            # True, so use the (lazily-cached) rephrased question. In the
+            # non-deferred path the cache is already populated.
+            return {
+                "id": vectorstore_id,
+                "src_k": chunks_per_source,
+                "score_threshold": None,
+                "question": self._get_rephrased_question(),
+            }
+        src_k = max(1, int(src_cfg.chunks))
+        # Prescreen fetches a larger candidate set up front; the Dispatcher's
+        # prescreen stage trims back to max_keep afterwards. Raise the fetch
+        # size to candidate_k here.
+        ps_cfg = (
+            src_cfg.prescreen_config() if hasattr(src_cfg, "prescreen_config") else None
+        )
+        if ps_cfg is not None:
+            src_k = max(src_k, int(ps_cfg.candidate_k))
+        return {
+            "id": vectorstore_id,
+            "src_k": src_k,
+            "score_threshold": src_cfg.score_threshold,
+            "question": (
+                self._get_rephrased_question()
+                if src_cfg.rephrase_query
+                else self.original_question
+            ),
+        }
+
+    def _plan_sources(self, chunks_per_source: int) -> List[Dict[str, Any]]:
+        """Resolve every source's fetch parameters, in source order.
+
+        Runs on the calling thread: the lazy rephrase behind it is an LLM
+        side-call that must happen once, not once per worker.
+        """
+        plans = []
+        for vectorstore_id in self.vectorstores:
+            if not vectorstore_id:
+                continue
+            try:
+                plans.append(self._resolve_source(vectorstore_id, chunks_per_source))
+            except Exception as e:
+                logger.error(
+                    f"Error searching vectorstore {vectorstore_id}: {e}", exc_info=True
+                )
+        return plans
+
+    @staticmethod
+    def _store_embeddings(docsearch):
+        """Return the embeddings object a vector store searches with, if any."""
+        return store_embeddings(docsearch)
+
+    def _embed_questions(
+        self, docsearch, questions: List[str]
+    ) -> Dict[str, List[float]]:
+        """Embed each distinct query once for the whole retrieval."""
+        return embed_questions(docsearch, questions)
+
+    def _search_source(
+        self,
+        plan: Dict[str, Any],
+        docsearch=None,
+        query_vector: Optional[List[float]] = None,
+    ) -> Optional[Tuple[Any, Optional[str]]]:
+        """Search one source, returning ``(candidates, score_kind)``.
+
+        Builds the vector store when not supplied, so each worker thread owns
+        its own store instance (and therefore its own DB connection). Errors are
+        logged and reported as ``None`` so one bad source cannot take the rest
+        of the retrieval down with it.
+        """
+        try:
+            if docsearch is None:
+                docsearch = VectorCreator.create_vectorstore(
+                    settings.VECTOR_STORE, plan["id"], settings.EMBEDDINGS_KEY
+                )
+            docs_temp = self._fetch_candidates(
+                docsearch,
+                plan["question"],
+                plan["src_k"],
+                plan["score_threshold"],
+                query_vector=query_vector,
+            )
+            score_kind = self._score_kind(docsearch) if self.include_scores else None
+            return docs_temp, score_kind
+        except Exception as e:
+            logger.error(
+                f"Error searching vectorstore {plan['id']}: {e}", exc_info=True
+            )
+            return None
+
+    def _fetch_all(
+        self, plans: List[Dict[str, Any]]
+    ) -> List[Optional[Tuple[Any, Optional[str]]]]:
+        """Fetch every source's candidates, one embedding and one fan-out.
+
+        The first store is built on the calling thread because its embeddings
+        object supplies the shared query vector — and priming the embeddings
+        singleton there keeps the workers off a concurrent model load. The
+        searches then run on a bounded pool, and results come back in the
+        original source order so the merge stays deterministic.
+        """
+        first_store = None
+        try:
+            first_store = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, plans[0]["id"], settings.EMBEDDINGS_KEY
+            )
+        except Exception as e:
+            logger.error(
+                f"Error searching vectorstore {plans[0]['id']}: {e}", exc_info=True
+            )
+
+        if first_store is None:
+            # The first source is already a logged failure; the rest still run,
+            # each embedding its own query (there is no store to borrow one from).
+            return [None] + self._run_jobs([(plan, None, None) for plan in plans[1:]])
+
+        questions = list(dict.fromkeys(plan["question"] for plan in plans))
+        vectors = self._embed_questions(first_store, questions)
+        return self._run_jobs(
+            [
+                (plan, first_store if idx == 0 else None, vectors.get(plan["question"]))
+                for idx, plan in enumerate(plans)
+            ]
+        )
+
+    def _run_jobs(self, jobs: List[Tuple]) -> List[Optional[Tuple[Any, Optional[str]]]]:
+        """Run the per-source search jobs, concurrently when there are several."""
+        if not jobs:
+            return []
+        return run_source_jobs(
+            lambda job: self._search_source(*job),
+            jobs,
+            workers=_max_parallel_sources(len(jobs)),
+        )
+
     def _get_data(self):
         if self.chunks == 0 or not self.vectorstores:
             logger.info(
@@ -194,96 +368,64 @@ class ClassicRAG(BaseRetriever):
         token_budget = max(int(self.doc_token_limit * 0.9), 100)
         cumulative_tokens = 0
 
-        for vectorstore_id in self.vectorstores:
-            if vectorstore_id:
-                try:
-                    # Per-source overrides (set by the Dispatcher). Absent →
-                    # global behaviour, byte-identical to before.
-                    src_cfg = self.per_source_retrieval.get(vectorstore_id)
-                    if src_cfg is not None:
-                        src_k = max(1, int(src_cfg.chunks))
-                        # Prescreen fetches a larger candidate set up front; the
-                        # Dispatcher's prescreen stage trims back to max_keep
-                        # afterwards. Raise the fetch size to candidate_k here.
-                        ps_cfg = (
-                            src_cfg.prescreen_config()
-                            if hasattr(src_cfg, "prescreen_config")
-                            else None
-                        )
-                        if ps_cfg is not None:
-                            src_k = max(src_k, int(ps_cfg.candidate_k))
-                        score_threshold = src_cfg.score_threshold
-                        question = (
-                            self._get_rephrased_question()
-                            if src_cfg.rephrase_query
-                            else self.original_question
-                        )
-                    else:
-                        src_k = chunks_per_source
-                        score_threshold = None
-                        # No per-source override → the effective rephrase_query
-                        # defaults to True, so use the (lazily-cached) rephrased
-                        # question. In the non-deferred path the cache is already
-                        # populated, so this matches today's behaviour exactly.
-                        question = self._get_rephrased_question()
+        # Resolve every source, then fetch them all (one query embedding, one
+        # bounded fan-out). The merge below stays serial and in source order, so
+        # dedupe/budget/trim semantics are exactly what they were.
+        plans = self._plan_sources(chunks_per_source)
+        results = self._fetch_all(plans) if plans else []
 
-                    docsearch = VectorCreator.create_vectorstore(
-                        settings.VECTOR_STORE, vectorstore_id, settings.EMBEDDINGS_KEY
-                    )
-                    docs_temp = self._fetch_candidates(
-                        docsearch, question, src_k, score_threshold
-                    )
-                    score_kind = (
-                        self._score_kind(docsearch) if self.include_scores else None
-                    )
+        for plan, result in zip(plans, results):
+            if result is None:
+                continue
+            vectorstore_id = plan["id"]
+            src_k = plan["src_k"]
+            docs_temp, score_kind = result
+            try:
+                # ``_fetch_candidates`` over-fetches (k >= 20) so a prescreen
+                # stage has candidates to filter; trim back to src_k so
+                # ``chunks`` is the final top-k it claims to be. With
+                # prescreen on, src_k is already raised to candidate_k above,
+                # so the stage still sees its full candidate set.
+                kept = 0
 
-                    # ``_fetch_candidates`` over-fetches (k >= 20) so a prescreen
-                    # stage has candidates to filter; trim back to src_k so
-                    # ``chunks`` is the final top-k it claims to be. With
-                    # prescreen on, src_k is already raised to candidate_k above,
-                    # so the stage still sees its full candidate set.
-                    kept = 0
-
-                    for doc in docs_temp:
-                        if kept >= src_k or cumulative_tokens >= token_budget:
-                            break
-
-                        score = None
-                        if isinstance(doc, tuple):
-                            doc, score = doc
-
-                        if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
-                            page_content = doc.page_content
-                            metadata = doc.metadata
-                        else:
-                            page_content = doc.get("text", doc.get("page_content", ""))
-                            metadata = doc.get("metadata", {})
-
-                        labels = labels_from_metadata(
-                            metadata, page_content, vectorstore_id
-                        )
-
-                        doc_text_with_header = f"{labels['filename']}\n{page_content}"
-                        doc_tokens = num_tokens_from_string(doc_text_with_header)
-
-                        if cumulative_tokens + doc_tokens < token_budget:
-                            entry = {"text": page_content, **labels}
-                            if self.include_scores:
-                                entry["score"] = score
-                                entry["score_kind"] = score_kind
-                            all_docs.append(entry)
-                            cumulative_tokens += doc_tokens
-                            kept += 1
-
-                    if cumulative_tokens >= token_budget:
+                for doc in docs_temp:
+                    if kept >= src_k or cumulative_tokens >= token_budget:
                         break
 
-                except Exception as e:
-                    logger.error(
-                        f"Error searching vectorstore {vectorstore_id}: {e}",
-                        exc_info=True,
-                    )
-                    continue
+                    score = None
+                    if isinstance(doc, tuple):
+                        doc, score = doc
+
+                    if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
+                        page_content = doc.page_content
+                        metadata = doc.metadata
+                    else:
+                        page_content = doc.get("text", doc.get("page_content", ""))
+                        metadata = doc.get("metadata", {})
+
+                    labels = labels_from_metadata(metadata, page_content, vectorstore_id)
+
+                    doc_text_with_header = f"{labels['filename']}\n{page_content}"
+                    doc_tokens = num_tokens_from_string(doc_text_with_header)
+
+                    if cumulative_tokens + doc_tokens < token_budget:
+                        entry = {"text": page_content, **labels}
+                        if self.include_scores:
+                            entry["score"] = score
+                            entry["score_kind"] = score_kind
+                        all_docs.append(entry)
+                        cumulative_tokens += doc_tokens
+                        kept += 1
+
+                if cumulative_tokens >= token_budget:
+                    break
+
+            except Exception as e:
+                logger.error(
+                    f"Error searching vectorstore {vectorstore_id}: {e}",
+                    exc_info=True,
+                )
+                continue
 
         # ``chunks_per_source`` has a floor of 1 so no attached source is
         # starved, which means N sources always yield at least N documents —

application/retriever/fanout.py

@@ -0,0 +1,127 @@
+"""Shared helpers for fanning one query out across several vector stores.
+
+Attaching N sources used to cost N query embeddings and N serial round trips.
+The pieces here collapse that to one embedding plus one bounded, order-
+preserving thread pool, and are reused by :mod:`application.retriever.classic_rag`
+and :mod:`application.services.search_service` so both paths behave the same.
+"""
+
+from __future__ import annotations
+
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar
+
+from application.core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+# Fallback fan-out width when ``RETRIEVAL_MAX_PARALLEL_SOURCES`` is unset.
+DEFAULT_MAX_PARALLEL_SOURCES = 4
+
+# Attributes the vector stores keep their embeddings object on, in the order
+# they are probed. Stores that build embeddings inline (elasticsearch, lancedb)
+# expose none of them and are handled by the ``_get_embeddings`` fallback.
+EMBEDDING_ATTRS = ("_embedding", "_embeddings", "embeddings")
+
+Job = TypeVar("Job")
+Result = TypeVar("Result")
+
+
+def max_parallel_sources(n_sources: int, settings_obj: Any = None) -> int:
+    """Worker count for the per-source fan-out, bounded by the source count.
+
+    Args:
+        n_sources: Number of sources about to be searched.
+        settings_obj: Settings to read ``RETRIEVAL_MAX_PARALLEL_SOURCES`` from;
+            defaults to the global settings.
+
+    Returns:
+        int: Worker count, at least 1 and never above ``n_sources``.
+    """
+    config = settings if settings_obj is None else settings_obj
+    configured = getattr(
+        config, "RETRIEVAL_MAX_PARALLEL_SOURCES", DEFAULT_MAX_PARALLEL_SOURCES
+    )
+    try:
+        configured = int(configured)
+    except (TypeError, ValueError):
+        configured = DEFAULT_MAX_PARALLEL_SOURCES
+    return max(1, min(configured, n_sources))
+
+
+def store_embeddings(docsearch) -> Optional[Any]:
+    """Return the embeddings object a vector store searches with, if any."""
+    for attr in EMBEDDING_ATTRS:
+        embedder = getattr(docsearch, attr, None)
+        if embedder is not None and hasattr(embedder, "embed_query"):
+            return embedder
+    getter = getattr(docsearch, "_get_embeddings", None)
+    if callable(getter):
+        try:
+            return getter(settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY)
+        except Exception as e:
+            logger.debug("Could not resolve store embeddings: %s", e)
+    return None
+
+
+def embed_questions(docsearch, questions: Iterable[str]) -> Dict[str, List[float]]:
+    """Embed each distinct query once for the whole retrieval.
+
+    Every source of one retrieval shares a single embeddings config
+    (``settings.EMBEDDINGS_NAME`` is global — nothing per-source overrides it),
+    so a vector computed with one store's embedder is valid for all of them.
+    That turns N sources from N query embeddings into one.
+
+    Args:
+        docsearch: Any store of the retrieval; only its embedder is used.
+        questions: Queries to embed (duplicates are embedded once).
+
+    Returns:
+        dict: ``{question: vector}``, or ``{}`` when no embedder is reachable
+        or embedding fails — in which case every store embeds its own query,
+        exactly as before.
+    """
+    embedder = store_embeddings(docsearch)
+    if embedder is None:
+        return {}
+    try:
+        return {
+            question: embedder.embed_query(question)
+            for question in dict.fromkeys(questions)
+        }
+    except Exception as e:
+        logger.warning(
+            "Query embedding failed (%s); each store will embed its own.", e
+        )
+        return {}
+
+
+def run_source_jobs(
+    fn: Callable[[Job], Result],
+    jobs: Iterable[Job],
+    workers: Optional[int] = None,
+) -> List[Result]:
+    """Run one job per source, concurrently when there are several.
+
+    Args:
+        fn: Per-source callable; it must swallow its own errors, since a raise
+            here aborts the whole fan-out.
+        jobs: Job payloads, in source order.
+        workers: Pool size; defaults to :func:`max_parallel_sources`.
+
+    Returns:
+        list: ``fn`` results in the original job order, so the caller's merge
+        stays deterministic no matter which source finishes first.
+    """
+    jobs = list(jobs)
+    if not jobs:
+        return []
+    if workers is None:
+        workers = max_parallel_sources(len(jobs))
+    if workers == 1:
+        return [fn(job) for job in jobs]
+    with ThreadPoolExecutor(
+        max_workers=workers, thread_name_prefix="rag-source"
+    ) as pool:
+        return list(pool.map(fn, jobs))

application/retriever/graph_rag.py

@@ -7,7 +7,11 @@ reused) rephrase.
 
 Composes :class:`ClassicRAG` rather than subclassing: PPR doesn't fit the
 ``_fetch_candidates`` hook, but the composed instance supplies the rephrase, the
-token-budget loop, and the per-source fallback when a source has no graph.
+token-budget loop, and the fallback for sources that have no graph.
+
+Per request the whole source group costs one node-count query, one query
+embedding, and one ClassicRAG run for all the graphless sources together — all
+over a single pooled connection shared with the vector store.
 """
 
 from __future__ import annotations
@@ -25,7 +29,7 @@ from application.retriever.base import BaseRetriever
 from application.retriever.classic_rag import ClassicRAG
 from application.retriever.labels import labels_from_metadata
 from application.utils import num_tokens_from_string
-from application.vectorstore.base import EmbeddingsSingleton
+from application.vectorstore.base import get_embeddings
 
 SEED_NODES = 10
 SUBGRAPH_HOPS = 1
@@ -88,9 +92,7 @@ class GraphRAGRetriever(BaseRetriever):
         self.per_source_retrieval = {}
 
     def _embed_query(self, question: str) -> List[float]:
-        embedding = EmbeddingsSingleton.get_instance(
-            settings.EMBEDDINGS_NAME, settings.EMBEDDINGS_KEY
-        )
+        embedding = get_embeddings()
         return embedding.embed_query(question)
 
     def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]:
@@ -157,10 +159,17 @@ class GraphRAGRetriever(BaseRetriever):
         base = self.base_chunks if self.base_chunks is not None else self.chunks
         return max(1, base // max(1, len(self.vectorstores)))
 
-    def _graph_docs_for_source(self, store, source_id) -> List[Dict[str, Any]]:
-        """Local PPR retrieval for one source (caller guarantees it has a graph)."""
-        question = self._classic._get_rephrased_question()
-        query_embedding = self._embed_query(question)
+    def _graph_docs_for_source(
+        self, store, source_id, query_embedding: List[float]
+    ) -> List[Dict[str, Any]]:
+        """Local PPR retrieval for one source (caller guarantees it has a graph).
+
+        Args:
+            store: Open :class:`GraphStore` shared by every source of this run.
+            source_id: Source to retrieve from.
+            query_embedding: Embedding of the rephrased question, computed once
+                by the caller for the whole retrieval.
+        """
         seed_rows = store.search_nodes_by_embedding(
             source_id, query_embedding, k=SEED_NODES
         )
@@ -204,15 +213,26 @@ class GraphRAGRetriever(BaseRetriever):
             cumulative_tokens += doc_tokens
         return docs
 
-    def _classic_for_source(self, source_id) -> List[Dict[str, Any]]:
-        """Reuse the composed ClassicRAG to retrieve one source's chunks."""
+    def _classic_for_sources(self, source_ids) -> List[Dict[str, Any]]:
+        """Reuse the composed ClassicRAG to retrieve a whole batch of sources.
+
+        One inner run for every graphless source instead of one run per source:
+        ClassicRAG then embeds the query once and fans the sources out itself.
+        The trade is that the per-source chunk split and the shared doc-token
+        budget apply across the batch — i.e. exactly plain ClassicRAG semantics
+        over those sources, rather than each source getting its own full budget.
+        """
+        source_ids = [source_id for source_id in source_ids if source_id]
+        if not source_ids:
+            return []
+        wanted = set(source_ids)
         original = self._classic.vectorstores
         original_overrides = self._classic.per_source_retrieval
         original_base = self._classic.base_chunks
         try:
-            self._classic.vectorstores = [source_id]
+            self._classic.vectorstores = list(source_ids)
             self._classic.per_source_retrieval = {
-                k: v for k, v in self.per_source_retrieval.items() if k == source_id
+                k: v for k, v in self.per_source_retrieval.items() if k in wanted
             }
             # The Dispatcher sets these on *this* object; the inner retriever is
             # the one that reads them.
@@ -223,8 +243,79 @@ class GraphRAGRetriever(BaseRetriever):
             self._classic.per_source_retrieval = original_overrides
             self._classic.base_chunks = original_base
 
+    def _classic_for_source(self, source_id) -> List[Dict[str, Any]]:
+        """Retrieve one source through the batched classic path."""
+        return self._classic_for_sources([source_id])
+
+    def _retrieve_with_store(self, store, sources) -> List[Dict[str, Any]]:
+        """Split ``sources`` by graph presence, then batch each half.
+
+        Graph sources keep their own slot in source order; every graphless
+        source collapses into a single ClassicRAG run that occupies the slot of
+        the first graphless source. Sources whose PPR retrieval raises are
+        collected and retried as one more classic batch, appended at the end.
+        """
+        try:
+            counts = store.count_nodes_many(sources)
+        except Exception as e:
+            logging.error(f"GraphRAG count_nodes failed for {sources}: {e}")
+            counts = {}
+
+        segments: List[List[Dict[str, Any]]] = []
+        graph_slots: Dict[str, int] = {}
+        graphed: List[str] = []
+        graphless: List[str] = []
+        classic_slot = None
+        for source_id in sources:
+            if counts.get(source_id, 0) > 0:
+                if source_id in graph_slots:
+                    continue
+                graph_slots[source_id] = len(segments)
+                segments.append([])
+                graphed.append(source_id)
+            else:
+                if classic_slot is None:
+                    classic_slot = len(segments)
+                    segments.append([])
+                graphless.append(source_id)
+
+        failed: List[str] = []
+        query_embedding = None
+        if graphed:
+            # Embedded once for the whole retrieval, not once per graph source.
+            try:
+                query_embedding = self._embed_query(
+                    self._classic._get_rephrased_question()
+                )
+            except Exception as e:
+                logging.error(
+                    f"GraphRAG query embedding failed, falling back: {e}",
+                    exc_info=True,
+                )
+                failed, graphed = list(graphed), []
+
+        for source_id in graphed:
+            try:
+                segments[graph_slots[source_id]] = self._graph_docs_for_source(
+                    store, source_id, query_embedding
+                )
+            except Exception as e:
+                logging.error(
+                    f"GraphRAG retrieval failed for {source_id}, falling back: {e}",
+                    exc_info=True,
+                )
+                failed.append(source_id)
+
+        if graphless:
+            segments[classic_slot] = self._classic_for_sources(graphless)
+        if failed:
+            segments.append(self._classic_for_sources(failed))
+
+        return [doc for segment in segments for doc in segment]
+
     def _get_data(self) -> List[Dict[str, Any]]:
-        if not self.vectorstores:
+        sources = [source_id for source_id in self.vectorstores if source_id]
+        if not sources:
             return []
 
         store = None
@@ -235,31 +326,17 @@ class GraphRAGRetriever(BaseRetriever):
                 logging.error(f"GraphRAG store unavailable, falling back: {e}")
                 store = None
 
-        all_docs: List[Dict[str, Any]] = []
-        for source_id in self.vectorstores:
-            if not source_id:
-                continue
-            has_graph = False
-            if store is not None:
-                try:
-                    has_graph = store.count_nodes(source_id) > 0
-                except Exception as e:
-                    logging.error(f"GraphRAG count_nodes failed for {source_id}: {e}")
-                    has_graph = False
-
-            if not has_graph:
-                all_docs.extend(self._classic_for_source(source_id))
-                continue
+        if store is None:
+            return self._classic_for_sources(sources)
 
+        try:
+            return self._retrieve_with_store(store, sources)
+        finally:
+            # Hand the pooled connection back; the store is per-request.
             try:
-                all_docs.extend(self._graph_docs_for_source(store, source_id))
+                store.close()
             except Exception as e:
-                logging.error(
-                    f"GraphRAG retrieval failed for {source_id}, falling back: {e}",
-                    exc_info=True,
-                )
-                all_docs.extend(self._classic_for_source(source_id))
-        return all_docs
+                logging.debug("Error closing GraphRAG store: %s", e)
 
     def search(self, query: str = "") -> List[Dict[str, Any]]:
         if query:

application/retriever/hybrid_rag.py

@@ -55,16 +55,22 @@ class HybridRetriever(ClassicRAG):
         carry their own label."""
         return "rrf"
 
-    def _fetch_candidates(self, docsearch, question, src_k, score_threshold):
+    def _fetch_candidates(
+        self, docsearch, question, src_k, score_threshold, query_vector=None
+    ):
         """Return RRF-fused vector+keyword hits for one vector store.
 
         Inherits the per-source resolution and budgeting from
         :meth:`ClassicRAG._get_data`; only candidate sourcing differs.
         RRF scores are not cosine similarities, so ``score_threshold`` is
-        intentionally not applied to the fused list.
+        intentionally not applied to the fused list. ``query_vector`` is the
+        retrieval's single query embedding — the keyword half never needs one.
         """
         candidate_k = min(max(src_k * 2, 20), 500)
-        vector_hits = docsearch.search(question, k=candidate_k)
+        vector_kwargs = {"k": candidate_k}
+        if query_vector is not None:
+            vector_kwargs["query_vector"] = query_vector
+        vector_hits = docsearch.search(question, **vector_kwargs)
         keyword_hits = docsearch.keyword_search(question, k=candidate_k)
         fused = fuse_with_scores(vector_hits, keyword_hits)
         if self.include_scores:

application/services/search_service.py

@@ -8,9 +8,10 @@ MCP error responses, etc.).
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from application.core.settings import settings
+from application.retriever.fanout import embed_questions, run_source_jobs
 from application.storage.db.repositories.agents import AgentsRepository
 from application.storage.db.session import db_readonly
 from application.vectorstore.vector_creator import VectorCreator
@@ -89,6 +90,79 @@ def _authorized_source_ids(conn, agent: Dict[str, Any], source_ids: List[str]) -
     return allowed
 
 
+def _search_one(
+    source_id: str,
+    docsearch: Any,
+    query: str,
+    k: int,
+    query_vector: Optional[List[float]],
+) -> Optional[List[Any]]:
+    """Search one source, returning its hits or ``None`` when it fails.
+
+    Builds the vector store when not supplied, so each worker thread owns its
+    own store instance (and therefore its own DB connection). Errors are logged
+    and reported as ``None`` so one broken index cannot take the rest down.
+    """
+    try:
+        if docsearch is None:
+            docsearch = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+            )
+        search_kwargs: Dict[str, Any] = {"k": k}
+        if query_vector is not None:
+            search_kwargs["query_vector"] = query_vector
+        return docsearch.search(query, **search_kwargs)
+    except Exception as e:
+        logger.error(
+            f"Error searching vectorstore {source_id}: {e}",
+            exc_info=True,
+        )
+        return None
+
+
+def _fetch_sources(
+    query: str, source_ids: List[str], k: int
+) -> List[Optional[List[Any]]]:
+    """Fetch every source's hits: one query embedding, one bounded fan-out.
+
+    The first store is built on the calling thread because its embeddings
+    object supplies the shared query vector — and priming the embeddings
+    singleton there keeps the workers off a concurrent model load. Results come
+    back in the original source order so the merge stays deterministic.
+    """
+
+    def _job(job) -> Optional[List[Any]]:
+        source_id, docsearch, query_vector = job
+        return _search_one(source_id, docsearch, query, k, query_vector)
+
+    first_store = None
+    try:
+        first_store = VectorCreator.create_vectorstore(
+            settings.VECTOR_STORE, source_ids[0], settings.EMBEDDINGS_KEY
+        )
+    except Exception as e:
+        logger.error(
+            f"Error searching vectorstore {source_ids[0]}: {e}",
+            exc_info=True,
+        )
+
+    if first_store is None:
+        # The first source is already a logged failure; the rest still run,
+        # each embedding its own query (there is no store to borrow one from).
+        return [None] + run_source_jobs(
+            _job, [(sid, None, None) for sid in source_ids[1:]]
+        )
+
+    vector = embed_questions(first_store, [query]).get(query)
+    return run_source_jobs(
+        _job,
+        [
+            (sid, first_store if idx == 0 else None, vector)
+            for idx, sid in enumerate(source_ids)
+        ],
+    )
+
+
 def _search_sources(
     query: str, source_ids: List[str], chunks: int
 ) -> List[Dict[str, Any]]:
@@ -101,19 +175,24 @@ def _search_sources(
         return []
 
     results: List[Dict[str, Any]] = []
+    # Blank ids build no store but still count here, so the per-source budget
+    # matches what the serial implementation handed each real source.
     chunks_per_source = max(1, chunks // len(source_ids))
     seen_texts: set[int] = set()
 
-    for source_id in source_ids:
-        if not source_id or not source_id.strip():
+    active_ids = [sid for sid in source_ids if sid and sid.strip()]
+    if not active_ids:
+        return []
+
+    # Fetch every source up front, then merge serially in source order so the
+    # dedupe / cap semantics are exactly what they were.
+    fetched = _fetch_sources(query, active_ids, chunks_per_source * 2)
+
+    for source_id, docs in zip(active_ids, fetched):
+        if docs is None:
             continue
 
         try:
-            docsearch = VectorCreator.create_vectorstore(
-                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
-            )
-            docs = docsearch.search(query, k=chunks_per_source * 2)
-
             for doc in docs:
                 if len(results) >= chunks:
                     break

application/storage/db/bootstrap.py

@@ -7,9 +7,15 @@ On app startup the Flask factory (and Celery worker init) can call
    configured role to have ``CREATEDB`` privilege).
 2. Apply every pending Alembic migration up to ``head``.
 
-Both steps are gated by settings that default ON for dev convenience and
-can be turned off in prod (``AUTO_CREATE_DB`` / ``AUTO_MIGRATE``) where
-schema is managed out-of-band by a deploy pipeline.
+:func:`ensure_vector_schema` does the same job for the *vector* database
+(``pgvector``): it owns the ``documents`` table and, when GraphRAG is on,
+the graph tables, so the retrieval hot path never runs DDL. That database
+may be a separate cluster, which is why it has no Alembic migration.
+
+Every step is gated by a setting that defaults ON for dev convenience and
+can be turned off in prod (``AUTO_CREATE_DB`` / ``AUTO_MIGRATE`` /
+``AUTO_VECTOR_SCHEMA``) where schema is managed out-of-band by a deploy
+pipeline.
 
 All heavy imports (alembic, psycopg, sqlalchemy.exc sub-symbols) are
 deferred to inside the function so merely importing this module has no
@@ -19,6 +25,7 @@ side effects and is cheap for test collection.
 from __future__ import annotations
 
 import logging
+import time
 from typing import Optional
 
 
@@ -69,6 +76,134 @@ def ensure_database_ready(
         _run_migrations(log)
 
 
+def ensure_vector_schema(*, logger: Optional[logging.Logger] = None) -> None:
+    """Create the pgvector schema once at boot and verify its dimension.
+
+    Owning the DDL here is what lets ``PGVectorStore`` construction stay free:
+    the retriever builds one store per source per request, and each used to run
+    ``CREATE EXTENSION`` / ``CREATE TABLE`` / two ``CREATE INDEX`` statements
+    before its first SELECT. A dimension mismatch between the table and the
+    configured embeddings model is fatal here rather than silent garbage
+    retrieval at query time.
+
+    Args:
+        logger: Optional logger. Defaults to this module's logger.
+
+    Raises:
+        RuntimeError: The existing ``documents`` table was built for a
+            different embedding width than the configured model produces.
+        Exception: Any connection or DDL failure is re-raised so the app fails
+            fast, matching :func:`ensure_database_ready`.
+    """
+    log = logger or logging.getLogger(__name__)
+
+    from application.core.settings import settings
+
+    store_kind = (settings.VECTOR_STORE or "").lower()
+    if store_kind != "pgvector":
+        log.debug(
+            "ensure_vector_schema: VECTOR_STORE is %r, not pgvector; nothing to do.",
+            settings.VECTOR_STORE,
+        )
+        return
+
+    dsn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
+    if not dsn and getattr(settings, "POSTGRES_URI", None):
+        from application.core.db_uri import normalize_pgvector_connection_string
+
+        dsn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
+    if not dsn:
+        log.info(
+            "ensure_vector_schema: no pgvector connection string configured "
+            "(PGVECTOR_CONNECTION_STRING / POSTGRES_URI); skipping."
+        )
+        return
+
+    import psycopg
+
+    from application.vectorstore.pgvector import (
+        DEFAULT_EMBEDDING_DIM,
+        SCHEMA_LOCK_KEY,
+        PGVectorStore,
+    )
+
+    # Loading the model here is deliberate: the process loads it on first
+    # retrieval anyway, and EmbeddingsSingleton caches it.
+    dim: Optional[int] = None
+    try:
+        from application.vectorstore.base import get_embeddings
+
+        dim = getattr(get_embeddings(), "dimension", None)
+    except Exception as exc:  # noqa: BLE001 — never block boot on the model
+        log.warning(
+            "ensure_vector_schema: could not load the embeddings model (%s); "
+            "creating the table with %d dimensions and skipping the dimension "
+            "check.",
+            exc,
+            DEFAULT_EMBEDDING_DIM,
+        )
+    if dim is None:
+        log.warning(
+            "ensure_vector_schema: the embeddings model exposes no dimension; "
+            "using %d and skipping the dimension check.",
+            DEFAULT_EMBEDDING_DIM,
+        )
+
+    graph_enabled = bool(getattr(settings, "GRAPHRAG_ENABLED", False))
+    started = time.monotonic()
+    # A plain connection, never the store's pool: this can run pre-fork under
+    # ``gunicorn --preload``, and an inherited pooled socket is a broken one.
+    conn = psycopg.connect(dsn)
+    try:
+        cursor = conn.cursor()
+        try:
+            # Serialize concurrent workers; released when this transaction ends.
+            cursor.execute(
+                "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+            )
+        finally:
+            cursor.close()
+
+        PGVectorStore.create_schema(conn, dimension=dim or DEFAULT_EMBEDDING_DIM)
+        if graph_enabled:
+            from application.graphrag.store import (
+                DEFAULT_NAME_EMBEDDING_DIM,
+                GraphStore,
+            )
+
+            GraphStore.create_schema(
+                conn, dimension=dim or DEFAULT_NAME_EMBEDDING_DIM
+            )
+        conn.commit()
+
+        actual = PGVectorStore.table_dimension(conn)
+        if dim and actual and actual != dim:
+            raise RuntimeError(
+                f"pgvector table 'documents' has vector({actual}) but "
+                f"EMBEDDINGS_NAME={settings.EMBEDDINGS_NAME} produces {dim}-dim "
+                "vectors; re-ingest into a fresh table or point "
+                "PGVECTOR_CONNECTION_STRING at the matching database."
+            )
+
+        log.info(
+            "ensure_vector_schema: table 'documents' ready "
+            "(dimension=%s, graph tables=%s) in %d ms.",
+            actual or dim or DEFAULT_EMBEDDING_DIM,
+            "yes" if graph_enabled else "no",
+            int((time.monotonic() - started) * 1000),
+        )
+    except Exception as exc:
+        try:
+            conn.rollback()
+        except Exception:
+            # Connection already gone; nothing left to roll back.
+            pass
+        log.error("ensure_vector_schema: %s", exc)
+        raise
+    finally:
+        conn.close()
+
+
 def _ensure_database_exists(uri: str, log: logging.Logger) -> None:
     """Create the target database if a connection reveals it's missing.
 

application/vectorstore/base.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import requests
 
@@ -196,18 +197,111 @@ class EmbeddingsSingleton:
         }
 
         if embeddings_name in embeddings_factory:
-            return embeddings_factory[embeddings_name](*args, **kwargs)
+            if args or kwargs:
+                logging.debug(
+                    "Dropping %d positional and %d keyword argument(s) for pinned "
+                    "embeddings model %s: its factory takes none.",
+                    len(args),
+                    len(kwargs),
+                    embeddings_name,
+                )
+            return embeddings_factory[embeddings_name]()
         else:
             return EmbeddingsWrapper(embeddings_name, *args, **kwargs)
 
 
+def _azure_configured() -> bool:
+    """True when the Azure OpenAI deployment settings are all present."""
+    return bool(
+        settings.OPENAI_API_BASE
+        and settings.OPENAI_API_VERSION
+        and settings.AZURE_DEPLOYMENT_NAME
+    )
+
+
+def get_embeddings(
+    embeddings_name: Optional[str] = None, embeddings_key: Optional[str] = None
+):
+    """Resolve the configured embeddings instance. The single entry point.
+
+    Callers that reach for :meth:`EmbeddingsSingleton.get_instance` directly
+    reproduce neither the bundled local-model path (the Docker image ships
+    ``/app/models/all-mpnet-base-v2``) nor the OpenAI/Azure key handling: they
+    download a second copy of the model from the hub, and passing the key
+    positionally to the pinned HuggingFace names raises ``TypeError`` because
+    those factories take no arguments. Route every caller through here.
+
+    Args:
+        embeddings_name: Model name; defaults to ``settings.EMBEDDINGS_NAME``.
+        embeddings_key: API key; defaults to ``settings.EMBEDDINGS_KEY``.
+
+    Returns:
+        The shared embeddings instance for the resolved model.
+    """
+    embeddings_name = embeddings_name or settings.EMBEDDINGS_NAME
+    embeddings_key = (
+        embeddings_key if embeddings_key is not None else settings.EMBEDDINGS_KEY
+    )
+
+    # Check for remote embeddings first
+    if settings.EMBEDDINGS_BASE_URL:
+        logging.info(
+            f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}"
+        )
+        return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
+
+    if embeddings_name == "openai_text-embedding-ada-002":
+        if _azure_configured():
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
+            )
+        else:
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name, openai_api_key=embeddings_key
+            )
+    elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
+        possible_paths = [
+            "/app/models/all-mpnet-base-v2",  # Docker absolute path
+            "./models/all-mpnet-base-v2",  # Relative path
+        ]
+        local_model_path = None
+        for path in possible_paths:
+            if os.path.exists(path):
+                local_model_path = path
+                logging.info(f"Found local model at path: {path}")
+                break
+            else:
+                logging.info(f"Path does not exist: {path}")
+        if local_model_path:
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                local_model_path,
+            )
+        else:
+            logging.warning(
+                f"Local model not found in any of the paths: {possible_paths}. Falling back to HuggingFace download."
+            )
+            embedding_instance = EmbeddingsSingleton.get_instance(
+                embeddings_name,
+            )
+    else:
+        embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
+    return embedding_instance
+
+
 class BaseVectorStore(ABC):
     def __init__(self):
         pass
 
     @abstractmethod
     def search(self, *args, **kwargs):
-        """Search for similar documents/chunks in the vectorstore"""
+        """Search for similar documents/chunks in the vectorstore.
+
+        Implementations accept an optional ``query_vector`` kwarg: the query
+        already embedded by the caller, so a multi-source retrieval embeds once
+        instead of once per store. A store that cannot use it must still swallow
+        the kwarg (every signature here ends in ``**kwargs``) and embed the
+        question itself.
+        """
         pass
 
     def keyword_search(self, question, k=10):
@@ -279,53 +373,9 @@ class BaseVectorStore(ABC):
         return deleted
 
     def is_azure_configured(self):
-        return (
-            settings.OPENAI_API_BASE
-            and settings.OPENAI_API_VERSION
-            and settings.AZURE_DEPLOYMENT_NAME
-        )
+        """Kept for compatibility; delegates to the module-level check."""
+        return _azure_configured()
 
     def _get_embeddings(self, embeddings_name, embeddings_key=None):
-        # Check for remote embeddings first
-        if settings.EMBEDDINGS_BASE_URL:
-            logging.info(
-                f"Using remote embeddings API at: {settings.EMBEDDINGS_BASE_URL}"
-            )
-            return EmbeddingsSingleton._remote_instance(embeddings_name, embeddings_key)
-
-        if embeddings_name == "openai_text-embedding-ada-002":
-            if self.is_azure_configured():
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name, model=settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME
-                )
-            else:
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name, openai_api_key=embeddings_key
-                )
-        elif embeddings_name == "huggingface_sentence-transformers/all-mpnet-base-v2":
-            possible_paths = [
-                "/app/models/all-mpnet-base-v2",  # Docker absolute path
-                "./models/all-mpnet-base-v2",  # Relative path
-            ]
-            local_model_path = None
-            for path in possible_paths:
-                if os.path.exists(path):
-                    local_model_path = path
-                    logging.info(f"Found local model at path: {path}")
-                    break
-                else:
-                    logging.info(f"Path does not exist: {path}")
-            if local_model_path:
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    local_model_path,
-                )
-            else:
-                logging.warning(
-                    f"Local model not found in any of the paths: {possible_paths}. Falling back to HuggingFace download."
-                )
-                embedding_instance = EmbeddingsSingleton.get_instance(
-                    embeddings_name,
-                )
-        else:
-            embedding_instance = EmbeddingsSingleton.get_instance(embeddings_name)
-        return embedding_instance
+        """Resolve embeddings for this store; see :func:`get_embeddings`."""
+        return get_embeddings(embeddings_name, embeddings_key)

application/vectorstore/elasticsearch.py

@@ -73,9 +73,27 @@ class ElasticsearchStore(BaseVectorStore):
 
         return es_client
 
-    def search(self, question, k=2, index_name=settings.ELASTIC_INDEX, *args, **kwargs):
-        embeddings = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key)
-        vector = embeddings.embed_query(question)
+    def search(
+        self,
+        question,
+        k=2,
+        index_name=settings.ELASTIC_INDEX,
+        *args,
+        query_vector=None,
+        **kwargs,
+    ):
+        """Search by kNN + full text, fused with RRF.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        vector = query_vector
+        if vector is None:
+            embeddings = self._get_embeddings(
+                settings.EMBEDDINGS_NAME, self.embeddings_key
+            )
+            vector = embeddings.embed_query(question)
         knn = {
             "filter": [{"match": {"metadata.source_id.keyword": self.source_id}}],
             "field": "vector",

application/vectorstore/faiss.py

@@ -174,16 +174,28 @@ class FaissStore(BaseVectorStore):
         return [doc for doc, _ in self.search_with_scores(question, k, *args, **kwargs)]
 
     def search_with_scores(
-        self, question: str, k: int = 4, *args, **kwargs
+        self,
+        question: str,
+        k: int = 4,
+        *args,
+        query_vector: Optional[List[float]] = None,
+        **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Same search as :meth:`search`, pairing each hit with its L2 distance."""
+        """Same search as :meth:`search`, pairing each hit with its L2 distance.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
         # FAISS has no relevance-threshold knob; drop it so the per-source
         # score_threshold is safely ignored rather than crashing the forward.
         kwargs.pop("score_threshold", None)
         if self.index is None or self.index.ntotal == 0:
             return []
 
-        vector = np.array([self.embeddings.embed_query(question)], dtype=np.float32)
+        if query_vector is None:
+            query_vector = self.embeddings.embed_query(question)
+        vector = np.array([query_vector], dtype=np.float32)
         distances, rows = self.index.search(vector, min(k, self.index.ntotal))
 
         results = []

application/vectorstore/lancedb.py

@@ -80,10 +80,19 @@ class LanceDBVectorStore(BaseVectorStore):
         self.ensure_table_exists()
         self.docsearch.add(vectors)
 
-    def search(self, query: str, k: int = 2, *args, **kwargs):
-        """Search LanceDB for the top k most similar vectors."""
+    def search(self, query: str, k: int = 2, *args, query_vector=None, **kwargs):
+        """Search LanceDB for the top k most similar vectors.
+
+        Args:
+            query_vector: Precomputed embedding of ``query``; when given the
+                store skips embedding the query itself.
+        """
         self.ensure_table_exists()
-        query_embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, self.embeddings_key).embed_query(query)
+        query_embedding = query_vector
+        if query_embedding is None:
+            query_embedding = self._get_embeddings(
+                settings.EMBEDDINGS_NAME, self.embeddings_key
+            ).embed_query(query)
         results = self.docsearch.search(query_embedding).limit(k).to_list()
         return [(result["_distance"], result["text"], result["metadata"]) for result in results]
 

application/vectorstore/milvus.py

@@ -110,12 +110,20 @@ class MilvusStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: Optional[float] = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Search, pairing each hit with its cosine similarity."""
+        """Search, pairing each hit with its cosine similarity.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        if query_vector is None:
+            query_vector = self._embeddings.embed_query(question)
         results = self._client.search(
             collection_name=self._collection,
-            data=[self._embeddings.embed_query(question)],
+            data=[query_vector],
             filter=self._filter,
             limit=k,
             output_fields=["text", "metadata"],

application/vectorstore/mongodb.py

@@ -54,7 +54,9 @@ class MongoDBVectorStore(BaseVectorStore):
 
     score_kind = "cosine_similarity"
 
-    def search(self, question, k=2, *args, score_threshold=None, **kwargs):
+    def search(
+        self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
+    ):
         """Search via Atlas ``$vectorSearch``.
 
         Args:
@@ -62,21 +64,32 @@ class MongoDBVectorStore(BaseVectorStore):
             k: Maximum number of results.
             score_threshold: Optional ``vectorSearchScore`` floor in ``[0, 1]``;
                 results scoring below it are dropped.
+            query_vector: Precomputed embedding of ``question``, so a caller
+                searching several sources embeds the query only once.
         """
         return [
             doc
             for doc, _ in self.search_with_scores(
-                question, k, *args, score_threshold=score_threshold, **kwargs
+                question,
+                k,
+                *args,
+                score_threshold=score_threshold,
+                query_vector=query_vector,
+                **kwargs,
             )
         ]
 
-    def search_with_scores(self, question, k=2, *args, score_threshold=None, **kwargs):
+    def search_with_scores(
+        self, question, k=2, *args, score_threshold=None, query_vector=None, **kwargs
+    ):
         """Same search as :meth:`search`, pairing each hit with its score.
 
         The score is Atlas' ``vectorSearchScore`` — the same quantity
-        ``score_threshold`` is compared against.
+        ``score_threshold`` is compared against. ``query_vector`` skips the
+        per-store query embedding when the caller already has one.
         """
-        query_vector = self._embedding.embed_query(question)
+        if query_vector is None:
+            query_vector = self._embedding.embed_query(question)
 
         pipeline = [
             {

application/vectorstore/pgconn.py

@@ -0,0 +1,112 @@
+"""Process-wide psycopg connection pools keyed by DSN.
+
+Every component that talks to the pgvector database — the vector store and the
+GraphRAG graph store — checks out of the *same* pool for a given DSN, so a
+request that touches N sources costs pooled checkouts instead of N fresh TCP +
+TLS + auth handshakes. Keeping the registry in its own module (rather than in
+``pgvector``) is what lets ``application.graphrag.store`` join it without
+importing the vector store.
+
+The boot-time schema hook deliberately does not use this: pools are built
+lazily on first checkout, i.e. after the fork in a Celery/gunicorn worker, so a
+preloading parent never hands a live socket to a child.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any, Dict
+
+# Seconds a caller waits for a free pooled connection before giving up.
+POOL_TIMEOUT_SECONDS = 30.0
+
+# dsn -> psycopg_pool.ConnectionPool, one per process.
+_POOLS: Dict[str, Any] = {}
+_POOLS_LOCK = threading.Lock()
+
+
+def configure_pooled_connection(conn) -> None:
+    """Register pgvector's type adapters once per physical pooled connection.
+
+    Tolerates a database where ``CREATE EXTENSION vector`` has not run yet:
+    ``register_vector`` looks the type up in the catalog and raises when it is
+    absent, which would leave the pool unable to open the very connection the
+    schema bootstrap needs. The write path re-registers once the type exists.
+    """
+    from pgvector.psycopg import register_vector
+
+    try:
+        register_vector(conn)
+    except Exception as e:
+        logging.debug("pgvector types not registered yet: %s", e)
+
+
+def pool_for(dsn: str, max_size: int):
+    """Return this process's connection pool for ``dsn``, creating it lazily.
+
+    The pool is keyed by DSN alone and is built on first checkout — i.e. after
+    the fork in a Celery/gunicorn worker, so a preloading parent never hands a
+    live socket to a child. The boot-time schema hook must therefore use a
+    direct connection rather than this pool.
+
+    Args:
+        dsn: libpq connection string; the pool's only identity.
+        max_size: Upper bound on physical connections, used on creation only.
+
+    Returns:
+        psycopg_pool.ConnectionPool: The shared pool for ``dsn``.
+    """
+    pool = _POOLS.get(dsn)
+    if pool is not None:
+        return pool
+    with _POOLS_LOCK:
+        pool = _POOLS.get(dsn)
+        if pool is None:
+            from psycopg_pool import ConnectionPool
+
+            pool = ConnectionPool(
+                conninfo=dsn,
+                min_size=1,
+                max_size=max_size,
+                open=True,
+                timeout=POOL_TIMEOUT_SECONDS,
+                name="docsgpt-pgvector",
+                configure=configure_pooled_connection,
+                check=ConnectionPool.check_connection,
+            )
+            _POOLS[dsn] = pool
+    return pool
+
+
+def release(dsn: str, conn, pooled: bool) -> None:
+    """Hand ``conn`` back: to ``dsn``'s pool when pooled, else close it.
+
+    A pooled connection is rolled back first when it is still in a transaction,
+    so the next borrower gets a clean session. Never raises — releasing is
+    always best-effort cleanup on a path that is often a destructor.
+
+    Args:
+        dsn: Connection string the connection was opened for.
+        conn: The psycopg connection being released; ``None`` is a no-op.
+        pooled: Whether the connection came from :func:`pool_for`.
+    """
+    if conn is None:
+        return
+    try:
+        if pooled:
+            try:
+                if conn.info.transaction_status.name != "IDLE":
+                    conn.rollback()
+            except Exception:
+                # Connection already broken; the pool discards it.
+                pass
+            pool = _POOLS.get(dsn)
+            if pool is not None:
+                pool.putconn(conn)
+            elif not conn.closed:
+                conn.close()
+        elif not conn.closed:
+            conn.close()
+    except Exception as e:
+        logging.debug("Error releasing pooled connection: %s", e)

application/vectorstore/pgvector.py

@@ -6,12 +6,28 @@ from typing import List, Optional, Any, Dict
 from psycopg.types.json import Jsonb
 
 from application.core.settings import settings
+from application.vectorstore import pgconn
 from application.vectorstore.base import BaseVectorStore
 from application.vectorstore.document_class import Document
 
 # table name -> IVFFlat ``lists`` (None when the table has no such index)
 _IVFFLAT_LISTS_CACHE: Dict[str, Optional[int]] = {}
 
+DEFAULT_EMBEDDING_DIM = 768
+DEFAULT_POOL_MAX_SIZE = 8
+# Advisory-lock key shared with the boot hook so concurrent workers serialize DDL.
+SCHEMA_LOCK_KEY = "docsgpt:vectors:ddl"
+
+# The connection pools moved to ``application.vectorstore.pgconn`` so the graph
+# store can share them without importing this module. These names stay bound
+# here — same objects, not copies — because callers and tests reach for
+# ``pgvector._POOLS`` / patch ``pgvector._pool_for``.
+POOL_TIMEOUT_SECONDS = pgconn.POOL_TIMEOUT_SECONDS
+_POOLS = pgconn._POOLS
+_POOLS_LOCK = pgconn._POOLS_LOCK
+_configure_pooled_connection = pgconn.configure_pooled_connection
+_pool_for = pgconn.pool_for
+
 
 class PGVectorStore(BaseVectorStore):
     def __init__(
@@ -63,17 +79,64 @@ class PGVectorStore(BaseVectorStore):
         self._psycopg = psycopg
         self._register_vector = register_vector
         self._connection = None
-        self._ensure_table_exists()
+        self._pooled = False
+        self._pool_max_size = self._resolve_pool_max_size()
+        # No DDL here. The retriever builds one store per source per request, so
+        # construction must stay free: schema is owned at boot by
+        # ``ensure_vector_schema``, and the write path below re-checks it once as
+        # a safety net for processes that never ran the boot hook.
+        self._schema_ensured = False
+
+    @staticmethod
+    def _resolve_pool_max_size() -> int:
+        """Pool size from settings, defensively — 0 means one direct connection."""
+        value = getattr(settings, "PGVECTOR_POOL_MAX_SIZE", DEFAULT_POOL_MAX_SIZE)
+        if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+            return value
+        return DEFAULT_POOL_MAX_SIZE
 
     def _get_connection(self):
-        """Get or create database connection"""
+        """Get or create this store's connection, pooled unless pooling is off."""
         if self._connection is None or self._connection.closed:
-            self._connection = self._psycopg.connect(self._connection_string)
-            # Register pgvector types
-            self._register_vector(self._connection)
-            self._apply_ivfflat_probes(self._connection)
+            if self._pool_max_size > 0:
+                self._connection = _pool_for(
+                    self._connection_string, self._pool_max_size
+                ).getconn()
+                self._pooled = True
+            else:
+                self._connection = self._psycopg.connect(self._connection_string)
+                # Register pgvector types
+                self._register_pgvector_types(self._connection)
+                self._pooled = False
+            self._apply_probes_once(self._connection)
         return self._connection
 
+    def _register_pgvector_types(self, conn) -> None:
+        """Register pgvector's adapters, tolerating a not-yet-created extension."""
+        try:
+            self._register_vector(conn)
+        except Exception as e:
+            logging.debug("pgvector types not registered yet: %s", e)
+
+    def _apply_probes_once(self, conn) -> None:
+        """Set ``ivfflat.probes`` once per physical connection and table.
+
+        A pooled connection outlives the store that checked it out, and the
+        ``SET`` is session-level, so the probe lookup must not repeat on every
+        checkout. The marker lives on the connection object itself because that
+        is what the pool recycles.
+        """
+        if getattr(conn, "_docsgpt_probes_table", None) == self._table_name:
+            return
+        self._apply_ivfflat_probes(conn)
+        try:
+            # ``SET`` (not SET LOCAL) is undone by a rollback, so make it stick
+            # before any query runs on this connection.
+            conn.commit()
+            conn._docsgpt_probes_table = self._table_name
+        except Exception as e:  # never let session tuning break a query
+            logging.debug("Could not persist ivfflat.probes: %s", e)
+
     def _ivfflat_lists(self, conn) -> Optional[int]:
         """Return the ``lists`` value of this table's IVFFlat index, if any.
 
@@ -125,30 +188,48 @@ class PGVectorStore(BaseVectorStore):
         except Exception as e:  # older pgvector / no index — search still works
             logging.debug("Could not set ivfflat.probes: %s", e)
 
-    def _ensure_table_exists(self):
-        """Create table and enable pgvector extension if they don't exist"""
-        conn = self._get_connection()
+    @staticmethod
+    def create_schema(
+        conn,
+        *,
+        table_name: str = "documents",
+        vector_column: str = "embedding",
+        text_column: str = "text",
+        metadata_column: str = "metadata",
+        dimension: int = DEFAULT_EMBEDDING_DIM,
+    ) -> None:
+        """Create the extension, table and indexes on ``conn`` without committing.
+
+        Shared by the boot hook (``ensure_vector_schema``) and the store's own
+        write-path safety net; the caller owns the transaction.
+
+        Args:
+            conn: Open psycopg connection.
+            table_name: Documents table to create.
+            vector_column: Embedding column name.
+            text_column: Chunk-text column name.
+            metadata_column: JSONB metadata column name.
+            dimension: Width of the embedding vectors.
+        """
         cursor = conn.cursor()
-        
+
         try:
             # Enable pgvector extension
             cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
-            
-            embedding_dim = getattr(self._embedding, 'dimension', 768)
-            
+
             # Create table with vector column
             create_table_query = f"""
-            CREATE TABLE IF NOT EXISTS {self._table_name} (
+            CREATE TABLE IF NOT EXISTS {table_name} (
                 id SERIAL PRIMARY KEY,
-                {self._text_column} TEXT NOT NULL,
-                {self._vector_column} vector({embedding_dim}),
-                {self._metadata_column} JSONB,
+                {text_column} TEXT NOT NULL,
+                {vector_column} vector({dimension}),
+                {metadata_column} JSONB,
                 source_id TEXT NOT NULL,
                 created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
             );
             """
             cursor.execute(create_table_query)
-            
+
             # NO vector index is created here, deliberately.
             #
             # This runs when the table is first created, i.e. on an EMPTY
@@ -166,25 +247,88 @@ class PGVectorStore(BaseVectorStore):
             # data, once a corpus is large enough to need one.
             # Create index for source_id filtering
             source_index_query = f"""
-            CREATE INDEX IF NOT EXISTS {self._table_name}_source_id_idx
-            ON {self._table_name} (source_id);
+            CREATE INDEX IF NOT EXISTS {table_name}_source_id_idx
+            ON {table_name} (source_id);
             """
             cursor.execute(source_index_query)
 
             # Functional GIN index backing keyword_search full-text queries.
             fts_index_query = f"""
-            CREATE INDEX IF NOT EXISTS {self._table_name}_text_fts_idx
-            ON {self._table_name} USING gin(to_tsvector('english', {self._text_column}));
+            CREATE INDEX IF NOT EXISTS {table_name}_text_fts_idx
+            ON {table_name} USING gin(to_tsvector('english', {text_column}));
             """
             cursor.execute(fts_index_query)
+        finally:
+            cursor.close()
 
+    @staticmethod
+    def table_dimension(
+        conn, table_name: str = "documents", vector_column: str = "embedding"
+    ) -> Optional[int]:
+        """Width declared by the table's vector column, or ``None`` if unknown.
+
+        ``None`` means the table is absent or the column is not a ``vector``.
+        """
+        cursor = conn.cursor()
+        try:
+            cursor.execute("SELECT to_regclass(%s);", (table_name,))
+            row = cursor.fetchone()
+            if not row or row[0] is None:
+                return None
+            cursor.execute(
+                "SELECT format_type(a.atttypid, a.atttypmod) FROM pg_attribute a "
+                "WHERE a.attrelid = %s::regclass AND a.attname = %s "
+                "AND NOT a.attisdropped",
+                (table_name, vector_column),
+            )
+            row = cursor.fetchone()
+            if not row or not row[0]:
+                return None
+            match = re.search(r"vector\((\d+)\)", str(row[0]))
+            return int(match.group(1)) if match else None
+        finally:
+            cursor.close()
+
+    def _ensure_table_exists(self) -> None:
+        """Create this store's schema under an advisory lock, then commit."""
+        conn = self._get_connection()
+        try:
+            cursor = conn.cursor()
+            try:
+                cursor.execute(
+                    "SELECT pg_advisory_xact_lock(hashtext(%s));", (SCHEMA_LOCK_KEY,)
+                )
+            finally:
+                cursor.close()
+            self.create_schema(
+                conn,
+                table_name=self._table_name,
+                vector_column=self._vector_column,
+                text_column=self._text_column,
+                metadata_column=self._metadata_column,
+                dimension=getattr(self._embedding, "dimension", DEFAULT_EMBEDDING_DIM),
+            )
             conn.commit()
+            # The extension may have just been created; pick up its adapters so
+            # the insert that follows can bind a vector.
+            self._register_pgvector_types(conn)
         except Exception as e:
             conn.rollback()
             logging.error(f"Error creating table: {e}")
             raise
-        finally:
-            cursor.close()
+
+    def _ensure_schema_once(self) -> None:
+        """Create the schema on this instance's first write, at most once.
+
+        Readers never create schema — boot owns it. This is the safety net for a
+        process that never ran the boot hook (scripts, tests, the first ingest on
+        a brand-new deployment); a read against a missing table surfaces
+        psycopg's own error, like every other store.
+        """
+        if self._schema_ensured:
+            return
+        self._ensure_table_exists()
+        self._schema_ensured = True
 
     score_kind = "cosine_similarity"
 
@@ -194,6 +338,7 @@ class PGVectorStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: float = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Document]:
         """Search for similar documents using vector similarity.
@@ -204,11 +349,19 @@ class PGVectorStore(BaseVectorStore):
             score_threshold: Optional cosine-similarity floor in ``[0, 1]``.
                 Cosine distance = ``1 - similarity``; rows with similarity below
                 the threshold (distance above ``1 - threshold``) are dropped.
+            query_vector: Precomputed embedding of ``question``. Supplied by a
+                caller searching several sources with one query, so the query is
+                embedded once instead of once per store.
         """
         return [
             doc
             for doc, _ in self.search_with_scores(
-                question, k, *args, score_threshold=score_threshold, **kwargs
+                question,
+                k,
+                *args,
+                score_threshold=score_threshold,
+                query_vector=query_vector,
+                **kwargs,
             )
         ]
 
@@ -291,6 +444,7 @@ class PGVectorStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: float = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[tuple]:
         """Same search as :meth:`search`, pairing each hit with its similarity.
@@ -298,8 +452,13 @@ class PGVectorStore(BaseVectorStore):
         The score is the cosine similarity (``1 - cosine_distance``) — the exact
         quantity ``score_threshold`` is compared against, so a caller can read a
         result's score and pick a threshold from it directly.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
         """
-        query_vector = self._embedding.embed_query(question)
+        if query_vector is None:
+            query_vector = self._embedding.embed_query(question)
 
         conn = self._get_connection()
         cursor = conn.cursor()
@@ -334,6 +493,11 @@ class PGVectorStore(BaseVectorStore):
                     (Document(page_content=text, metadata=metadata), score)
                 )
 
+            # End the read transaction. On a persistent (pooled) connection an
+            # uncommitted SELECT leaves the backend "idle in transaction",
+            # pinning a snapshot and blocking VACUUM for as long as the store
+            # lives.
+            conn.commit()
             return documents
 
         except Exception as e:
@@ -381,10 +545,16 @@ class PGVectorStore(BaseVectorStore):
                 metadata = metadata or {}
                 documents.append(Document(page_content=text, metadata=metadata))
 
+            conn.commit()
             return documents
 
         except Exception as e:
             logging.error(f"Error in keyword search: {e}", exc_info=True)
+            try:
+                conn.rollback()
+            except Exception:
+                # Connection already gone; nothing left to roll back.
+                pass
             return []
         finally:
             cursor.close()
@@ -402,7 +572,8 @@ class PGVectorStore(BaseVectorStore):
 
         embeddings = self._embedding.embed_documents(texts)
         metadatas = metadatas or [{}] * len(texts)
-        
+
+        self._ensure_schema_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         
@@ -474,11 +645,17 @@ class PGVectorStore(BaseVectorStore):
                     "text": text,
                     "metadata": metadata or {}
                 })
-            
+
+            conn.commit()
             return chunks
-            
+
         except Exception as e:
             logging.error(f"Error getting chunks: {e}")
+            try:
+                conn.rollback()
+            except Exception:
+                # Connection already gone; nothing left to roll back.
+                pass
             return []
         finally:
             cursor.close()
@@ -495,7 +672,8 @@ class PGVectorStore(BaseVectorStore):
 
         if not embeddings:
             raise ValueError("Could not generate embedding for chunk")
-        
+
+        self._ensure_schema_once()
         conn = self._get_connection()
         cursor = conn.cursor()
         
@@ -570,7 +748,22 @@ class PGVectorStore(BaseVectorStore):
         finally:
             cursor.close()
 
+    def close(self) -> None:
+        """Release this store's connection: back to the pool, or closed outright.
+
+        A pooled connection is rolled back first when it is still in a
+        transaction, so the next borrower gets a clean session.
+        """
+        conn = getattr(self, "_connection", None)
+        if conn is None:
+            return
+        self._connection = None
+        pgconn.release(self._connection_string, conn, getattr(self, "_pooled", False))
+
     def __del__(self):
-        """Close database connection when object is destroyed"""
-        if hasattr(self, '_connection') and self._connection and not self._connection.closed:
-            self._connection.close()
+        """Release the connection when the object is destroyed. Never raises."""
+        try:
+            self.close()
+        except Exception:
+            # Interpreter teardown can null out module globals; never raise here.
+            pass

application/vectorstore/qdrant.py

@@ -109,12 +109,20 @@ class QdrantStore(BaseVectorStore):
         k: int = 2,
         *args,
         score_threshold: Optional[float] = None,
+        query_vector: Optional[List[float]] = None,
         **kwargs,
     ) -> List[Tuple[Document, float]]:
-        """Search, pairing each hit with its cosine similarity."""
+        """Search, pairing each hit with its cosine similarity.
+
+        Args:
+            query_vector: Precomputed embedding of ``question``; when given the
+                store skips embedding the query itself.
+        """
+        if query_vector is None:
+            query_vector = self._embeddings.embed_query(question)
         hits = self._client.query_points(
             collection_name=self._collection,
-            query=self._embeddings.embed_query(question),
+            query=query_vector,
             query_filter=self._filter,
             limit=k,
             with_payload=True,

tests/agents/test_base_agent.py

@@ -3,6 +3,13 @@ from unittest.mock import Mock, patch
 
 import pytest
 from application.agents.classic_agent import ClassicAgent
+from application.llm.anthropic import AnthropicLLM
+from application.llm.docsgpt_provider import DocsGPTAPILLM
+from application.llm.google_ai import GoogleLLM
+from application.llm.groq import GroqLLM
+from application.llm.novita import NovitaLLM
+from application.llm.open_router import OpenRouterLLM
+from application.llm.openai import OpenAILLM
 
 
 @pytest.mark.unit
@@ -620,24 +627,23 @@ class TestBaseAgentLLMGeneration:
     def test_llm_gen_with_json_schema(
         self,
         agent_base_params,
-        mock_llm,
-        mock_llm_creator,
         mock_llm_handler_creator,
         log_context,
     ):
-        mock_llm._supports_structured_output = Mock(return_value=True)
-        mock_llm.prepare_structured_output_format = Mock(
-            return_value={"schema": "test"}
-        )
+        # A real OpenAILLM: the format kwarg is chosen from the LLM class,
+        # so a bare Mock deliberately gets no structured output.
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        llm.gen_stream = Mock()
 
         agent_base_params["json_schema"] = {"type": "object"}
         agent_base_params["llm_name"] = "openai"
+        agent_base_params["llm"] = llm
         agent = ClassicAgent(**agent_base_params)
 
         messages = [{"role": "user", "content": "test"}]
         agent._llm_gen(messages, log_context)
 
-        call_args = mock_llm.gen_stream.call_args[1]
+        call_args = llm.gen_stream.call_args[1]
         assert "response_format" in call_args
 
 
@@ -1163,24 +1169,21 @@ class TestLLMGenAdvanced:
     def test_llm_gen_google_structured_output(
         self,
         agent_base_params,
-        mock_llm,
-        mock_llm_creator,
         mock_llm_handler_creator,
         log_context,
     ):
-        mock_llm._supports_structured_output = Mock(return_value=True)
-        mock_llm.prepare_structured_output_format = Mock(
-            return_value={"schema": "test"}
-        )
+        llm = GoogleLLM(api_key="g-test", user_api_key=None)
+        llm.gen_stream = Mock()
 
         agent_base_params["json_schema"] = {"type": "object"}
         agent_base_params["llm_name"] = "google"
+        agent_base_params["llm"] = llm
         agent = ClassicAgent(**agent_base_params)
 
         messages = [{"role": "user", "content": "test"}]
         agent._llm_gen(messages, log_context)
 
-        call_kwargs = mock_llm.gen_stream.call_args[1]
+        call_kwargs = llm.gen_stream.call_args[1]
         assert "response_schema" in call_kwargs
 
     def test_llm_gen_no_tools_when_unsupported(
@@ -1529,3 +1532,295 @@ class TestBaseAgentDocumentBudgetOrdering:
         assert "<documents>" in user, "documents must survive a long question"
         assert agent.retrieved_docs, "documents must not all be shed"
         assert len(user) < len(huge_question), "question must be truncated"
+
+
+# ---------------------------------------------------------------------------
+# _llm_gen — structured-output provider gating
+# ---------------------------------------------------------------------------
+
+
+SCHEMA = {"type": "object", "properties": {"answer": {"type": "string"}}}
+
+
+def _build_agent_with_llm(agent_base_params, llm, llm_name, **overrides):
+    """Agent wired to a real LLM instance whose ``gen_stream`` is captured."""
+    llm.gen_stream = Mock()
+    agent_base_params["llm"] = llm
+    agent_base_params["llm_name"] = llm_name
+    agent_base_params.update(overrides)
+    return ClassicAgent(**agent_base_params)
+
+
+@pytest.mark.unit
+class TestLLMGenStructuredOutputProviderGating:
+    """Structured output must key off the LLM class, not ``llm_name``.
+
+    Catalog models loaded from MODELS_CONFIG_DIR dispatch under
+    ``openai_compatible`` (and ``groq`` / ``novita`` / ``openrouter`` /
+    ``docsgpt``), all of which instantiate an ``OpenAILLM`` subclass. A
+    string comparison against "openai" dropped guided decoding for every
+    one of them.
+    """
+
+    @pytest.mark.parametrize(
+        "llm_name,llm_factory",
+        [
+            ("openai", lambda: OpenAILLM(api_key="sk-test", user_api_key=None)),
+            (
+                "openai_compatible",
+                lambda: OpenAILLM(
+                    api_key="sk-test",
+                    user_api_key=None,
+                    base_url="https://vllm.example.invalid/v1",
+                ),
+            ),
+            ("groq", lambda: GroqLLM(api_key="gsk-test", user_api_key=None)),
+            ("novita", lambda: NovitaLLM(api_key="nv-test", user_api_key=None)),
+            (
+                "openrouter",
+                lambda: OpenRouterLLM(api_key="or-test", user_api_key=None),
+            ),
+            ("docsgpt", lambda: DocsGPTAPILLM(user_api_key=None)),
+        ],
+    )
+    def test_openai_wire_providers_get_response_format(
+        self,
+        agent_base_params,
+        mock_llm_handler_creator,
+        log_context,
+        llm_name,
+        llm_factory,
+    ):
+        llm = llm_factory()
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, llm_name, json_schema=SCHEMA
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert call_kwargs["response_format"]["type"] == "json_schema"
+        assert "response_schema" not in call_kwargs
+
+    @pytest.mark.parametrize(
+        "llm_name",
+        ["openai", "openai_compatible", "groq", "novita", "openrouter", "docsgpt"],
+    )
+    def test_openai_wire_providers_get_json_object_mode(
+        self, agent_base_params, mock_llm_handler_creator, log_context, llm_name
+    ):
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, llm_name, json_object=True
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert call_kwargs["response_format"] == {"type": "json_object"}
+
+    def test_google_llm_gets_response_schema(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        llm = GoogleLLM(api_key="g-test", user_api_key=None)
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "google", json_schema=SCHEMA
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert "response_schema" in call_kwargs
+        assert "response_format" not in call_kwargs
+
+    def test_google_llm_gets_no_json_object_mode(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        """``{"type": "json_object"}`` is an OpenAI-wire shape only."""
+        llm = GoogleLLM(api_key="g-test", user_api_key=None)
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "google", json_object=True
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+        assert "response_schema" not in call_kwargs
+
+    def test_non_openai_wire_llm_gets_neither(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        """Even a capability-claiming Anthropic LLM takes no format kwarg."""
+        llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
+        llm._supports_structured_output = Mock(return_value=True)
+        llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "anthropic", json_schema=SCHEMA
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+        assert "response_schema" not in call_kwargs
+
+    def test_non_openai_wire_llm_gets_no_json_object_mode(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
+        llm._supports_structured_output = Mock(return_value=True)
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "anthropic", json_object=True
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        assert "response_format" not in llm.gen_stream.call_args[1]
+
+    def test_llm_double_without_class_declaration_gets_neither(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        """The kwarg name is read off the LLM *class*.
+
+        A ``Mock`` answers every attribute lookup with a fresh truthy Mock, so
+        reading it off the instance would smuggle a Mock-named kwarg into the
+        gen call. Its class declares nothing, so it stays unstructured.
+        """
+        llm = Mock()
+        llm._supports_structured_output = Mock(return_value=True)
+        llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "custom", json_schema=SCHEMA
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+        assert "response_schema" not in call_kwargs
+        assert all(isinstance(key, str) for key in call_kwargs)
+
+    def test_declaration_comes_from_the_llm_class_attribute(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        """``BaseLLM.structured_output_kwarg`` is the single source of truth —
+        the agent no longer isinstance-checks provider classes."""
+        llm = AnthropicLLM(api_key="ant-test", user_api_key=None)
+        llm._supports_structured_output = Mock(return_value=True)
+        llm.prepare_structured_output_format = Mock(return_value={"schema": "x"})
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "anthropic", json_schema=SCHEMA
+        )
+        assert agent._structured_output_kwarg() is None
+
+        with patch.object(
+            type(llm), "structured_output_kwarg", "response_format", create=True
+        ):
+            assert agent._structured_output_kwarg() == "response_format"
+
+    def test_capability_flag_still_vetoes_openai_compatible(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        """A model whose registry caps deny structured output stays clean."""
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        llm._supports_structured_output = Mock(return_value=False)
+        agent = _build_agent_with_llm(
+            agent_base_params, llm, "openai_compatible", json_schema=SCHEMA
+        )
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        call_kwargs = llm.gen_stream.call_args[1]
+        assert "response_format" not in call_kwargs
+        assert "response_schema" not in call_kwargs
+
+
+# ---------------------------------------------------------------------------
+# _llm_gen — tools gate
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestLLMGenToolsGate:
+    """``_supports_tools`` must be *called*, not merely looked up.
+
+    Testing the bound method for truthiness always passed, so the
+    agent-level gate was dead and only the provider-level drop kept an
+    unsupported ``tools`` payload off the wire.
+    """
+
+    def test_tools_dropped_when_llm_reports_unsupported(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        llm.gen_stream = Mock()
+        llm._supports_tools = Mock(return_value=False)
+        agent_base_params["llm"] = llm
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        assert "tools" not in llm.gen_stream.call_args[1]
+
+    def test_tools_attached_when_llm_reports_supported(
+        self, agent_base_params, mock_llm_handler_creator, log_context
+    ):
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        llm.gen_stream = Mock()
+        llm._supports_tools = Mock(return_value=True)
+        agent_base_params["llm"] = llm
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        agent._llm_gen([{"role": "user", "content": "test"}], log_context)
+
+        assert llm.gen_stream.call_args[1]["tools"] == agent.tools
+
+    def test_bool_attribute_double_is_honored(
+        self,
+        agent_base_params,
+        mock_llm,
+        mock_llm_creator,
+        mock_llm_handler_creator,
+    ):
+        """Test doubles set ``_supports_tools`` as a plain bool."""
+        mock_llm._supports_tools = True
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        agent._llm_gen([{"role": "user", "content": "test"}])
+
+        assert mock_llm.gen_stream.call_args[1]["tools"] == agent.tools
+
+    def test_missing_capability_check_drops_tools(
+        self, agent_base_params, mock_llm_handler_creator
+    ):
+        llm = Mock(spec=["gen_stream", "model_id"])
+        llm.gen_stream = Mock()
+        llm.model_id = "gpt-4"
+        agent_base_params["llm"] = llm
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        agent._llm_gen([{"role": "user", "content": "test"}])
+
+        assert "tools" not in llm.gen_stream.call_args[1]
+
+    def test_unimplemented_capability_check_still_sends_tools(
+        self, agent_base_params, mock_llm_handler_creator
+    ):
+        """``BaseLLM._supports_tools`` raises; the provider drops them."""
+        llm = OpenAILLM(api_key="sk-test", user_api_key=None)
+        llm.gen_stream = Mock()
+        llm._supports_tools = Mock(
+            side_effect=NotImplementedError("Subclass must implement")
+        )
+        agent_base_params["llm"] = llm
+        agent = ClassicAgent(**agent_base_params)
+        agent.tools = [{"type": "function", "function": {"name": "test"}}]
+
+        agent._llm_gen([{"role": "user", "content": "test"}])
+
+        assert llm.gen_stream.call_args[1]["tools"] == agent.tools

tests/agents/test_workflow_input_documents.py

@@ -592,3 +592,218 @@ def test_node_extract_cap_bounds_parse_attempts_even_when_every_parse_times_out(
     assert parse_calls["n"] == 2  # not 4 -- failed parses still consume cap budget
     notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID]
     assert len(notes) == 1
+
+
+# ---------------------------------------------------------------------------
+# Reuse of the text ``store_attachment`` already extracted (no double OCR)
+# ---------------------------------------------------------------------------
+
+DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+_OK_EXTRACTION = {
+    "status": "ok",
+    "parser": "DoclingParser",
+    "truncated": False,
+    "original_tokens": 12,
+    "stored_tokens": 12,
+}
+
+
+def _with_extraction(attachment: dict, content, extraction=_OK_EXTRACTION) -> dict:
+    """Add the ``content`` + ``metadata.extraction`` shape ``store_attachment`` persists."""
+    attachment["content"] = content
+    attachment["metadata"] = {} if extraction is None else {"extraction": dict(extraction)}
+    return attachment
+
+
+def _forbid_parse(monkeypatch) -> None:
+    """Make any enqueue of the parsing worker a hard failure."""
+    import application.api.user.tasks as tasks
+
+    def _apply_async(*a, **k):
+        raise AssertionError("the document must not be re-parsed")
+
+    monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
+
+
+def _count_parses(monkeypatch, calls: dict) -> None:
+    """Stub the parsing worker so each blocking parse succeeds and is counted."""
+    import application.api.user.tasks as tasks
+
+    class _R:
+        def get(self, timeout=None, disable_sync_subtasks=True):
+            return {"status": "ok", "content": "PARSED"}
+
+    def _apply_async(args=None, queue=None, **kw):
+        calls["n"] = calls.get("n", 0) + 1
+        return _R()
+
+    monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
+
+
+def _bridged_engine(pg_engine, tmp_path, monkeypatch, attachments):
+    """Run the bridge for ``attachments`` and return (engine, refs)."""
+    wf_id = _make_workflow(pg_engine)
+    agent = _agent(wf_id, attachments)
+    _patch_engine(monkeypatch)
+    list(agent._gen_inner("summarize", log_context=None))
+    engine = _RecordingEngine.instances[-1]
+    return engine, engine.captured_inputs["input_documents"]
+
+
+def test_preextracted_attachment_text_is_reused_without_reparsing(
+    pg_engine, tmp_path, monkeypatch
+):
+    """An attachment already parsed at upload time is inlined verbatim; the worker is never called."""
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    att = _with_extraction(
+        _stage_attachment(storage, b"PK-docx-bytes", "scan.docx", DOCX_MIME),
+        "PRE-EXTRACTED OCR TEXT",
+    )
+    engine, refs = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att])
+
+    # The text rides on the engine only -- never in the persisted run state/refs.
+    aid = refs[0]["artifact_id"]
+    assert engine.preextracted_text == {aid: "PRE-EXTRACTED OCR TEXT"}
+    assert set(refs[0]) == {"artifact_id", "ref", "filename", "mime_type"}
+    assert "content" not in str(engine.state["input_documents"])
+
+    _forbid_parse(monkeypatch)
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+
+    assert out == [{"id": aid, "mime_type": "text/plain", "content": "PRE-EXTRACTED OCR TEXT"}]
+
+
+def test_preextracted_text_is_bounded_head_and_tail(pg_engine, tmp_path, monkeypatch):
+    """Reused text goes through the same head+tail window as the inline-text path."""
+    from application.parser.document_reader import _TEXT_MAX_BYTES
+
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    big = "A" * (_TEXT_MAX_BYTES * 3)
+    att = _with_extraction(
+        _stage_attachment(storage, b"PK-docx-bytes", "huge.docx", DOCX_MIME), big
+    )
+    engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att])
+
+    _forbid_parse(monkeypatch)
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+
+    content = out[0]["content"]
+    assert "...[truncated" in content
+    assert len(content.encode("utf-8")) < _TEXT_MAX_BYTES + 200
+
+
+def test_truncated_attachment_text_falls_through_to_parse(pg_engine, tmp_path, monkeypatch):
+    """A row cut at ATTACHMENT_MAX_TOKENS lost its tail, so the node re-parses instead."""
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    att = _with_extraction(
+        _stage_attachment(storage, b"PK-docx-bytes", "long.docx", DOCX_MIME),
+        "HEAD ONLY",
+        extraction={**_OK_EXTRACTION, "truncated": True},
+    )
+    engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att])
+
+    assert engine.preextracted_text == {}
+    calls: dict = {}
+    _count_parses(monkeypatch, calls)
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+
+    assert calls["n"] == 1
+    assert out[0]["content"] == "PARSED"
+
+
+def test_failed_or_empty_attachment_text_falls_through_to_parse(
+    pg_engine, tmp_path, monkeypatch
+):
+    """An empty/whitespace ``content`` (or a failed extraction) is never reused."""
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    empty = _with_extraction(
+        _stage_attachment(storage, b"PK-a", "empty.docx", DOCX_MIME), "   "
+    )
+    failed = _with_extraction(
+        _stage_attachment(storage, b"PK-b", "failed.docx", DOCX_MIME),
+        None,
+        extraction={"status": "failed", "parser": None, "truncated": False, "error": "boom"},
+    )
+    engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, [empty, failed])
+
+    assert engine.preextracted_text == {}
+    calls: dict = {}
+    _count_parses(monkeypatch, calls)
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+
+    assert calls["n"] == 2
+    assert [a["content"] for a in out] == ["PARSED", "PARSED"]
+
+
+def test_legacy_row_without_extraction_metadata_is_trusted(pg_engine, tmp_path, monkeypatch):
+    """A pre-``extraction`` row with non-empty content is reused (nothing records otherwise)."""
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    att = _stage_attachment(storage, b"PK-docx-bytes", "legacy.docx", DOCX_MIME)
+    att["content"] = "LEGACY TEXT"  # no ``metadata`` key at all
+    engine, refs = _bridged_engine(pg_engine, tmp_path, monkeypatch, [att])
+
+    assert engine.preextracted_text == {refs[0]["artifact_id"]: "LEGACY TEXT"}
+    _forbid_parse(monkeypatch)
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+    assert out[0]["content"] == "LEGACY TEXT"
+
+
+def test_preextracted_docs_do_not_consume_the_extract_cap(pg_engine, tmp_path, monkeypatch):
+    """Reused text costs no blocking parse, so it must not spend the per-node parse budget."""
+    from application.core.settings import settings
+
+    storage = _wire(pg_engine, tmp_path, monkeypatch)
+    attachments = [
+        _with_extraction(
+            _stage_attachment(storage, b"PK-0", "pre0.docx", DOCX_MIME), "REUSED 0"
+        ),
+        _with_extraction(
+            _stage_attachment(storage, b"PK-1", "pre1.docx", DOCX_MIME), "REUSED 1"
+        ),
+        _stage_attachment(storage, b"PK-2", "raw2.docx", DOCX_MIME),
+        _stage_attachment(storage, b"PK-3", "raw3.docx", DOCX_MIME),
+    ]
+    engine, _ = _bridged_engine(pg_engine, tmp_path, monkeypatch, attachments)
+
+    monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False)
+    calls: dict = {}
+    _count_parses(monkeypatch, calls)
+
+    out = engine._materialize_node_attachments(
+        AgentNodeConfig(input_documents=["*"]), "Reviewer", supported_types=[]
+    )
+
+    # All four documents made it in: the two reused ones did not spend cap budget,
+    # so both remaining documents still fit the cap of 2 blocking parses.
+    assert [a["content"] for a in out] == ["REUSED 0", "REUSED 1", "PARSED", "PARSED"]
+    assert calls["n"] == 2
+    assert not [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID]
+
+
+@pytest.mark.parametrize(
+    "attachment,expected",
+    [
+        ({"content": "text", "metadata": {"extraction": {"status": "ok", "truncated": False}}}, "text"),
+        ({"content": "text", "metadata": {"extraction": {"status": "ok", "truncated": True}}}, None),
+        ({"content": "text", "metadata": {"extraction": {"status": "failed"}}}, None),
+        ({"content": "text", "metadata": {}}, "text"),
+        ({"content": "text"}, "text"),
+        ({"content": "text", "metadata": {"extraction": "nope"}}, None),
+        ({"content": "", "metadata": {"extraction": {"status": "ok", "truncated": False}}}, None),
+        ({"content": None}, None),
+        ({}, None),
+    ],
+)
+def test_usable_attachment_text_matrix(attachment, expected):
+    """Only whole, successfully-extracted, non-empty text is eligible for reuse."""
+    assert WorkflowAgent._usable_attachment_text(attachment) == expected

tests/agents/test_workflow_node_attachments.py

@@ -23,6 +23,7 @@ from application.agents.workflows.schemas import (
     WorkflowNode,
 )
 from application.agents.workflows.workflow_engine import WorkflowEngine
+from application.core.settings import settings
 from application.storage.db.repositories.artifacts import ArtifactsRepository
 
 RUN_ID = "11111111-1111-1111-1111-111111111111"
@@ -231,7 +232,7 @@ def test_extract_non_text_uses_parse_worker(monkeypatch):
     _patch_repo(monkeypatch, {aid: rec})
     eng = _engine(monkeypatch)
     monkeypatch.setattr(
-        WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED MD"
+        WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED MD"
     )
     cfg = _node_config(input_documents=[aid], file_passing="extract")
 
@@ -295,7 +296,7 @@ def test_native_file_cap_bounds_native_then_extracts(monkeypatch):
     )
     # extract of a non-text image routes through the parsing worker; stub it so it returns text.
     monkeypatch.setattr(
-        WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "fallback text"
+        WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "fallback text"
     )
     cfg = _node_config(input_documents=ids, file_passing="auto")
 
@@ -383,6 +384,93 @@ def test_duplicate_refs_attach_once(monkeypatch):
     assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}]
 
 
+# ---------------------------------------------------------------------------
+# Size-scaled parse window
+# ---------------------------------------------------------------------------
+
+
+def _capture_parse(monkeypatch) -> dict:
+    """Patch parse_document.apply_async so nothing touches a broker; capture the call."""
+    import application.api.user.tasks as tasks
+
+    captured: Dict[str, Any] = {}
+
+    class _R:
+        def get(self, timeout=None, disable_sync_subtasks=True):
+            captured["timeout"] = timeout
+            return {"status": "ok", "content": "PARSED"}
+
+    def _apply_async(args=None, queue=None, **kw):
+        captured["kwargs"] = kw
+        return _R()
+
+    monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
+    return captured
+
+
+@pytest.mark.unit
+def test_parse_window_scales_with_the_document_size(monkeypatch):
+    """The version's ``size`` reaches the parse: a longer await AND matching Celery limits."""
+    from application.api.user.tasks import parse_task_time_limits, parse_timeout_for_size
+
+    size = 8 * 1024 * 1024
+    aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=size)
+    _patch_repo(monkeypatch, {aid: rec})
+    eng = _engine(monkeypatch)
+    captured = _capture_parse(monkeypatch)
+
+    out = eng._materialize_node_attachments(
+        _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
+    )
+
+    expected = parse_timeout_for_size(size)
+    assert out == [{"id": aid, "mime_type": "text/plain", "content": "PARSED"}]
+    assert captured["timeout"] == expected
+    # Without the per-call limits the worker self-terminates at the base timeout,
+    # so the longer await would buy nothing.
+    assert captured["kwargs"] == parse_task_time_limits(expected)
+    assert captured["kwargs"]["soft_time_limit"] > settings.DOCUMENT_PARSE_TIMEOUT
+
+
+@pytest.mark.unit
+def test_parse_window_falls_back_to_the_base_timeout_without_a_size(monkeypatch):
+    """A NULL/missing version size keeps the configured base window as the floor."""
+    aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=None)
+    _patch_repo(monkeypatch, {aid: rec})
+    eng = _engine(monkeypatch)
+    captured = _capture_parse(monkeypatch)
+
+    eng._materialize_node_attachments(
+        _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
+    )
+
+    assert captured["timeout"] == float(settings.DOCUMENT_PARSE_TIMEOUT)
+    assert captured["kwargs"]["soft_time_limit"] == settings.DOCUMENT_PARSE_TIMEOUT
+
+
+@pytest.mark.unit
+def test_preextracted_text_skips_the_parse_worker(monkeypatch):
+    """Text already extracted at upload time is inlined; the parsing worker is never called."""
+    aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf")
+    _patch_repo(monkeypatch, {aid: rec})
+    eng = _engine(monkeypatch)
+    eng.preextracted_text[aid] = "ALREADY OCRED"
+
+    import application.api.user.tasks as tasks
+
+    monkeypatch.setattr(
+        tasks.parse_document,
+        "apply_async",
+        lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not re-parse")),
+    )
+
+    out = eng._materialize_node_attachments(
+        _node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
+    )
+
+    assert out == [{"id": aid, "mime_type": "text/plain", "content": "ALREADY OCRED"}]
+
+
 # ---------------------------------------------------------------------------
 # _execute_agent_node wiring
 # ---------------------------------------------------------------------------
@@ -484,7 +572,7 @@ def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch):
     _patch_repo(monkeypatch, {aid: rec})
     _patch_capabilities(monkeypatch)
     monkeypatch.setattr(
-        WorkflowEngine, "_parse_document_text", lambda self, artifact_id: "EXTRACTED"
+        WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None: "EXTRACTED"
     )
     eng = _engine(monkeypatch)
 

tests/agents/tools/test_read_document_tool.py

@@ -8,6 +8,10 @@ validation, and the surfaced action params.
 
 from __future__ import annotations
 
+import logging
+import signal
+import threading
+import time
 import uuid
 from typing import Any, Dict, Optional
 
@@ -22,7 +26,10 @@ _ART_ID = str(uuid.uuid4())
 # ---------------------------------------------------------------------------
 # Run-scoped input resolution (mocks the repo gate)
 # ---------------------------------------------------------------------------
-def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str]):
+def _stub_repo(
+    monkeypatch, *, found: bool, conv: Optional[str], run: Optional[str],
+    size: Optional[int] = None,
+):
     class _Repo:
         def __init__(self, conn):
             pass
@@ -45,6 +52,9 @@ def _stub_repo(monkeypatch, *, found: bool, conv: Optional[str], run: Optional[s
                 return None
             return {"id": artifact_id, "current_version": 1, "title": "statement.pdf"}
 
+        def get_version(self, artifact_id, version):
+            return {"storage_path": "s/1", "mime_type": "application/pdf", "size": size}
+
     class _Conn:
         def __enter__(self):
             return object()
@@ -83,6 +93,7 @@ def _patch_task(monkeypatch, *, payload=None, exc=None):
     def _apply_async(args=None, queue=None, **kw):
         captured["args"] = args
         captured["queue"] = queue
+        captured["limits"] = kw
         result = _FakeAsyncResult(payload=payload, exc=exc)
         captured["result"] = result
         return result
@@ -208,6 +219,42 @@ def test_artifact_ref_sets_last_artifact_id(monkeypatch):
     assert tool.get_artifact_id("read_document") == "new-art"
 
 
+# ---------------------------------------------------------------------------
+# Size-scaled parse window
+# ---------------------------------------------------------------------------
+@pytest.mark.unit
+def test_parse_window_scales_with_the_input_size(monkeypatch):
+    """A large input widens the await AND the task's per-call Celery time limits."""
+    from application.api.user.tasks import parse_task_time_limits, parse_timeout_for_size
+
+    size = 8 * 1024 * 1024
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None, size=size)
+    captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "x"})
+
+    out = _tool().execute_action("read_document", input=_ART_ID, persist=False)
+
+    expected = parse_timeout_for_size(size)
+    assert out["status"] == "ok"
+    assert captured["result"].get_kwargs["timeout"] == expected
+    # The task's import-time soft limit is the BASE timeout, so the per-call limits
+    # must be raised too or the worker kills the parse the caller is still awaiting.
+    assert captured["limits"] == parse_task_time_limits(expected)
+    assert captured["limits"]["soft_time_limit"] > rd.settings.DOCUMENT_PARSE_TIMEOUT
+
+
+@pytest.mark.unit
+def test_parse_window_falls_back_to_the_base_timeout_without_a_size(monkeypatch):
+    """An unknown input size keeps the configured base timeout as the floor."""
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None, size=None)
+    captured = _patch_task(monkeypatch, payload={"status": "ok", "content": "x"})
+
+    _tool().execute_action("read_document", input=_ART_ID, persist=False)
+
+    base = float(rd.settings.DOCUMENT_PARSE_TIMEOUT)
+    assert captured["result"].get_kwargs["timeout"] == base
+    assert captured["limits"]["soft_time_limit"] == int(base)
+
+
 # ---------------------------------------------------------------------------
 # Cross-tenant: rejected BEFORE enqueue
 # ---------------------------------------------------------------------------
@@ -354,3 +401,216 @@ def test_dispatch_enqueues_when_not_in_worker(monkeypatch):
     assert out["status"] == "ok" and out["content"] == "queued"
     assert captured["args"][0] == _ART_ID
     assert captured["args"][1] == {"conversation_id": "conv-1"}
+
+
+# ---------------------------------------------------------------------------
+# The inline (in-worker) parse is time-bounded too
+#
+# The inline branch has no Celery time limit of its own, so an unbounded parse
+# would pin the agent's worker slot until the OUTER task's limit (if any) kills
+# the whole run. It must degrade on the SAME window the dispatch branch awaits.
+# ---------------------------------------------------------------------------
+_TIMED_OUT = "document parsing timed out after"
+
+
+def _inline(monkeypatch, run_parse, *, timeout=0.2) -> ReadDocumentTool:
+    """Drive the inline branch (current_task truthy) with a patched parse window."""
+    _stub_repo(monkeypatch, found=True, conv="conv-1", run=None)
+    monkeypatch.setattr(rd, "current_task", object())
+
+    import application.api.user.tasks as tasks
+    monkeypatch.setattr(
+        tasks.parse_document, "apply_async",
+        lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not enqueue inside a worker")),
+    )
+    monkeypatch.setattr(tasks, "parse_timeout_for_size", lambda size: timeout)
+
+    import application.worker as worker
+    monkeypatch.setattr(worker, "run_parse_document", run_parse)
+    return _tool()
+
+
+@pytest.mark.unit
+def test_inline_parse_within_the_window_returns_its_result(monkeypatch):
+    seen: Dict[str, Any] = {}
+
+    def _run(artifact_id, parent, user_id, options):
+        seen["thread"] = threading.current_thread()
+        seen["armed"] = signal.getitimer(signal.ITIMER_REAL)
+        time.sleep(0.05)
+        return {"status": "ok", "content": "inline", "truncated": False}
+
+    before = signal.getsignal(signal.SIGALRM)
+    tool = _inline(monkeypatch, _run, timeout=0.5)
+    out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    assert out["status"] == "ok" and out["content"] == "inline"
+    # Signal strategy: the parse runs on the calling (main) thread under an armed timer,
+    assert seen["thread"] is threading.main_thread()
+    assert seen["armed"][0] > 0
+    # and the handler + timer are restored even though the parse returned early.
+    assert signal.getsignal(signal.SIGALRM) is before
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)
+
+
+@pytest.mark.unit
+def test_inline_parse_times_out_on_the_signal_path(monkeypatch, caplog):
+    seen: Dict[str, Any] = {}
+
+    def _run(*args, **kwargs):
+        seen["thread"] = threading.current_thread()
+        time.sleep(0.5)
+        seen["finished"] = True
+        return {"status": "ok", "content": "too late"}
+
+    before = signal.getsignal(signal.SIGALRM)
+    tool = _inline(monkeypatch, _run, timeout=0.2)
+    started = time.monotonic()
+    with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"):
+        out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+    elapsed = time.monotonic() - started
+
+    # Same error shape the worker returns on its soft limit.
+    assert out == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."}
+    # SIGALRM actually INTERRUPTS the parse, freeing the slot (it does not merely abandon it).
+    assert elapsed < 0.45
+    assert seen["thread"] is threading.main_thread()
+    assert "finished" not in seen
+    assert signal.getsignal(signal.SIGALRM) is before
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)
+    assert any("timed out" in rec.getMessage() for rec in caplog.records if rec.levelname == "WARNING")
+
+
+@pytest.mark.unit
+def test_inline_parse_times_out_on_the_thread_path(monkeypatch, caplog):
+    """Off the main thread (gevent/threads pools, Windows) the timer is unusable."""
+    seen: Dict[str, Any] = {}
+
+    def _run(*args, **kwargs):
+        seen["thread"] = threading.current_thread()
+        time.sleep(0.5)
+        seen["finished"] = True
+        return {"status": "ok", "content": "too late"}
+
+    tool = _inline(monkeypatch, _run, timeout=0.2)
+    main_handler = signal.getsignal(signal.SIGALRM)
+    box: Dict[str, Any] = {}
+
+    def _call():
+        box["out"] = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"):
+        caller = threading.Thread(target=_call, name="fake-worker-pool-thread")
+        caller.start()
+        caller.join(2.0)
+
+    assert not caller.is_alive()
+    assert box["out"] == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."}
+    # The parse ran in the helper executor, not on the caller's thread...
+    assert seen["thread"].name.startswith("read-document-inline")
+    assert "finished" not in seen
+    # ...and the orphan cannot be interrupted, which the warning must say.
+    assert any("keeps running" in rec.getMessage() for rec in caplog.records if rec.levelname == "WARNING")
+    # The main thread's signal state was never touched from the worker thread.
+    assert signal.getsignal(signal.SIGALRM) is main_handler
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)
+
+
+@pytest.mark.unit
+def test_inline_parse_falls_back_to_a_thread_when_sigalrm_is_taken(monkeypatch):
+    """A foreign SIGALRM handler must not be clobbered; use the thread strategy instead."""
+    seen: Dict[str, Any] = {}
+
+    def _run(*args, **kwargs):
+        seen["thread"] = threading.current_thread()
+        return {"status": "ok", "content": "threaded", "truncated": False}
+
+    def _foreign(signum, frame):  # pragma: no cover - must never be invoked
+        raise AssertionError("foreign SIGALRM handler fired")
+
+    tool = _inline(monkeypatch, _run, timeout=0.5)
+    previous = signal.signal(signal.SIGALRM, _foreign)
+    try:
+        out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+        still_installed = signal.getsignal(signal.SIGALRM)
+    finally:
+        signal.signal(signal.SIGALRM, previous)
+
+    assert out["status"] == "ok" and out["content"] == "threaded"
+    assert seen["thread"].name.startswith("read-document-inline")
+    assert still_installed is _foreign
+
+
+@pytest.mark.unit
+def test_inline_signal_timer_does_not_fire_after_a_fast_parse(monkeypatch):
+    before = signal.getsignal(signal.SIGALRM)
+    tool = _inline(monkeypatch, lambda *a, **k: {"status": "ok", "content": "fast"}, timeout=0.2)
+    out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    assert out["status"] == "ok" and out["content"] == "fast"
+    # A leaked timer would raise inside this sleep, well past the window.
+    time.sleep(0.3)
+    assert signal.getsignal(signal.SIGALRM) is before
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)
+
+
+@pytest.mark.unit
+def test_inline_parse_failure_still_degrades_to_error(monkeypatch):
+    def _boom(*args, **kwargs):
+        raise RuntimeError("docling blew up")
+
+    tool = _inline(monkeypatch, _boom, timeout=0.5)
+    out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    assert out["status"] == "error"
+    assert "document parsing failed" in out["error"] and "docling blew up" in out["error"]
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("window", [0, None])
+def test_inline_parse_runs_unbounded_without_a_positive_window(monkeypatch, window):
+    seen: Dict[str, Any] = {}
+
+    def _run(*args, **kwargs):
+        seen["thread"] = threading.current_thread()
+        seen["armed"] = signal.getitimer(signal.ITIMER_REAL)
+        seen["handler"] = signal.getsignal(signal.SIGALRM)
+        return {"status": "ok", "content": "unbounded", "truncated": False}
+
+    before = signal.getsignal(signal.SIGALRM)
+    tool = _inline(monkeypatch, _run, timeout=window)
+    out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    assert out["status"] == "ok" and out["content"] == "unbounded"
+    # Called directly: no timer armed, no handler swapped, no helper thread.
+    assert seen["armed"] == (0.0, 0.0)
+    assert seen["handler"] is before
+    assert seen["thread"] is threading.main_thread()
+
+
+@pytest.mark.unit
+def test_inline_timeout_is_not_swallowed_by_the_parser_catch_all(monkeypatch, caplog):
+    """``parse_document_bytes`` wraps the parse in ``except Exception`` and returns an error dict.
+
+    The interrupt must fly straight through that catch-all (and still run its ``finally``
+    cleanup) so the caller reports the shared timed-out shape, not "parsing failed: ...".
+    """
+    seen: Dict[str, Any] = {}
+
+    def _run(*args, **kwargs):
+        try:
+            time.sleep(0.5)
+            return {"status": "ok", "content": "too late"}
+        except Exception as exc:  # mirrors parser/document_reader.py's catch-all
+            return {"status": "error", "error": f"parsing failed: {type(exc).__name__}: {exc}"}
+        finally:
+            seen["cleaned_up"] = True
+
+    tool = _inline(monkeypatch, _run, timeout=0.2)
+    with caplog.at_level(logging.WARNING, logger="application.agents.tools.read_document"):
+        out = tool.execute_action("read_document", input=_ART_ID, persist=False)
+
+    assert out == {"status": "error", "error": f"{_TIMED_OUT} {int(0.2)}s."}
+    assert seen["cleaned_up"] is True
+    assert signal.getitimer(signal.ITIMER_REAL) == (0.0, 0.0)

tests/api/user/test_tasks.py

@@ -830,3 +830,67 @@ def test_bare_worker_consumes_app_and_parsing_queues():
     names = {queue.name for queue in celeryconfig.task_queues}
     assert "docsgpt" in names
     assert settings.DOCUMENT_PARSE_QUEUE in names
+
+
+class TestParseTimeoutForSize:
+    """The awaited parse window is floored at the base timeout, scales with size, and is capped."""
+
+    @pytest.mark.unit
+    def test_unknown_size_uses_the_base_timeout(self):
+        from application.api.user.tasks import parse_timeout_for_size
+        from application.core.settings import settings
+
+        base = float(settings.DOCUMENT_PARSE_TIMEOUT)
+        assert parse_timeout_for_size(None) == base
+        assert parse_timeout_for_size(0) == base
+        # A negative/garbage size can never shrink the window below the floor.
+        assert parse_timeout_for_size(-1) == base
+        assert parse_timeout_for_size("nope") == base
+
+    @pytest.mark.unit
+    def test_window_grows_with_the_document_size(self, monkeypatch):
+        from application.api.user import tasks
+        from application.core.settings import settings
+
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False)
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 60, raising=False)
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", 900, raising=False)
+
+        assert tasks.parse_timeout_for_size(1024 * 1024) == 180.0
+        assert tasks.parse_timeout_for_size(5 * 1024 * 1024) == 420.0
+        # Fractional sizes scale proportionally, never rounded down to the floor.
+        assert tasks.parse_timeout_for_size(512 * 1024) == 150.0
+
+    @pytest.mark.unit
+    def test_window_is_capped(self, monkeypatch):
+        from application.api.user import tasks
+        from application.core.settings import settings
+
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False)
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 60, raising=False)
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_MAX", 900, raising=False)
+
+        # A huge document cannot pin a parsing-worker slot indefinitely.
+        assert tasks.parse_timeout_for_size(500 * 1024 * 1024) == 900.0
+
+    @pytest.mark.unit
+    def test_scaling_disabled_by_zero_per_mb(self, monkeypatch):
+        from application.api.user import tasks
+        from application.core.settings import settings
+
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT", 120, raising=False)
+        monkeypatch.setattr(settings, "DOCUMENT_PARSE_TIMEOUT_PER_MB", 0, raising=False)
+
+        assert tasks.parse_timeout_for_size(50 * 1024 * 1024) == 120.0
+
+    @pytest.mark.unit
+    def test_task_time_limits_track_the_awaited_window(self):
+        from application.api.user.tasks import parse_document, parse_task_time_limits
+
+        limits = parse_task_time_limits(420.0)
+        assert limits == {"soft_time_limit": 420, "time_limit": 450}
+        # Same grace as the import-time binding, so per-call limits stay comparable.
+        grace = parse_document.time_limit - parse_document.soft_time_limit
+        assert limits["time_limit"] - limits["soft_time_limit"] == grace
+        # Never zero/negative, whatever the caller computed.
+        assert parse_task_time_limits(0)["soft_time_limit"] == 1

tests/conftest.py

@@ -43,6 +43,7 @@ import os
 # opt back in by setting the env var explicitly.
 os.environ.setdefault("AUTO_MIGRATE", "false")
 os.environ.setdefault("AUTO_CREATE_DB", "false")
+os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false")
 
 import subprocess
 import sys

tests/graphrag/test_extraction.py

@@ -2,8 +2,8 @@
 
 The LLM and the embeddings model are mocked in every test so the suite makes no
 real model or network calls. A live ``GraphStore`` is exercised against the
-pgvector DB with a unique temp ``source_id`` and torn down via
-``delete_by_source``; if no DB is reachable the live tests skip.
+ephemeral pytest-postgresql cluster (never the operator's dev DB) with a unique
+temp ``source_id``; if pgvector is unavailable there the live tests skip.
 """
 
 from __future__ import annotations
@@ -16,54 +16,55 @@ import pytest
 import application.graphrag.extraction as extraction_module
 from application.graphrag.store import GraphStore
 from application.storage.db.source_config import SourceConfig
+from application.vectorstore import pgconn
 
 extract_graph_for_source = extraction_module.extract_graph_for_source
 
 TEST_EMBEDDING_DIM = 8
 
 
-def _resolve_connection_string():
-    from application.core.settings import settings
-
-    conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
-    if not conn and getattr(settings, "POSTGRES_URI", None):
-        from application.core.db_uri import normalize_pgvector_connection_string
-
-        conn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
-    return conn
+@pytest.fixture(autouse=True)
+def _close_pools():
+    """Never leak a pool into another test; an ephemeral DSN dies with its DB."""
+    yield
+    for dsn, pool in list(pgconn._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgconn._POOLS.pop(dsn, None)
 
 
-_GRAPH_TABLES = (
-    "graph_node_chunks",
-    "graph_edges",
-    "graph_nodes",
-    "graph_ingest_progress",
-)
+def _ephemeral_dsn(info) -> str:
+    """libpq DSN for the ephemeral pytest-postgresql database."""
+    password = f":{info.password}" if info.password else ""
+    return (
+        f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+    )
 
 
-def _drop_graph_tables(conn_string):
-    import psycopg
+def _live_store(monkeypatch, info):
+    """Graph store on a fresh ephemeral database, schema created up front.
 
-    with psycopg.connect(conn_string) as conn:
-        with conn.cursor() as cursor:
-            cursor.execute(
-                f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;"
-            )
-        conn.commit()
-
-
-def _live_store(monkeypatch):
+    Construction runs no DDL any more (boot owns the schema), so the tables are
+    created explicitly here — what ``ensure_vector_schema`` does in production.
+    """
     monkeypatch.setattr(
         GraphStore, "_embedding_dim", lambda self: TEST_EMBEDDING_DIM
     )
-    conn = _resolve_connection_string()
-    if not conn:
-        pytest.skip("No pgvector connection string configured")
+    dsn = _ephemeral_dsn(info)
+    # The pipeline builds its own GraphStore() from settings, so point those at
+    # the ephemeral cluster too — never at the operator's configured DB.
+    from application.core import settings as settings_module
+
+    monkeypatch.setattr(
+        settings_module.settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False
+    )
+    store = GraphStore(connection_string=dsn)
     try:
-        _drop_graph_tables(conn)
-        store = GraphStore(connection_string=conn)
+        store._ensure_tables()
     except Exception as exc:
-        pytest.skip(f"pgvector DB not reachable: {exc}")
+        pytest.skip(f"pgvector extension unavailable: {exc}")
     return store
 
 
@@ -102,6 +103,11 @@ class _StubEmbedding:
 
 @pytest.fixture
 def stub_embedding(monkeypatch):
+    from application.core.settings import settings
+
+    # The resolver short-circuits to the remote API when this is configured,
+    # which would bypass the stub on a dev machine that sets it.
+    monkeypatch.setattr(settings, "EMBEDDINGS_BASE_URL", None)
     embedding = _StubEmbedding()
     monkeypatch.setattr(
         extraction_module.EmbeddingsSingleton,
@@ -135,10 +141,10 @@ def _extraction_json(entities, relationships):
 @pytest.mark.integration
 class TestExtractionLive:
     @pytest.fixture
-    def store(self, monkeypatch):
-        store = _live_store(monkeypatch)
+    def store(self, monkeypatch, postgresql):
+        store = _live_store(monkeypatch, postgresql.info)
         yield store
-        _drop_graph_tables(store._connection_string)
+        store.close()
 
     @pytest.fixture
     def source_id(self):
@@ -419,3 +425,37 @@ class TestParsing:
         assert extraction_module._chunk_id({"chunk_id": "abc"}) == "abc"
         assert extraction_module._chunk_id({"id": 9}) == "9"
         assert extraction_module._chunk_id({"text": "no id"}) is None
+
+
+@pytest.mark.unit
+class TestEmbeddingsResolution:
+    def test_extraction_uses_shared_resolver(self, monkeypatch):
+        """Extraction must resolve embeddings through ``get_embeddings``."""
+        from unittest.mock import MagicMock
+
+        fake_store = MagicMock()
+        fake_store.pending_chunks.return_value = []
+        monkeypatch.setattr(
+            "application.graphrag.store.GraphStore", lambda *a, **k: fake_store
+        )
+        _install_stub_llm(monkeypatch, _StubLLM([]))
+
+        calls = []
+        fake_embedding = MagicMock()
+
+        def _resolver(*args, **kwargs):
+            calls.append((args, kwargs))
+            return fake_embedding
+
+        monkeypatch.setattr(extraction_module, "get_embeddings", _resolver)
+
+        summary = extract_graph_for_source(
+            str(uuid.uuid4()),
+            user="owner-1",
+            chunks=[],
+            config=SourceConfig(),
+            request_id="req-1",
+        )
+
+        assert calls == [((), {})]
+        assert summary["chunks_processed"] == 0

tests/graphrag/test_retriever_live.py

@@ -0,0 +1,213 @@
+"""Live GraphRAG retrieval: real graph store, real pool, batched classic fallback.
+
+Runs against the ephemeral pytest-postgresql cluster (never the operator's dev
+database) with an 8-dimension stub embeddings model, so the whole
+``GraphRAGRetriever._get_data`` path — count batching, PPR over a real graph,
+one pooled connection, one ClassicRAG call for the graphless sources — is
+exercised end to end in a few seconds. Skips when the cluster's Postgres has no
+pgvector build installed.
+"""
+
+from __future__ import annotations
+
+import uuid
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+
+from application.retriever.graph_rag import GraphRAGRetriever
+from application.vectorstore import pgconn
+
+pytestmark = pytest.mark.integration
+
+STUB_DIM = 8
+
+
+class _StubEmbeddings:
+    """Deterministic tiny embeddings; no model download, no network."""
+
+    dimension = STUB_DIM
+
+    def embed_documents(self, texts):
+        return [self._vector(1.0) for _ in texts]
+
+    def embed_query(self, query):
+        return self._vector(1.0)
+
+    @staticmethod
+    def _vector(seed: float):
+        vector = [0.0] * STUB_DIM
+        vector[0] = seed
+        return vector
+
+
+def _dsn(info) -> str:
+    password = f":{info.password}" if info.password else ""
+    return (
+        f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+    )
+
+
+@pytest.fixture(autouse=True)
+def _close_pools():
+    """Never leak a pool into another test; the ephemeral DSN dies with its DB."""
+    yield
+    for dsn, pool in list(pgconn._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgconn._POOLS.pop(dsn, None)
+
+
+@pytest.fixture
+def stub_embeddings():
+    stub = _StubEmbeddings()
+    with patch(
+        "application.vectorstore.base.get_embeddings", return_value=stub
+    ), patch(
+        "application.vectorstore.base.BaseVectorStore._get_embeddings",
+        return_value=stub,
+    ), patch(
+        "application.retriever.graph_rag.get_embeddings", return_value=stub
+    ):
+        yield stub
+
+
+@pytest.fixture
+def live_dsn(postgresql, monkeypatch):
+    """DSN of an empty ephemeral DB with GraphRAG enabled and pooling on."""
+    try:
+        with postgresql.cursor() as cursor:
+            cursor.execute("CREATE EXTENSION vector;")
+        postgresql.commit()
+    except Exception as exc:
+        postgresql.rollback()
+        pytest.skip(f"pgvector extension unavailable: {exc}")
+
+    dsn = _dsn(postgresql.info)
+    from application.core import settings as settings_module
+
+    settings = settings_module.settings
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False)
+    return dsn
+
+
+def _seed(dsn, graph_source_id):
+    """Ingest one chunk and one graph node for ``graph_source_id``."""
+    from application.graphrag.store import GraphStore
+    from application.vectorstore.pgvector import PGVectorStore
+
+    vector_store = PGVectorStore(
+        source_id=graph_source_id, connection_string=dsn
+    )
+    try:
+        chunk_ids = vector_store.add_texts(
+            ["Ada Lovelace wrote the first algorithm."],
+            [{"title": "Ada", "source": "/docs/ada.txt"}],
+        )
+    finally:
+        vector_store.close()
+
+    store = GraphStore(connection_string=dsn)
+    try:
+        store._ensure_tables()
+        node_id = store.upsert_node(
+            graph_source_id,
+            "Ada Lovelace",
+            "ada lovelace",
+            "person",
+            "A mathematician.",
+            _StubEmbeddings._vector(1.0),
+        )
+        store.link_node_chunk(graph_source_id, node_id, chunk_ids[0])
+    finally:
+        store.close()
+    return chunk_ids[0]
+
+
+def _retriever(sources):
+    with patch(
+        "application.retriever.classic_rag.LLMCreator.create_llm",
+        Mock(return_value=MagicMock()),
+    ):
+        return GraphRAGRetriever(
+            source={"question": "who wrote the first algorithm?",
+                    "active_docs": list(sources)},
+            chat_history=None,
+            prompt="",
+            chunks=2,
+            doc_token_limit=50000,
+            model_id="test-model",
+            llm_name="openai",
+            api_key="fake",
+            decoded_token={"sub": "user1"},
+        )
+
+
+class TestGraphRAGRetrieverLive:
+    def test_one_pool_one_classic_batch_and_the_connection_goes_back(
+        self, live_dsn, stub_embeddings
+    ):
+        graph_source = str(uuid.uuid4())
+        graphless_source = str(uuid.uuid4())
+        _seed(live_dsn, graph_source)
+
+        rag = _retriever([graphless_source, graph_source])
+        seen = []
+
+        def _classic():
+            seen.append(list(rag._classic.vectorstores))
+            return [
+                {
+                    "title": "classic",
+                    "text": "classic chunk",
+                    "source": graphless_source,
+                    "filename": "classic",
+                }
+            ]
+
+        rag._classic._get_data = Mock(side_effect=_classic)
+
+        docs = rag._get_data()
+
+        # The graphless source is retrieved by exactly one ClassicRAG run.
+        assert rag._classic._get_data.call_count == 1
+        assert seen == [[graphless_source]]
+
+        # The graph source came back through PPR, not the fallback.
+        texts = [doc["text"] for doc in docs]
+        assert texts == ["classic chunk", "Ada Lovelace wrote the first algorithm."]
+        assert docs[1]["source"] == "/docs/ada.txt"
+
+        # One pool for the DSN, shared with the vector store that seeded it,
+        # and the retriever handed its connection back when it finished.
+        assert list(pgconn._POOLS) == [live_dsn]
+        stats = pgconn._POOLS[live_dsn].get_stats()
+        assert stats["pool_size"] <= 4
+        assert stats["pool_available"] == stats["pool_size"]
+
+    def test_all_graphless_sources_take_one_classic_call(
+        self, live_dsn, stub_embeddings
+    ):
+        from application.graphrag.store import GraphStore
+
+        store = GraphStore(connection_string=live_dsn)
+        store._ensure_tables()  # tables exist, but no source has a graph
+        store.close()
+
+        sources = [str(uuid.uuid4()) for _ in range(3)]
+        rag = _retriever(sources)
+        seen = []
+        rag._classic._get_data = Mock(
+            side_effect=lambda: seen.append(list(rag._classic.vectorstores)) or []
+        )
+
+        assert rag._get_data() == []
+        assert seen == [sources]
+        stats = pgconn._POOLS[live_dsn].get_stats()
+        assert stats["pool_available"] == stats["pool_size"]

tests/graphrag/test_store.py

@@ -16,16 +16,20 @@ and the helpers build matching ones.
 from __future__ import annotations
 
 import uuid
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 
 import application.graphrag.store as store_module
+from application.vectorstore import pgconn
+from application.vectorstore import pgvector as pgvector_module
 
 GraphStore = store_module.GraphStore
 
 TEST_EMBEDDING_DIM = 8
 
+POOL_DSN = "postgresql://u:p@localhost/graphpool"
+
 _REAL_EMBEDDING_DIM = GraphStore._embedding_dim
 
 
@@ -36,46 +40,28 @@ def _mock_embedding_dim(monkeypatch):
     )
 
 
-def _resolve_connection_string():
-    from application.core.settings import settings
-
-    conn = getattr(settings, "PGVECTOR_CONNECTION_STRING", None)
-    if not conn and getattr(settings, "POSTGRES_URI", None):
-        from application.core.db_uri import normalize_pgvector_connection_string
-
-        conn = normalize_pgvector_connection_string(settings.POSTGRES_URI)
-    return conn
+@pytest.fixture(autouse=True)
+def _close_pools():
+    """Never leak a pool into another test; an ephemeral DSN dies with its DB."""
+    yield
+    for dsn, pool in list(pgconn._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgconn._POOLS.pop(dsn, None)
 
 
-_GRAPH_TABLES = (
-    "graph_node_chunks",
-    "graph_edges",
-    "graph_nodes",
-    "graph_ingest_progress",
-)
+def _ephemeral_dsn(info) -> str:
+    """libpq DSN for the ephemeral pytest-postgresql database.
 
-
-def _drop_graph_tables(conn_string):
-    import psycopg
-
-    with psycopg.connect(conn_string) as conn:
-        with conn.cursor() as cursor:
-            cursor.execute(
-                f"DROP TABLE IF EXISTS {', '.join(_GRAPH_TABLES)} CASCADE;"
-            )
-        conn.commit()
-
-
-def _live_store():
-    conn = _resolve_connection_string()
-    if not conn:
-        pytest.skip("No pgvector connection string configured")
-    try:
-        _drop_graph_tables(conn)
-        store = GraphStore(connection_string=conn)
-    except Exception as exc:
-        pytest.skip(f"pgvector DB not reachable: {exc}")
-    return store
+    Deliberately not the operator's ``POSTGRES_URI``: these tests create and
+    drop graph tables, and the dev database is not theirs to rewrite.
+    """
+    password = f":{info.password}" if info.password else ""
+    return (
+        f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+    )
 
 
 def _embedding(seed: float) -> list:
@@ -87,10 +73,20 @@ def _embedding(seed: float) -> list:
 @pytest.mark.integration
 class TestGraphStoreLive:
     @pytest.fixture
-    def store(self):
-        store = _live_store()
+    def store(self, postgresql):
+        """Graph store on a fresh ephemeral database.
+
+        Construction no longer creates tables, so the fixture calls
+        ``_ensure_tables`` explicitly — exactly what ``ensure_vector_schema``
+        does at boot in production — and read-before-write tests still pass.
+        """
+        store = GraphStore(connection_string=_ephemeral_dsn(postgresql.info))
+        try:
+            store._ensure_tables()
+        except Exception as exc:
+            pytest.skip(f"pgvector extension unavailable: {exc}")
         yield store
-        _drop_graph_tables(store._connection_string)
+        store.close()
 
     @pytest.fixture
     def source_id(self):
@@ -344,6 +340,61 @@ class TestGraphStoreLive:
         finally:
             store.delete_by_source(source_id)
 
+    def test_count_nodes_many_batches_and_zero_fills(self, store):
+        """One query for N sources; a source with no graph still gets an entry."""
+        a, b, c = (str(uuid.uuid4()) for _ in range(3))
+        try:
+            store.upsert_node(a, "A1", "a1")
+            store.upsert_node(a, "A2", "a2")
+            store.upsert_node(b, "B1", "b1")
+
+            counts = store.count_nodes_many([a, b, c])
+
+            assert counts == {a: 2, b: 1, c: 0}
+            # Agrees with the per-source query it replaces.
+            assert [store.count_nodes(s) for s in (a, b, c)] == [2, 1, 0]
+            assert store.count_nodes_many([]) == {}
+        finally:
+            store.delete_by_source(a)
+            store.delete_by_source(b)
+
+    def test_pooled_connection_is_returned_to_the_shared_pool(self, store):
+        """The live store borrows from the shared pool and gives the socket back."""
+        source_id = str(uuid.uuid4())
+        assert store.count_nodes_many([source_id]) == {source_id: 0}
+        assert store._pooled is True
+        assert list(pgconn._POOLS) == [store._connection_string]
+
+        pool = pgconn._POOLS[store._connection_string]
+        store.close()
+
+        assert store._connection is None
+        stats = pool.get_stats()
+        assert stats["pool_available"] == stats["pool_size"]
+
+    def test_the_vector_store_reuses_the_graph_store_pool(self, store, postgresql):
+        """Same DSN, one pool: the graph store does not double the connections."""
+        from application.vectorstore.pgvector import PGVectorStore
+
+        stub = MagicMock()
+        stub.dimension = TEST_EMBEDDING_DIM
+        stub.embed_query.return_value = [0.0] * TEST_EMBEDDING_DIM
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings",
+            return_value=stub,
+        ):
+            vector_store = PGVectorStore(
+                source_id="live-source", connection_string=store._connection_string
+            )
+        try:
+            store.count_nodes_many([str(uuid.uuid4())])
+            vector_store._get_connection()
+
+            assert list(pgconn._POOLS) == [store._connection_string]
+            assert vector_store._pooled is True
+        finally:
+            vector_store.close()
+
     def test_delete_by_source_isolation(self, store):
         keep = str(uuid.uuid4())
         drop = str(uuid.uuid4())
@@ -381,6 +432,8 @@ class TestGraphStoreParameterization:
         conn.cursor.return_value = cursor
         store._connection = conn
         store._get_connection = lambda: conn
+        # Boot owns the schema; the write-path safety net has its own tests.
+        store._tables_ensured = True
         return store, cursor
 
     def test_delete_by_source_binds_source_id(self):
@@ -448,6 +501,7 @@ class TestEmbeddingDim:
     def test_uses_configured_model_dimension(self, monkeypatch):
         from application.vectorstore import base as base_module
 
+        monkeypatch.setattr(base_module.settings, "EMBEDDINGS_BASE_URL", None)
         fake_embedding = MagicMock()
         fake_embedding.dimension = 1536
         monkeypatch.setattr(
@@ -463,6 +517,7 @@ class TestEmbeddingDim:
     def test_falls_back_to_default_dimension(self, monkeypatch):
         from application.vectorstore import base as base_module
 
+        monkeypatch.setattr(base_module.settings, "EMBEDDINGS_BASE_URL", None)
         fake_embedding = object()
         monkeypatch.setattr(
             base_module.EmbeddingsSingleton,
@@ -473,3 +528,321 @@ class TestEmbeddingDim:
 
         store = GraphStore.__new__(GraphStore)
         assert store._embedding_dim() == store_module.DEFAULT_NAME_EMBEDDING_DIM
+
+
+@pytest.mark.unit
+class TestEmbeddingDimResolution:
+    def test_uses_shared_resolver(self, monkeypatch):
+        """The dimension probe must not build its own embeddings instance."""
+        from unittest.mock import patch
+
+        fake_embedding = MagicMock()
+        fake_embedding.dimension = 1536
+        monkeypatch.setattr(GraphStore, "_embedding_dim", _REAL_EMBEDDING_DIM)
+
+        with patch(
+            "application.vectorstore.base.get_embeddings",
+            return_value=fake_embedding,
+        ) as mock_resolver:
+            store = GraphStore.__new__(GraphStore)
+            assert store._embedding_dim() == 1536
+
+        mock_resolver.assert_called_once_with()
+
+
+@pytest.mark.unit
+class TestGraphSchemaIsBootOwned:
+    """Construction must not run DDL: reads happen once per query, per source."""
+
+    def _mock_store(self):
+        store = GraphStore.__new__(GraphStore)
+        cursor = MagicMock()
+        cursor.fetchone.return_value = [str(uuid.uuid4())]
+        cursor.fetchall.return_value = []
+        conn = MagicMock()
+        conn.cursor.return_value = cursor
+        store._connection = conn
+        store._get_connection = lambda: conn
+        store._tables_ensured = False
+        store._ensure_tables = MagicMock()
+        return store, cursor
+
+    def test_init_opens_no_connection_and_creates_no_tables(self):
+        from unittest.mock import patch
+
+        with patch.dict(
+            "sys.modules",
+            {
+                "psycopg": MagicMock(),
+                "pgvector": MagicMock(),
+                "pgvector.psycopg": MagicMock(),
+            },
+        ), patch.object(GraphStore, "_ensure_tables") as ensure, patch.object(
+            GraphStore, "_get_connection"
+        ) as get_conn:
+            store = GraphStore(connection_string="postgresql://u:p@localhost/db")
+
+        ensure.assert_not_called()
+        get_conn.assert_not_called()
+        assert store._tables_ensured is False
+
+    @pytest.mark.parametrize(
+        "call",
+        [
+            lambda s: s.upsert_node("sid", "N", "n"),
+            lambda s: s.add_edge("sid", "a", "b"),
+            lambda s: s.link_node_chunk("sid", "n", "c1"),
+            lambda s: s.apply_chunk("sid", "c1", [], [], {}),
+            lambda s: s.set_node_degrees("sid"),
+            lambda s: s.mark_chunk("sid", "c1", "done"),
+            lambda s: s.delete_by_source("sid"),
+        ],
+        ids=[
+            "upsert_node",
+            "add_edge",
+            "link_node_chunk",
+            "apply_chunk",
+            "set_node_degrees",
+            "mark_chunk",
+            "delete_by_source",
+        ],
+    )
+    def test_writes_ensure_tables_once(self, call):
+        store, _ = self._mock_store()
+
+        call(store)
+        store._tables_ensured = True  # what the real _ensure_tables_once sets
+        call(store)
+
+        assert store._ensure_tables.call_count == 1
+
+    @pytest.mark.parametrize(
+        "call",
+        [
+            lambda s: s.count_nodes("sid"),
+            lambda s: s.count_nodes_many(["sid"]),
+            lambda s: s.get_node_by_normalized("sid", "n"),
+            lambda s: s.search_nodes_by_embedding("sid", _embedding(1.0)),
+            lambda s: s.get_subgraph("sid", ["n"]),
+            lambda s: s.get_graph_overview("sid"),
+            lambda s: s.get_chunk_ids_for_nodes("sid", ["n"]),
+            lambda s: s.pending_chunks("sid", ["c1"]),
+            lambda s: s.get_progress("sid"),
+        ],
+        ids=[
+            "count_nodes",
+            "count_nodes_many",
+            "get_node_by_normalized",
+            "search_nodes_by_embedding",
+            "get_subgraph",
+            "get_graph_overview",
+            "get_chunk_ids_for_nodes",
+            "pending_chunks",
+            "get_progress",
+        ],
+    )
+    def test_reads_never_create_tables(self, call):
+        store, _ = self._mock_store()
+
+        call(store)
+
+        store._ensure_tables.assert_not_called()
+
+    def test_create_schema_emits_the_ddl_without_committing(self):
+        conn, cursor = MagicMock(), MagicMock()
+        conn.cursor.return_value = cursor
+
+        GraphStore.create_schema(conn, dimension=8)
+
+        statements = " ".join(str(c) for c in cursor.execute.call_args_list)
+        assert "CREATE EXTENSION IF NOT EXISTS vector" in statements
+        for table in (
+            "graph_nodes",
+            "graph_edges",
+            "graph_node_chunks",
+            "graph_ingest_progress",
+        ):
+            assert f"CREATE TABLE IF NOT EXISTS {table}" in statements
+        assert "name_embedding vector(8)" in statements
+        assert statements.count("CREATE INDEX IF NOT EXISTS") == 5
+        conn.commit.assert_not_called()
+
+    def test_ensure_tables_locks_then_commits(self):
+        store, cursor = self._mock_store()
+        del store._ensure_tables  # exercise the real method
+
+        store._ensure_tables()
+
+        statements = " ".join(str(c) for c in cursor.execute.call_args_list)
+        assert "pg_advisory_xact_lock" in statements
+        assert "CREATE TABLE IF NOT EXISTS graph_nodes" in statements
+        store._connection.commit.assert_called_once()
+
+
+@pytest.mark.unit
+class TestGraphStorePooling:
+    """The graph store borrows from the same per-DSN pool as ``PGVectorStore``."""
+
+    def _store(self, dsn=POOL_DSN, pool_max_size=4):
+        store = GraphStore.__new__(GraphStore)
+        store._connection_string = dsn
+        store._connection = None
+        store._pooled = False
+        store._pool_max_size = pool_max_size
+        store._psycopg = MagicMock()
+        store._register_vector = MagicMock()
+        store._tables_ensured = True
+        return store
+
+    def _fake_pool(self):
+        pooled_conn = MagicMock()
+        pooled_conn.closed = False
+        pool = MagicMock()
+        pool.getconn.return_value = pooled_conn
+        return pool, pooled_conn
+
+    def test_get_connection_checks_out_of_the_pool(self, monkeypatch):
+        pool, pooled_conn = self._fake_pool()
+        monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool)
+        store = self._store()
+
+        conn = store._get_connection()
+
+        assert conn is pooled_conn
+        assert store._pooled is True
+        pool.getconn.assert_called_once()
+        store._psycopg.connect.assert_not_called()
+        # The pool's ``configure`` hook already registered the adapters.
+        store._register_vector.assert_not_called()
+
+    def test_close_rolls_back_and_returns_the_connection(self, monkeypatch):
+        pool, pooled_conn = self._fake_pool()
+        monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool)
+        monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool)
+        store = self._store()
+        store._get_connection()
+        pooled_conn.info.transaction_status.name = "INTRANS"
+
+        store.close()
+
+        pooled_conn.rollback.assert_called_once()
+        pool.putconn.assert_called_once_with(pooled_conn)
+        pooled_conn.close.assert_not_called()
+        assert store._connection is None
+
+    def test_close_does_not_roll_back_an_idle_connection(self, monkeypatch):
+        pool, pooled_conn = self._fake_pool()
+        monkeypatch.setattr(store_module.pgconn, "pool_for", lambda dsn, n: pool)
+        monkeypatch.setitem(pgconn._POOLS, POOL_DSN, pool)
+        store = self._store()
+        store._get_connection()
+        pooled_conn.info.transaction_status.name = "IDLE"
+
+        store.close()
+
+        pooled_conn.rollback.assert_not_called()
+        pool.putconn.assert_called_once_with(pooled_conn)
+
+    def test_legacy_path_connects_directly_and_closes(self, monkeypatch):
+        def _never(dsn, n):
+            raise AssertionError("pooling is off; no pool must be built")
+
+        monkeypatch.setattr(store_module.pgconn, "pool_for", _never)
+        store = self._store(pool_max_size=0)
+        direct = MagicMock()
+        direct.closed = False
+        store._psycopg.connect.return_value = direct
+
+        conn = store._get_connection()
+
+        assert conn is direct
+        assert store._pooled is False
+        store._register_vector.assert_called_once_with(direct)
+
+        store.close()
+        direct.close.assert_called_once()
+
+    def test_del_never_raises(self):
+        store = self._store()
+        broken = MagicMock()
+        broken.closed = False
+        broken.close.side_effect = RuntimeError("already gone")
+        store._connection = broken
+
+        store.__del__()  # must not propagate
+
+    def test_the_graph_store_and_the_vector_store_share_one_pool(self):
+        """One DSN, one pool object — reached from either module."""
+        pool, _ = self._fake_pool()
+        store = self._store()
+
+        with patch("psycopg_pool.ConnectionPool", return_value=pool) as pool_cls:
+            store._get_connection()
+            # ``PGVectorStore``'s own entry point resolves to the same object.
+            assert pgvector_module._pool_for(POOL_DSN, 4) is pool
+
+        assert pool_cls.call_count == 1
+        assert pgconn._POOLS[POOL_DSN] is pool
+        assert pgvector_module._POOLS is pgconn._POOLS
+
+    @pytest.mark.parametrize(
+        "value,expected",
+        [(0, 0), (2, 2), (None, 8), ("4", 8), (True, 8), (-1, 8)],
+    )
+    def test_pool_size_is_resolved_defensively(self, monkeypatch, value, expected):
+        monkeypatch.setattr(
+            store_module.settings, "PGVECTOR_POOL_MAX_SIZE", value, raising=False
+        )
+        assert store_module._resolve_pool_max_size() == expected
+
+
+@pytest.mark.unit
+class TestCountNodesMany:
+    """One ``ANY(%s)`` query replaces the retriever's per-source count fan-out."""
+
+    def _store_with_mock_conn(self, rows):
+        store = GraphStore.__new__(GraphStore)
+        cursor = MagicMock()
+        cursor.fetchall.return_value = rows
+        conn = MagicMock()
+        conn.cursor.return_value = cursor
+        store._connection = conn
+        store._get_connection = lambda: conn
+        store._tables_ensured = True
+        return store, cursor, conn
+
+    def test_binds_the_ids_as_one_array_and_zero_fills(self):
+        store, cursor, _ = self._store_with_mock_conn([("a", 2), ("b", 1)])
+
+        counts = store.count_nodes_many(["a", "b", "c"])
+
+        assert counts == {"a": 2, "b": 1, "c": 0}
+        sql, params = (
+            cursor.execute.call_args.args[0],
+            cursor.execute.call_args.args[1],
+        )
+        assert "source_id = ANY(%s)" in sql
+        assert "GROUP BY source_id" in sql
+        assert cursor.execute.call_count == 1
+        assert params == (["a", "b", "c"],)
+
+    def test_empty_input_short_circuits(self):
+        store, cursor, _ = self._store_with_mock_conn([])
+
+        assert store.count_nodes_many([]) == {}
+        assert store.count_nodes_many([None, ""]) == {}
+        cursor.execute.assert_not_called()
+
+    def test_a_failed_query_reports_every_source_as_graphless(self):
+        store, cursor, conn = self._store_with_mock_conn([])
+        cursor.execute.side_effect = RuntimeError("no such table")
+
+        assert store.count_nodes_many(["a", "b"]) == {"a": 0, "b": 0}
+        conn.rollback.assert_called_once()
+
+    def test_the_callers_id_spelling_is_preserved(self):
+        """Postgres returns canonical lowercase UUID text; keys must still match."""
+        source_id = str(uuid.uuid4()).upper()
+        store, _, _ = self._store_with_mock_conn([(source_id.lower(), 3)])
+
+        assert store.count_nodes_many([source_id]) == {source_id: 3}

tests/llm/test_fallback.py

@@ -5,12 +5,18 @@ per-agent backup model is used before the global FALLBACK_* settings.
 """
 
 import copy
+import logging
+import types
 from unittest.mock import MagicMock
 
 import httpx
 import pytest
 
+from application.llm.anthropic import AnthropicLLM
 from application.llm.base import BaseLLM
+from application.llm.google_ai import GoogleLLM
+from application.llm.groq import GroqLLM
+from application.llm.openai import OpenAILLM
 
 
 # Concrete LLM stubs
@@ -41,6 +47,7 @@ class FakeLLM(BaseLLM):
         self.gen_stream_called = False
         self.last_model_received = None  # tracks the model kwarg passed to gen/gen_stream
         self.last_messages_received = None  # tracks the messages kwarg at the raw level
+        self.last_kwargs_received = None  # tracks the extra gen kwargs at the raw level
 
     # Track at the raw-method level. _execute_with_fallback applies
     # decorators to the fallback's raw method directly and
@@ -50,6 +57,7 @@ class FakeLLM(BaseLLM):
         self.gen_called = True
         self.last_model_received = model
         self.last_messages_received = messages
+        self.last_kwargs_received = dict(kwargs)
         if self.fail_at is not None:
             raise RuntimeError("primary model unavailable")
         return self.responses[0]
@@ -59,6 +67,7 @@ class FakeLLM(BaseLLM):
         self.stream_calls = getattr(self, "stream_calls", 0) + 1
         self.last_model_received = model
         self.last_messages_received = messages
+        self.last_kwargs_received = dict(kwargs)
         yielded = 0
         # Per-attempt failure schedule: `fail_schedule[n]` = fail_at value
         # for the n-th call (0-indexed). Falls back to the constant
@@ -1296,3 +1305,489 @@ class TestFallbackMessageReshaping:
         assert isinstance(user_content, str)
         assert "EXTRACTED REPORT TEXT" in user_content
         assert "assistant-abc123" not in user_content
+
+
+# Tests — cross-provider structured-output adaptation
+#
+# Structured output is provider-specific: OpenAI-wire classes take
+# ``response_format``, Google takes ``response_schema``. Forwarding the
+# primary's kwarg verbatim either loses enforcement silently (Google swallows
+# ``response_format`` in ``**kwargs``) or raises TypeError inside the OpenAI
+# SDK (``response_schema`` is not a Chat-Completions param) — which turned the
+# fallback into no fallback at all for every structured node.
+
+
+SCHEMA = {
+    "type": "object",
+    "properties": {
+        "answer": {"type": "string"},
+        "score": {"type": "integer"},
+    },
+    "required": ["answer"],
+}
+
+
+class _OpenAIWireFake(FakeLLM):
+    """OpenAI-wire double: real declaration + real preparer, fake transport."""
+
+    provider_name = "openai"
+    structured_output_kwarg = "response_format"
+    prepare_structured_output_format = OpenAILLM.prepare_structured_output_format
+
+    def _supports_structured_output(self):
+        return True
+
+
+class _GoogleFake(FakeLLM):
+    """Google double: real declaration + real preparer, fake transport."""
+
+    provider_name = "google"
+    structured_output_kwarg = "response_schema"
+    prepare_structured_output_format = GoogleLLM.prepare_structured_output_format
+
+    def _supports_structured_output(self):
+        return True
+
+
+class _AnthropicFake(FakeLLM):
+    """Provider with no structured-output kwarg at all."""
+
+    provider_name = "anthropic"
+
+
+def _openai_envelope(schema=SCHEMA, strict=True):
+    """An OpenAI ``response_format`` built the way a provider would."""
+    return OpenAILLM.prepare_structured_output_format(
+        _OpenAIWireFake(), schema, strict=strict
+    )
+
+
+def _google_schema(schema=SCHEMA):
+    """The Google ``response_schema`` conversion of ``schema``."""
+    return GoogleLLM.prepare_structured_output_format(_GoogleFake(), schema)
+
+
+@pytest.mark.unit
+class TestStructuredOutputDeclarations:
+    """One source of truth: the kwarg name lives on the LLM class."""
+
+    def test_openai_declares_response_format(self):
+        assert OpenAILLM.structured_output_kwarg == "response_format"
+
+    def test_openai_subclasses_inherit_the_declaration(self):
+        assert GroqLLM.structured_output_kwarg == "response_format"
+
+    def test_google_declares_response_schema(self):
+        assert GoogleLLM.structured_output_kwarg == "response_schema"
+
+    def test_anthropic_declares_nothing(self):
+        assert AnthropicLLM.structured_output_kwarg is None
+
+    def test_base_declares_nothing(self):
+        assert BaseLLM.structured_output_kwarg is None
+
+    def test_openai_prepare_records_source(self):
+        llm = _OpenAIWireFake()
+        llm.prepare_structured_output_format(SCHEMA, strict=False)
+        assert llm._structured_output_source == (SCHEMA, False)
+
+    def test_google_prepare_records_source(self):
+        llm = _GoogleFake()
+        llm.prepare_structured_output_format(SCHEMA)
+        assert llm._structured_output_source == (SCHEMA, True)
+
+    def test_empty_schema_clears_recorded_source(self):
+        llm = _OpenAIWireFake()
+        llm.prepare_structured_output_format(SCHEMA)
+        llm.prepare_structured_output_format(None)
+        assert llm._structured_output_source is None
+
+    def test_base_records_nothing_by_default(self):
+        assert FakeLLM()._structured_output_source is None
+
+
+@pytest.mark.unit
+class TestAdaptStructuredOutputKwargs:
+    """Unit-level contract of the adapter itself."""
+
+    def test_no_structured_kwargs_returns_equal_copy(self):
+        primary = _OpenAIWireFake()
+        kwargs = {"model": "m", "messages": [], "temperature": 0.2}
+
+        adapted = primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs)
+
+        assert adapted == kwargs
+        assert adapted is not kwargs
+
+    def test_does_not_mutate_the_callers_kwargs(self):
+        primary = _OpenAIWireFake()
+        primary.prepare_structured_output_format(SCHEMA)
+        kwargs = {"model": "m", "response_format": _openai_envelope()}
+        snapshot = copy.deepcopy(kwargs)
+
+        primary._adapt_structured_output_kwargs(_GoogleFake(), kwargs)
+
+        assert kwargs == snapshot
+
+    def test_recovers_schema_from_envelope_when_no_source_recorded(self):
+        """A hand-built ``response_format`` (research_agent) never went through
+        ``prepare_structured_output_format``, so nothing was recorded — the raw
+        schema is still readable out of the OpenAI envelope."""
+        primary = _OpenAIWireFake()
+        assert primary._structured_output_source is None
+
+        adapted = primary._adapt_structured_output_kwargs(
+            _GoogleFake(model_id="gemini-2.5-flash"),
+            {"model": "m", "response_format": _openai_envelope()},
+        )
+
+        assert "response_format" not in adapted
+        assert adapted["response_schema"]["type"] == "OBJECT"
+        assert set(adapted["response_schema"]["properties"]) == {"answer", "score"}
+        assert adapted["model"] == "m"
+
+    def test_google_schema_without_source_is_dropped(self, caplog):
+        """Google's conversion is lossy/type-mapped — not reversible."""
+        primary = _GoogleFake()
+        assert primary._structured_output_source is None
+
+        with caplog.at_level(logging.WARNING, logger="application.llm.base"):
+            adapted = primary._adapt_structured_output_kwargs(
+                _OpenAIWireFake(model_id="gpt-4o-mini"),
+                {"model": "m", "response_schema": _google_schema()},
+            )
+
+        assert "response_schema" not in adapted
+        assert "response_format" not in adapted
+        assert "gpt-4o-mini" in caplog.text
+
+    def test_fallback_without_structured_support_drops_and_warns(self, caplog):
+        primary = _OpenAIWireFake()
+        primary.prepare_structured_output_format(SCHEMA)
+        fallback = _GoogleFake(model_id="gemini-2.5-flash")
+        fallback._supports_structured_output = lambda: False
+
+        with caplog.at_level(logging.WARNING, logger="application.llm.base"):
+            adapted = primary._adapt_structured_output_kwargs(
+                fallback, {"model": "m", "response_format": _openai_envelope()}
+            )
+
+        assert "response_format" not in adapted
+        assert "response_schema" not in adapted
+        assert "cannot enforce structured output" in caplog.text
+
+    def test_non_callable_support_flag_is_honored(self):
+        """Test doubles sometimes set the capability as a plain bool."""
+        primary = _OpenAIWireFake()
+        primary.prepare_structured_output_format(SCHEMA)
+        fallback = _GoogleFake(model_id="gemini-2.5-flash")
+        fallback._supports_structured_output = False
+
+        adapted = primary._adapt_structured_output_kwargs(
+            fallback, {"response_format": _openai_envelope()}
+        )
+
+        assert adapted == {}
+
+    def test_preparer_returning_none_drops_the_kwarg(self, caplog):
+        class _NullPreparer(_GoogleFake):
+            def prepare_structured_output_format(self, json_schema, strict=True):
+                return None
+
+        primary = _OpenAIWireFake()
+        primary.prepare_structured_output_format(SCHEMA)
+
+        with caplog.at_level(logging.WARNING, logger="application.llm.base"):
+            adapted = primary._adapt_structured_output_kwargs(
+                _NullPreparer(model_id="gemini-2.5-flash"),
+                {"response_format": _openai_envelope()},
+            )
+
+        assert adapted == {}
+        assert "cannot enforce structured output" in caplog.text
+
+    def test_preparer_raising_never_breaks_the_fallback(self, caplog):
+        class _ExplodingPreparer(_GoogleFake):
+            def prepare_structured_output_format(self, json_schema, strict=True):
+                raise ValueError("boom")
+
+        primary = _OpenAIWireFake()
+        primary.prepare_structured_output_format(SCHEMA)
+
+        with caplog.at_level(logging.WARNING, logger="application.llm.base"):
+            adapted = primary._adapt_structured_output_kwargs(
+                _ExplodingPreparer(model_id="gemini-2.5-flash"),
+                {"response_format": _openai_envelope()},
+            )
+
+        assert adapted == {}
+        assert "Failed to prepare structured output" in caplog.text
+
+    def test_strict_flag_survives_the_translation(self):
+        primary = _GoogleFake()
+        primary.prepare_structured_output_format(SCHEMA)
+        primary._structured_output_source = (SCHEMA, False)
+
+        adapted = primary._adapt_structured_output_kwargs(
+            _OpenAIWireFake(model_id="gpt-4o-mini"),
+            {"response_schema": _google_schema()},
+        )
+
+        assert adapted["response_format"]["json_schema"]["strict"] is False
+        # strict=False leaves the schema untouched (no additionalProperties).
+        assert "additionalProperties" not in (
+            adapted["response_format"]["json_schema"]["schema"]
+        )
+
+
+@pytest.mark.integration
+class TestCrossProviderStructuredOutputFallback:
+    """End-to-end through ``gen`` / ``gen_stream``: the backup must receive
+    the schema in *its own* provider's kwarg."""
+
+    def test_gen_openai_primary_google_fallback_gets_response_schema(self):
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
+        primary._fallback_llm = fallback
+        response_format = primary.prepare_structured_output_format(SCHEMA)
+
+        result = primary.gen(**CALL_ARGS, response_format=response_format)
+
+        assert result == "fb"
+        assert fallback.last_kwargs_received["response_schema"] == _google_schema()
+        assert "response_format" not in fallback.last_kwargs_received
+
+    def test_stream_openai_primary_google_fallback_gets_response_schema(self):
+        primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0)
+        fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash")
+        primary._fallback_llm = fallback
+        response_format = primary.prepare_structured_output_format(SCHEMA)
+
+        chunks = list(primary.gen_stream(**CALL_ARGS, response_format=response_format))
+
+        assert chunks == ["fb"]
+        assert fallback.last_kwargs_received["response_schema"] == _google_schema()
+        assert "response_format" not in fallback.last_kwargs_received
+
+    def test_gen_google_primary_openai_fallback_gets_response_format(self):
+        primary = _GoogleFake(fail_at=0)
+        fallback = _OpenAIWireFake(responses=["fb"], model_id="gpt-4o-mini")
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        result = primary.gen(**CALL_ARGS, response_schema=response_schema)
+
+        assert result == "fb"
+        received = fallback.last_kwargs_received
+        assert "response_schema" not in received
+        assert received["response_format"]["type"] == "json_schema"
+        assert set(received["response_format"]["json_schema"]["schema"]["properties"]) == {
+            "answer",
+            "score",
+        }
+
+    def test_stream_google_primary_openai_fallback_gets_response_format(self):
+        primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
+        fallback = _OpenAIWireFake(stream_chunks=["fb"], model_id="gpt-4o-mini")
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
+
+        assert chunks == ["fb"]
+        received = fallback.last_kwargs_received
+        assert "response_schema" not in received
+        assert received["response_format"]["type"] == "json_schema"
+
+    def test_same_wire_family_passes_response_format_verbatim(self):
+        """OpenAI -> openai_compatible: no re-preparation, byte-identical."""
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b")
+        primary._fallback_llm = fallback
+        response_format = primary.prepare_structured_output_format(SCHEMA)
+
+        primary.gen(**CALL_ARGS, response_format=response_format)
+
+        assert fallback.last_kwargs_received["response_format"] is response_format
+        assert "response_schema" not in fallback.last_kwargs_received
+
+    def test_same_wire_family_passes_response_schema_verbatim(self):
+        primary = _GoogleFake(fail_at=0)
+        fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-pro")
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        primary.gen(**CALL_ARGS, response_schema=response_schema)
+
+        assert fallback.last_kwargs_received["response_schema"] is response_schema
+        assert "response_format" not in fallback.last_kwargs_received
+
+    def test_json_object_mode_kept_within_the_openai_family(self):
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _OpenAIWireFake(responses=["fb"], model_id="qwen3-4b")
+        primary._fallback_llm = fallback
+
+        primary.gen(**CALL_ARGS, response_format={"type": "json_object"})
+
+        assert fallback.last_kwargs_received["response_format"] == {
+            "type": "json_object"
+        }
+
+    def test_json_object_mode_dropped_for_google_fallback(self):
+        """Google has no json_object equivalent wired — drop, don't crash."""
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
+        primary._fallback_llm = fallback
+
+        result = primary.gen(**CALL_ARGS, response_format={"type": "json_object"})
+
+        assert result == "fb"
+        assert "response_format" not in fallback.last_kwargs_received
+        assert "response_schema" not in fallback.last_kwargs_received
+
+    def test_json_object_mode_dropped_for_google_fallback_streaming(self):
+        primary = _OpenAIWireFake(stream_chunks=["x"], fail_at=0)
+        fallback = _GoogleFake(stream_chunks=["fb"], model_id="gemini-2.5-flash")
+        primary._fallback_llm = fallback
+
+        chunks = list(
+            primary.gen_stream(**CALL_ARGS, response_format={"type": "json_object"})
+        )
+
+        assert chunks == ["fb"]
+        assert "response_format" not in fallback.last_kwargs_received
+        assert "response_schema" not in fallback.last_kwargs_received
+
+    def test_anthropic_fallback_gets_neither_kwarg(self):
+        """Anthropic has no structured-output kwarg: unstructured, not broken."""
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _AnthropicFake(responses=["fb"], model_id="claude-sonnet-4")
+        primary._fallback_llm = fallback
+        response_format = primary.prepare_structured_output_format(SCHEMA)
+
+        result = primary.gen(**CALL_ARGS, response_format=response_format)
+
+        assert result == "fb"
+        assert fallback.last_kwargs_received == {}
+
+    def test_anthropic_fallback_gets_neither_kwarg_streaming(self):
+        primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
+        fallback = _AnthropicFake(stream_chunks=["fb"], model_id="claude-sonnet-4")
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
+
+        assert chunks == ["fb"]
+        assert fallback.last_kwargs_received == {}
+
+    def test_unrelated_gen_kwargs_are_forwarded_untouched(self):
+        primary = _OpenAIWireFake(fail_at=0)
+        fallback = _GoogleFake(responses=["fb"], model_id="gemini-2.5-flash")
+        primary._fallback_llm = fallback
+        response_format = primary.prepare_structured_output_format(SCHEMA)
+
+        primary.gen(
+            **CALL_ARGS, response_format=response_format, temperature=0.3
+        )
+
+        assert fallback.last_kwargs_received["temperature"] == 0.3
+
+
+# The OpenAI SDK's ``chat.completions.create`` has an explicit keyword
+# signature: a forwarded ``response_schema`` raises TypeError before the
+# request is even built, so the Google -> OpenAI hop died with "Fallback LLM
+# also failed". These tests drive the *real* ``OpenAILLM`` raw methods against
+# a client double with the same strictness.
+
+
+class _StrictChatCompletions:
+    """Chat-Completions double that rejects kwargs the real SDK rejects."""
+
+    _ACCEPTED = {
+        "model",
+        "messages",
+        "stream",
+        "stream_options",
+        "tools",
+        "tool_choice",
+        "parallel_tool_calls",
+        "response_format",
+        "temperature",
+        "top_p",
+        "max_completion_tokens",
+        "reasoning_effort",
+        "presence_penalty",
+        "frequency_penalty",
+        "seed",
+        "stop",
+        "n",
+        "user",
+    }
+
+    def __init__(self):
+        self.last_kwargs = None
+
+    def create(self, **kwargs):
+        unexpected = sorted(set(kwargs) - self._ACCEPTED)
+        if unexpected:
+            raise TypeError(
+                f"Completions.create() got an unexpected keyword argument "
+                f"'{unexpected[0]}'"
+            )
+        self.last_kwargs = kwargs
+        if kwargs.get("stream"):
+            return [
+                _stream_line(content="fb"),
+                _stream_line(finish_reason="stop"),
+            ]
+        message = types.SimpleNamespace(content="fb answer", tool_calls=None)
+        return types.SimpleNamespace(
+            choices=[types.SimpleNamespace(message=message)], usage=None
+        )
+
+
+def _stream_line(content=None, finish_reason=None):
+    delta = types.SimpleNamespace(
+        content=content, reasoning_content=None, tool_calls=None
+    )
+    choice = types.SimpleNamespace(delta=delta, finish_reason=finish_reason)
+    return types.SimpleNamespace(choices=[choice], usage=None)
+
+
+def _strict_openai_llm():
+    llm = OpenAILLM(api_key="sk-test", user_api_key=None, model_id="gpt-4o-mini")
+    llm.client = types.SimpleNamespace(
+        chat=types.SimpleNamespace(completions=_StrictChatCompletions())
+    )
+    return llm
+
+
+@pytest.mark.integration
+class TestRealOpenAIFallbackRejectsForeignKwargs:
+
+    def test_gen_google_primary_real_openai_fallback_does_not_typeerror(self):
+        primary = _GoogleFake(fail_at=0)
+        fallback = _strict_openai_llm()
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        result = primary.gen(**CALL_ARGS, response_schema=response_schema)
+
+        assert result == "fb answer"
+        sent = fallback.client.chat.completions.last_kwargs
+        assert "response_schema" not in sent
+        assert sent["response_format"]["type"] == "json_schema"
+
+    def test_stream_google_primary_real_openai_fallback_does_not_typeerror(self):
+        primary = _GoogleFake(stream_chunks=["x"], fail_at=0)
+        fallback = _strict_openai_llm()
+        primary._fallback_llm = fallback
+        response_schema = primary.prepare_structured_output_format(SCHEMA)
+
+        chunks = list(primary.gen_stream(**CALL_ARGS, response_schema=response_schema))
+
+        assert chunks == ["fb"]
+        sent = fallback.client.chat.completions.last_kwargs
+        assert "response_schema" not in sent
+        assert sent["response_format"]["type"] == "json_schema"

tests/llm/test_openai.py

@@ -15,12 +15,19 @@ Extends coverage beyond test_openai_llm.py:
 """
 
 import base64
+import logging
 import types
 from unittest.mock import MagicMock, patch
 
+import httpx
 import pytest
+from openai import BadRequestError
 
-from application.llm.openai import OpenAILLM, _truncate_base64_for_logging
+from application.llm.openai import (
+    OpenAILLM,
+    _is_tools_unsupported_error,
+    _truncate_base64_for_logging,
+)
 
 
 # Fake client helpers
@@ -2080,3 +2087,235 @@ class TestKeylessConstruction:
             mock_settings.OPENAI_STT_MODEL = "whisper-1"
             stt = OpenAISTT(api_key=blank)
         assert stt.api_key == "sk-no-key"
+
+
+# Tool-calling fallback for endpoints that can't serve tools
+
+
+VLLM_TOOL_ERROR = (
+    '"auto" tool choice requires --enable-auto-tool-choice and '
+    "--tool-call-parser to be set"
+)
+
+
+def _bad_request(message):
+    """Build a real ``openai.BadRequestError`` carrying ``message``."""
+    request = httpx.Request("POST", "http://localhost:8000/v1/chat/completions")
+    response = httpx.Response(400, request=request)
+    return BadRequestError(message, response=response, body={"message": message})
+
+
+class _RecordingCreate:
+    """``create`` stub that raises on the first call, then succeeds."""
+
+    def __init__(self, error, result):
+        self.error = error
+        self.result = result
+        self.calls = []
+
+    def __call__(self, **kwargs):
+        self.calls.append(kwargs)
+        if len(self.calls) == 1 and self.error is not None:
+            raise self.error
+        return self.result
+
+
+@pytest.mark.unit
+class TestIsToolsUnsupportedError:
+
+    @pytest.mark.parametrize(
+        "message",
+        [
+            VLLM_TOOL_ERROR,
+            "Tool calling is not supported by this model",
+            "This endpoint does not support tools",
+            "function calling is unavailable for this deployment",
+            "Unsupported parameter: 'tools'",
+        ],
+    )
+    def test_matches_tool_messages(self, message):
+        assert _is_tools_unsupported_error(_bad_request(message)) is True
+
+    @pytest.mark.parametrize(
+        "message",
+        [
+            "This model's maximum context length is 8192 tokens",
+            "Invalid value for 'temperature': must be <= 2",
+            "Incorrect API key provided",
+        ],
+    )
+    def test_ignores_unrelated_messages(self, message):
+        assert _is_tools_unsupported_error(_bad_request(message)) is False
+
+    def test_matches_message_only_present_in_body(self):
+        request = httpx.Request("POST", "http://localhost:8000/v1/chat/completions")
+        error = BadRequestError(
+            "Error code: 400",
+            response=httpx.Response(400, request=request),
+            body={"message": VLLM_TOOL_ERROR},
+        )
+        assert _is_tools_unsupported_error(error) is True
+
+
+@pytest.mark.unit
+class TestToolCallingFallback:
+    """An OpenAI-compatible endpoint without tool support (vLLM started
+    without --enable-auto-tool-choice) 400s every chat, because agentless
+    chat always sends synthesized default tools. Degrade instead."""
+
+    TOOLS = [{"type": "function", "function": {"name": "t"}}]
+    MSGS = [{"role": "user", "content": "hi"}]
+
+    def test_raw_gen_retries_once_without_tools(self, llm):
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR),
+            _Response(choices=[_Choice(content="degraded answer")]),
+        )
+        llm.client.chat.completions.create = create
+
+        result = llm._raw_gen(
+            llm, model="local-model", messages=self.MSGS, stream=False,
+            tools=self.TOOLS, tool_choice="auto", parallel_tool_calls=True,
+        )
+
+        assert len(create.calls) == 2
+        assert create.calls[0]["tools"] == self.TOOLS
+        assert "tools" not in create.calls[1]
+        assert "tool_choice" not in create.calls[1]
+        assert "parallel_tool_calls" not in create.calls[1]
+        # The caller asked for tools, so the tool-shaped return contract holds.
+        assert result.message.content == "degraded answer"
+
+    def test_raw_gen_stream_retries_once_without_tools(self, llm):
+        retry_response = _Response(
+            lines=[_StreamLine([_Choice(delta="degraded")])]
+        )
+        create = _RecordingCreate(_bad_request(VLLM_TOOL_ERROR), retry_response)
+        llm.client.chat.completions.create = create
+
+        chunks = list(
+            llm._raw_gen_stream(
+                llm, model="local-model", messages=self.MSGS, stream=True,
+                tools=self.TOOLS, tool_choice="auto",
+            )
+        )
+
+        assert chunks == ["degraded"]
+        assert len(create.calls) == 2
+        assert create.calls[0]["tools"] == self.TOOLS
+        assert "tools" not in create.calls[1]
+        assert "tool_choice" not in create.calls[1]
+
+    def test_unrelated_bad_request_propagates(self, llm):
+        error = _bad_request("This model's maximum context length is 8192 tokens")
+        create = _RecordingCreate(error, _Response(choices=[_Choice(content="x")]))
+        llm.client.chat.completions.create = create
+
+        with pytest.raises(BadRequestError):
+            llm._raw_gen(
+                llm, model="local-model", messages=self.MSGS, stream=False,
+                tools=self.TOOLS,
+            )
+
+        assert len(create.calls) == 1
+        assert llm._supports_tools() is True
+
+    def test_unrelated_bad_request_propagates_from_stream(self, llm):
+        error = _bad_request("This model's maximum context length is 8192 tokens")
+        create = _RecordingCreate(error, _Response(lines=[]))
+        llm.client.chat.completions.create = create
+
+        with pytest.raises(BadRequestError):
+            list(
+                llm._raw_gen_stream(
+                    llm, model="local-model", messages=self.MSGS, stream=True,
+                    tools=self.TOOLS,
+                )
+            )
+
+        assert len(create.calls) == 1
+
+    def test_tool_error_without_tools_in_request_propagates(self, llm):
+        """No tools were sent, so there is nothing to degrade to."""
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR), _Response(choices=[_Choice(content="x")])
+        )
+        llm.client.chat.completions.create = create
+
+        with pytest.raises(BadRequestError):
+            llm._raw_gen(
+                llm, model="local-model", messages=self.MSGS, stream=False, tools=None
+            )
+
+        assert len(create.calls) == 1
+        assert llm._supports_tools() is True
+
+    def test_supports_tools_false_after_rejection(self, llm):
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR),
+            _Response(choices=[_Choice(content="degraded")]),
+        )
+        llm.client.chat.completions.create = create
+
+        assert llm._supports_tools() is True
+        llm._raw_gen(
+            llm, model="local-model", messages=self.MSGS, stream=False,
+            tools=self.TOOLS,
+        )
+        assert llm._supports_tools() is False
+
+    def test_later_calls_skip_tools_without_another_round_trip(self, llm):
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR),
+            _Response(choices=[_Choice(content="degraded")]),
+        )
+        llm.client.chat.completions.create = create
+
+        llm._raw_gen(
+            llm, model="local-model", messages=self.MSGS, stream=False,
+            tools=self.TOOLS,
+        )
+        llm._raw_gen(
+            llm, model="local-model", messages=self.MSGS, stream=False,
+            tools=self.TOOLS, tool_choice="auto",
+        )
+
+        # 2 for the first call (reject + retry), 1 for the second.
+        assert len(create.calls) == 3
+        assert "tools" not in create.calls[2]
+        assert "tool_choice" not in create.calls[2]
+
+    def test_rejection_does_not_leak_to_other_instances(self, llm):
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR),
+            _Response(choices=[_Choice(content="degraded")]),
+        )
+        llm.client.chat.completions.create = create
+        llm._raw_gen(
+            llm, model="local-model", messages=self.MSGS, stream=False,
+            tools=self.TOOLS,
+        )
+
+        other = OpenAILLM(api_key="sk-test", user_api_key=None)
+        assert other._supports_tools() is True
+
+    def test_warns_once_with_actionable_message(self, llm, caplog):
+        create = _RecordingCreate(
+            _bad_request(VLLM_TOOL_ERROR),
+            _Response(choices=[_Choice(content="degraded")]),
+        )
+        llm.client.chat.completions.create = create
+
+        with caplog.at_level(logging.WARNING):
+            llm._raw_gen(
+                llm, model="local-model", messages=self.MSGS, stream=False,
+                tools=self.TOOLS,
+            )
+            llm._note_tools_rejected("local-model", "second time")
+
+        warnings = [r for r in caplog.records if r.levelno == logging.WARNING]
+        assert len(warnings) == 1
+        message = warnings[0].getMessage()
+        assert "supports_tools: false" in message
+        assert "local-model" in message
+        assert "--enable-auto-tool-choice" in message

tests/llm/test_openai_responses.py

@@ -1471,3 +1471,82 @@ def test_to_responses_input_drops_out_of_order_pair_entirely(monkeypatch):
     items = llm._to_responses_input(messages)
     assert not [i for i in items if i.get("type") == "function_call"]
     assert not [i for i in items if i.get("type") == "function_call_output"]
+
+
+# ── tool-calling fallback on the Responses path ──────────────────────────────
+
+
+def _vllm_bad_request():
+    """400 a tool-less OpenAI-compatible endpoint returns for a tools request."""
+    import httpx
+    from openai import BadRequestError
+
+    message = (
+        '"auto" tool choice requires --enable-auto-tool-choice and '
+        "--tool-call-parser to be set"
+    )
+    request = httpx.Request("POST", "http://localhost:8000/v1/responses")
+    return BadRequestError(
+        message, response=httpx.Response(400, request=request), body={"message": message}
+    )
+
+
+class _RejectThenAccept:
+    def __init__(self, result):
+        self.result = result
+        self.calls = []
+
+    def __call__(self, **kwargs):
+        self.calls.append(kwargs)
+        if len(self.calls) == 1:
+            raise _vllm_bad_request()
+        return self.result
+
+
+@pytest.mark.unit
+def test_responses_gen_retries_without_tools(monkeypatch):
+    llm = _make_llm(monkeypatch, _responses_caps())
+    response = _ns(
+        id="resp_degraded",
+        output=[_ns(type="message", content=[_ns(type="output_text", text="plain")])],
+    )
+    create = _RejectThenAccept(response)
+    llm.client.responses.create = create
+
+    choice = llm._responses_gen(
+        "local-model",
+        [{"role": "user", "content": "hi"}],
+        tools=[{"type": "function", "function": {"name": "t", "parameters": {}}}],
+        tool_choice="auto",
+    )
+
+    assert len(create.calls) == 2
+    assert "tools" in create.calls[0]
+    assert "tools" not in create.calls[1]
+    assert "tool_choice" not in create.calls[1]
+    assert choice.message.content == "plain"
+    assert llm._supports_tools() is False
+
+
+@pytest.mark.unit
+def test_responses_gen_stream_retries_without_tools(monkeypatch):
+    llm = _make_llm(monkeypatch, _responses_caps())
+    events = [
+        _ns(type="response.output_text.delta", delta="plain"),
+        _ns(type="response.completed", response=_ns(id="resp_degraded")),
+    ]
+    create = _RejectThenAccept(events)
+    llm.client.responses.create = create
+
+    chunks = list(
+        llm._responses_gen_stream(
+            "local-model",
+            [{"role": "user", "content": "hi"}],
+            tools=[{"type": "function", "function": {"name": "t", "parameters": {}}}],
+        )
+    )
+
+    assert chunks == ["plain"]
+    assert len(create.calls) == 2
+    assert "tools" in create.calls[0]
+    assert "tools" not in create.calls[1]

tests/parser/file/test_docling_parser.py

@@ -259,7 +259,9 @@ class TestDoclingParserParseFile:
         mock_converter = MagicMock()
         mock_result = MagicMock()
         mock_doc = MagicMock()
-        mock_doc.export_to_markdown.return_value = "content"
+        # Long enough to clear the OCR chars-per-page floor, which a default
+        # (ocr_enabled=True) parser applies to PDFs.
+        mock_doc.export_to_markdown.return_value = "converted document content"
         mock_doc.texts = []
         mock_result.document = mock_doc
         mock_converter.convert.return_value = mock_result
@@ -268,7 +270,7 @@ class TestDoclingParserParseFile:
             parser._converter = mock_converter
             mock_init.return_value = {}
             result = parser.parse_file(Path("test.pdf"))
-            assert "content" in result
+            assert "converted document content" in result
 
     def test_parse_file_error_ignore_raises_instead_of_returning_error_text(self):
         """A failed conversion must never become the document's text.
@@ -1011,3 +1013,352 @@ class TestApplyInferenceSettings:
             "torch.compile must be disabled before PdfPipelineOptions is "
             f"constructed, got {events}"
         )
+
+
+# =====================================================================
+# OCR dropout guard
+# =====================================================================
+
+
+def _mock_conversion(markdown: str, pages: int = 1) -> MagicMock:
+    """Build a mock docling ConversionResult exporting ``markdown``."""
+    document = MagicMock()
+    document.export_to_markdown.return_value = markdown
+    document.texts = []
+    document.pages = {index: object() for index in range(pages)}
+    result = MagicMock()
+    result.document = document
+    return result
+
+
+def _set_threshold(monkeypatch, value: int) -> None:
+    """Point the OCR dropout guard at a specific chars-per-page floor."""
+    from application.core.settings import settings as real_settings
+
+    class _Stub:
+        def __getattr__(self, name):
+            return getattr(real_settings, name)
+
+    stub = _Stub()
+    stub.DOCLING_OCR_MIN_CHARS_PER_PAGE = value
+    monkeypatch.setattr("application.core.settings.settings", stub)
+
+
+@pytest.mark.unit
+class TestOCRDropoutGuard:
+    """A near-empty OCR parse must be loud, never stored as the document.
+
+    In production, after a long scanned PDF, docling's pipeline degraded
+    inside the worker and returned zero characters for every later scanned
+    page — with no error. ``parse_file`` returned that empty string and
+    ingestion reported success, so the source was indexed as empty.
+    """
+
+    def test_near_empty_first_pass_recovers_on_fresh_converter(self, caplog):
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        parser = DoclingPDFParser(ocr_enabled=True)
+        degraded = MagicMock()
+        degraded.convert.return_value = _mock_conversion(
+            "<!-- image -->\n\n<!-- image -->\n\n<!-- image -->", pages=3
+        )
+        parser._converter = degraded
+
+        fresh = MagicMock()
+        fresh.convert.return_value = _mock_conversion(
+            "Recovered scanned text for page one, page two and page three.",
+            pages=3,
+        )
+        forced_during_retry = []
+
+        def _create_converter():
+            forced_during_retry.append(parser.force_full_page_ocr)
+            return fresh
+
+        with patch.object(
+            parser, "_create_converter", side_effect=_create_converter
+        ) as create, caplog.at_level("WARNING"):
+            result = parser.parse_file(Path("scan.pdf"))
+
+        assert "Recovered scanned text" in result
+        # Exactly one retry, on a converter that did not exist before.
+        assert create.call_count == 1
+        assert fresh.convert.call_count == 1
+        assert degraded.convert.call_count == 1
+        # The retry forces full-page OCR...
+        assert forced_during_retry == [True]
+        # ...but that must not stick to the parser or its converter, or every
+        # later file in the worker pays for full-page OCR it does not need.
+        assert parser.force_full_page_ocr is False
+        assert parser._converter is None
+        messages = " ".join(caplog.messages)
+        assert "scan.pdf" in messages
+        assert "near-empty" in messages
+        assert "Recovered scan.pdf on retry" in messages
+
+    def test_near_empty_both_passes_raises_document_parse_error(self):
+        from application.parser.file.base_parser import DocumentParseError
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        parser = DoclingPDFParser(ocr_enabled=True)
+        degraded = MagicMock()
+        degraded.convert.return_value = _mock_conversion("<!-- image -->", pages=100)
+        parser._converter = degraded
+
+        fresh = MagicMock()
+        fresh.convert.return_value = _mock_conversion("", pages=100)
+
+        with patch.object(parser, "_create_converter", return_value=fresh):
+            with pytest.raises(DocumentParseError) as excinfo:
+                parser.parse_file(Path("scan.pdf"))
+
+        message = str(excinfo.value)
+        assert "scan.pdf" in message
+        assert "OCR produced 0 chars over 100 pages" in message
+        assert "OCR pipeline dropout" in message
+        assert "not indexed" in message
+        # The dropout diagnosis must survive: the generic handler must not
+        # re-wrap it into "Failed to parse ... with docling".
+        assert "Failed to parse" not in message
+        assert parser._converter is None
+
+    def test_dropout_error_reports_an_unread_text_layer(self, monkeypatch):
+        """A text layer docling ignored points at the pipeline, not the scan."""
+        from application.parser.file.base_parser import DocumentParseError
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        monkeypatch.setattr(
+            "application.parser.file.docling_parser._pdf_text_layer_probe",
+            lambda file: (12, 48_000),
+        )
+        parser = DoclingPDFParser(ocr_enabled=True)
+        parser._converter = MagicMock()
+        parser._converter.convert.return_value = _mock_conversion("", pages=12)
+
+        with patch.object(parser, "_create_converter", return_value=parser._converter):
+            with pytest.raises(DocumentParseError) as excinfo:
+                parser.parse_file(Path("scan.pdf"))
+
+        assert "The PDF carries a 48000-char text layer" in str(excinfo.value)
+
+    def test_ocr_disabled_returns_near_empty_content_untouched(self):
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        parser = DoclingPDFParser(ocr_enabled=False)
+        converter = MagicMock()
+        converter.convert.return_value = _mock_conversion("tiny", pages=40)
+        parser._converter = converter
+
+        with patch.object(parser, "_create_converter") as create:
+            assert parser.parse_file(Path("text.pdf")) == "tiny"
+
+        assert create.call_count == 0
+        assert converter.convert.call_count == 1
+        assert parser._converter is converter
+
+    def test_threshold_zero_disables_the_guard(self, monkeypatch):
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        _set_threshold(monkeypatch, 0)
+        parser = DoclingPDFParser(ocr_enabled=True)
+        converter = MagicMock()
+        converter.convert.return_value = _mock_conversion("", pages=100)
+        parser._converter = converter
+
+        with patch.object(parser, "_create_converter") as create:
+            assert parser.parse_file(Path("scan.pdf")) == ""
+
+        assert create.call_count == 0
+        assert parser._converter is converter
+
+    def test_healthy_ocr_parse_converts_once(self):
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        parser = DoclingPDFParser(ocr_enabled=True)
+        converter = MagicMock()
+        converter.convert.return_value = _mock_conversion(
+            "A page of genuinely scanned text. " * 10, pages=2
+        )
+        parser._converter = converter
+
+        with patch.object(parser, "_create_converter") as create:
+            result = parser.parse_file(Path("scan.pdf"))
+
+        assert "genuinely scanned text" in result
+        assert converter.convert.call_count == 1
+        assert create.call_count == 0
+        assert parser._converter is converter
+
+    def test_threshold_is_per_page_not_per_document(self, monkeypatch):
+        """200 chars is healthy for one page and a dropout for a hundred."""
+        from application.parser.file.base_parser import DocumentParseError
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        content = "x" * 200
+
+        parser = DoclingPDFParser(ocr_enabled=True)
+        parser._converter = MagicMock()
+        parser._converter.convert.return_value = _mock_conversion(content, pages=1)
+        with patch.object(parser, "_create_converter") as create:
+            assert parser.parse_file(Path("one-page.pdf")) == content
+        assert create.call_count == 0
+
+        long_doc = DoclingPDFParser(ocr_enabled=True)
+        long_doc._converter = MagicMock()
+        long_doc._converter.convert.return_value = _mock_conversion(content, pages=100)
+        with patch.object(
+            long_doc, "_create_converter", return_value=long_doc._converter
+        ):
+            with pytest.raises(DocumentParseError):
+                long_doc.parse_file(Path("hundred-page.pdf"))
+
+    def test_guard_does_not_apply_to_non_ocr_formats(self):
+        """DOCX parsers inherit ocr_enabled=True but never OCR anything."""
+        from application.parser.file.docling_parser import DoclingDocxParser
+
+        parser = DoclingDocxParser()
+        assert parser.ocr_enabled is True
+        converter = MagicMock()
+        converter.convert.return_value = _mock_conversion("Hi", pages=1)
+        parser._converter = converter
+
+        with patch.object(parser, "_create_converter") as create:
+            assert parser.parse_file(Path("memo.docx")) == "Hi"
+
+        assert create.call_count == 0
+
+    def test_guard_applies_to_images(self):
+        from application.parser.file.base_parser import DocumentParseError
+        from application.parser.file.docling_parser import DoclingImageParser
+
+        parser = DoclingImageParser(ocr_enabled=True)
+        converter = MagicMock()
+        converter.convert.return_value = _mock_conversion("<!-- image -->")
+        parser._converter = converter
+
+        with patch.object(parser, "_create_converter", return_value=converter):
+            with pytest.raises(DocumentParseError, match="scan.png"):
+                parser.parse_file(Path("scan.png"))
+
+    def test_retry_conversion_failure_is_still_a_parse_error(self):
+        from application.parser.file.base_parser import DocumentParseError
+        from application.parser.file.docling_parser import DoclingPDFParser
+
+        parser = DoclingPDFParser(ocr_enabled=True)
+        parser._converter = MagicMock()
+        parser._converter.convert.return_value = _mock_conversion("", pages=2)
+
+        fresh = MagicMock()
+        fresh.convert.side_effect = RuntimeError("OCR model unavailable")
+
+        with patch.object(parser, "_create_converter", return_value=fresh):
+            with pytest.raises(DocumentParseError, match="OCR model unavailable"):
+                parser.parse_file(Path("scan.pdf"))
+
+        assert parser.force_full_page_ocr is False
+        assert parser._converter is None
+
+
+@pytest.mark.unit
+class TestOCRDropoutHelpers:
+
+    def test_image_placeholders_are_not_text(self):
+        from application.parser.file.docling_parser import _text_char_count
+
+        assert _text_char_count(None) == 0
+        assert _text_char_count("") == 0
+        assert _text_char_count("<!-- image -->\n\n<!-- image -->") == 0
+        assert _text_char_count("  <!-- image -->  abc  ") == 3
+
+    def test_page_count_prefers_document_pages(self):
+        from application.parser.file.docling_parser import _result_page_count
+
+        result = _mock_conversion("x", pages=7)
+        assert _result_page_count(result, Path("a.pdf")) == 7
+
+    def test_page_count_falls_back_to_num_pages(self):
+        from application.parser.file.docling_parser import _result_page_count
+
+        result = _mock_conversion("x", pages=0)
+        result.document.num_pages.return_value = 5
+        assert _result_page_count(result, Path("a.pdf")) == 5
+
+    def test_page_count_falls_back_to_result_pages(self):
+        from application.parser.file.docling_parser import _result_page_count
+
+        result = _mock_conversion("x", pages=0)
+        result.document.num_pages.return_value = None
+        result.pages = [object(), object()]
+        assert _result_page_count(result, Path("a.pdf")) == 2
+
+    def test_page_count_defaults_to_one(self):
+        from application.parser.file.docling_parser import _result_page_count
+
+        result = _mock_conversion("x", pages=0)
+        result.document.num_pages.return_value = None
+        result.pages = []
+        assert _result_page_count(result, Path("a.pdf")) == 1
+
+    def test_images_are_always_one_page(self):
+        from application.parser.file.docling_parser import _result_page_count
+
+        result = _mock_conversion("x", pages=9)
+        assert _result_page_count(result, Path("a.png")) == 1
+
+    def test_threshold_reads_settings_with_a_default(self, monkeypatch):
+        from application.parser.file.docling_parser import _ocr_min_chars_per_page
+
+        assert _ocr_min_chars_per_page() == 20
+        _set_threshold(monkeypatch, 5)
+        assert _ocr_min_chars_per_page() == 5
+        _set_threshold(monkeypatch, "nonsense")
+        assert _ocr_min_chars_per_page() == 20
+
+    def test_text_layer_probe_survives_an_unreadable_pdf(self, tmp_path):
+        from application.parser.file.docling_parser import _pdf_text_layer_probe
+
+        assert _pdf_text_layer_probe(tmp_path / "missing.pdf") == (0, 0)
+        broken = tmp_path / "broken.pdf"
+        broken.write_bytes(b"not a pdf")
+        assert _pdf_text_layer_probe(broken) == (0, 0)
+
+
+@pytest.mark.unit
+class TestForceFullPageOCRWiring:
+    """``force_full_page_ocr`` must reach the pipeline for any OCR engine."""
+
+    def _pipeline_options(self, monkeypatch, **kwargs):
+        import docling.datamodel.pipeline_options as dpo
+
+        from application.parser.file.docling_parser import DoclingParser
+
+        built = []
+        real_options = dpo.PdfPipelineOptions
+
+        def _tracking_options(*args, **opts):
+            options = real_options(*args, **opts)
+            built.append(options)
+            return options
+
+        monkeypatch.setattr(dpo, "PdfPipelineOptions", _tracking_options)
+        monkeypatch.setattr("docling.document_converter.DocumentConverter", MagicMock())
+        DoclingParser(**kwargs)._create_converter()
+        return built[0]
+
+    def test_forced_without_rapidocr(self, monkeypatch):
+        options = self._pipeline_options(
+            monkeypatch,
+            ocr_enabled=True,
+            use_rapidocr=False,
+            force_full_page_ocr=True,
+        )
+        assert options.ocr_options.force_full_page_ocr is True
+
+    def test_not_forced_by_default(self, monkeypatch):
+        options = self._pipeline_options(
+            monkeypatch,
+            ocr_enabled=True,
+            use_rapidocr=False,
+            force_full_page_ocr=False,
+        )
+        assert options.ocr_options.force_full_page_ocr is False

tests/parser/test_chunking_strategies.py

@@ -140,6 +140,13 @@ class _FakeEmbeddings:
 
 @pytest.mark.unit
 class TestSemantic:
+    @pytest.fixture(autouse=True)
+    def _no_remote_embeddings(self, monkeypatch):
+        """The resolver short-circuits to the remote API when this is set."""
+        from application.core.settings import settings
+
+        monkeypatch.setattr(settings, "EMBEDDINGS_BASE_URL", None)
+
     def test_breakpoint_forces_split(self):
         # Two topics: sentences 0-1 vs 2-3, orthogonal embeddings between.
         text = "Alpha one. Alpha two. Beta one. Beta two."
@@ -250,3 +257,25 @@ class TestClassicByteIdentical:
         assert [(c.doc_id, c.text, c.extra_info) for c in via] == [
             (c.doc_id, c.text, c.extra_info) for c in direct
         ]
+
+
+@pytest.mark.unit
+class TestSemanticEmbeddingsResolution:
+    def test_uses_shared_resolver(self):
+        """Semantic chunking must resolve embeddings through ``get_embeddings``.
+
+        Calling the singleton directly with the configured name skips the
+        bundled local model and downloads a second copy from the hub.
+        """
+        vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]]
+        text = "Alpha one. Alpha two. Beta one. Beta two."
+        chunker = SemanticChunker(max_tokens=2000, min_tokens=0)
+
+        with patch(
+            "application.vectorstore.base.get_embeddings",
+            return_value=_FakeEmbeddings(vectors),
+        ) as mock_resolver:
+            out = chunker.chunk([Document(text=text, doc_id="d")])
+
+        mock_resolver.assert_called_once_with()
+        assert len(out) == 2

tests/retriever/test_fanout.py

@@ -0,0 +1,176 @@
+"""Unit tests for the shared multi-source fan-out helpers.
+
+``application/retriever/fanout.py`` holds the pieces both ClassicRAG and the
+search service reuse: the worker cap, the single query embedding, and the
+order-preserving pool runner. These tests pin them directly, independent of
+either caller.
+"""
+
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+
+from application.retriever.fanout import (
+    DEFAULT_MAX_PARALLEL_SOURCES,
+    embed_questions,
+    max_parallel_sources,
+    run_source_jobs,
+    store_embeddings,
+)
+
+
+@pytest.mark.unit
+class TestMaxParallelSources:
+    def _cap(self, n_sources, configured=None):
+        stub = SimpleNamespace()
+        if configured is not None:
+            stub.RETRIEVAL_MAX_PARALLEL_SOURCES = configured
+        return max_parallel_sources(n_sources, stub)
+
+    def test_default_when_setting_missing(self):
+        assert self._cap(10) == DEFAULT_MAX_PARALLEL_SOURCES
+
+    def test_never_more_workers_than_sources(self):
+        assert self._cap(2) == 2
+
+    def test_setting_is_honoured(self):
+        assert self._cap(10, configured=8) == 8
+
+    def test_floor_is_one(self):
+        assert self._cap(10, configured=0) == 1
+        assert self._cap(10, configured=-5) == 1
+
+    def test_garbage_setting_falls_back_to_default(self):
+        assert self._cap(10, configured="lots") == DEFAULT_MAX_PARALLEL_SOURCES
+        assert self._cap(10, configured=None) == DEFAULT_MAX_PARALLEL_SOURCES
+
+    def test_zero_sources_still_returns_one(self):
+        assert self._cap(0) == 1
+
+    def test_falls_back_to_module_settings(self, monkeypatch):
+        import application.retriever.fanout as fanout
+
+        monkeypatch.setattr(
+            fanout, "settings", SimpleNamespace(RETRIEVAL_MAX_PARALLEL_SOURCES=3)
+        )
+        assert fanout.max_parallel_sources(10) == 3
+
+
+@pytest.mark.unit
+class TestRunSourceJobs:
+    def test_empty_jobs_returns_empty(self):
+        assert run_source_jobs(Mock(side_effect=AssertionError), []) == []
+
+    def test_preserves_job_order(self):
+        import time
+
+        delays = {"a": 0.06, "b": 0.0, "c": 0.0}
+
+        def _fn(job):
+            time.sleep(delays[job])
+            return job.upper()
+
+        assert run_source_jobs(_fn, ["a", "b", "c"], workers=3) == ["A", "B", "C"]
+
+    def test_jobs_run_concurrently(self):
+        """The barrier only clears if three jobs overlap; a serial loop hangs."""
+        barrier = threading.Barrier(3, timeout=10)
+
+        def _fn(job):
+            barrier.wait()
+            return job
+
+        assert run_source_jobs(_fn, [1, 2, 3], workers=3) == [1, 2, 3]
+
+    def test_single_job_runs_inline(self):
+        main_thread = threading.current_thread().name
+        seen = []
+
+        def _fn(job):
+            seen.append(threading.current_thread().name)
+            return job
+
+        assert run_source_jobs(_fn, ["only"]) == ["only"]
+        assert seen == [main_thread]
+
+    def test_one_worker_runs_inline(self):
+        main_thread = threading.current_thread().name
+        seen = []
+
+        def _fn(job):
+            seen.append(threading.current_thread().name)
+            return job
+
+        run_source_jobs(_fn, ["a", "b"], workers=1)
+        assert seen == [main_thread, main_thread]
+
+    def test_worker_count_defaults_to_the_cap(self):
+        with patch(
+            "application.retriever.fanout.max_parallel_sources", return_value=1
+        ) as cap:
+            run_source_jobs(lambda job: job, ["a", "b", "c"])
+        cap.assert_called_once_with(3)
+
+    def test_accepts_any_iterable_of_jobs(self):
+        assert run_source_jobs(lambda job: job * 2, iter([1, 2])) == [2, 4]
+
+
+@pytest.mark.unit
+class TestStoreEmbeddings:
+    def test_finds_private_embedding_attribute(self):
+        embedder = Mock()
+        embedder.embed_query = Mock(return_value=[1.0])
+        store = SimpleNamespace(_embedding=embedder)
+        assert store_embeddings(store) is embedder
+
+    def test_probes_attributes_in_order(self):
+        first = Mock()
+        first.embed_query = Mock()
+        second = Mock()
+        second.embed_query = Mock()
+        store = SimpleNamespace(_embedding=first, _embeddings=second)
+        assert store_embeddings(store) is first
+
+    def test_skips_attributes_without_embed_query(self):
+        embedder = Mock(spec=["embed_query"])
+        store = SimpleNamespace(_embedding=object(), embeddings=embedder)
+        assert store_embeddings(store) is embedder
+
+    def test_falls_back_to_get_embeddings(self):
+        embedder = Mock()
+        store = SimpleNamespace(
+            _embedding=None, _get_embeddings=Mock(return_value=embedder)
+        )
+        assert store_embeddings(store) is embedder
+
+    def test_returns_none_when_nothing_is_reachable(self):
+        store = SimpleNamespace(_embedding=None)
+        assert store_embeddings(store) is None
+
+    def test_get_embeddings_failure_returns_none(self):
+        store = SimpleNamespace(
+            _embedding=None, _get_embeddings=Mock(side_effect=RuntimeError("boom"))
+        )
+        assert store_embeddings(store) is None
+
+
+@pytest.mark.unit
+class TestEmbedQuestions:
+    def test_embeds_each_distinct_question_once(self):
+        embedder = Mock()
+        embedder.embed_query = Mock(side_effect=lambda q: [len(q)])
+        store = SimpleNamespace(_embedding=embedder)
+
+        assert embed_questions(store, ["ab", "cde"]) == {"ab": [2], "cde": [3]}
+        assert embedder.embed_query.call_count == 2
+
+    def test_no_embedder_returns_empty_map(self):
+        assert embed_questions(SimpleNamespace(_embedding=None), ["q"]) == {}
+
+    def test_embedding_failure_returns_empty_map(self):
+        embedder = Mock()
+        embedder.embed_query = Mock(side_effect=RuntimeError("model down"))
+        store = SimpleNamespace(_embedding=embedder)
+        assert embed_questions(store, ["q"]) == {}

tests/retriever/test_graph_rag.py

@@ -56,7 +56,7 @@ class TestGraphRAGFallback:
         self, _avail, mock_store_cls, _patch_llm_creator
     ):
         store = MagicMock()
-        store.count_nodes.return_value = 0
+        store.count_nodes_many.return_value = {"src1": 0}
         mock_store_cls.return_value = store
 
         rag = _make_retriever()
@@ -68,6 +68,7 @@ class TestGraphRAGFallback:
         assert docs == classic_docs
         store.search_nodes_by_embedding.assert_not_called()
         store.get_subgraph.assert_not_called()
+        store.close.assert_called_once()
 
     @patch("application.retriever.graph_rag.GraphStore")
     @patch("application.retriever.graph_rag.graphrag_available", return_value=False)
@@ -101,6 +102,9 @@ def _store_with_graph(
 ):
     store = MagicMock()
     store.count_nodes.return_value = len(nodes)
+    store.count_nodes_many.side_effect = lambda ids: {
+        source_id: len(nodes) for source_id in ids
+    }
     store.search_nodes_by_embedding.return_value = seed_rows
     store.get_subgraph.return_value = {"nodes": nodes, "edges": edges}
     store.get_chunk_ids_for_nodes.return_value = node_chunks
@@ -166,7 +170,7 @@ class TestGraphRAGHappyPath:
         rag = _make_retriever(chunks=2)
         # Call the PPR path directly: _get_data would swallow a raise and fall
         # back to ClassicRAG, hiding the regression.
-        docs = rag._graph_docs_for_source(store, "src1")
+        docs = rag._graph_docs_for_source(store, "src1", [0.1, 0.2, 0.3])
 
         assert len(docs) >= 1
 
@@ -278,7 +282,7 @@ class TestGraphRAGHappyPath:
         self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
     ):
         store = _store_with_graph([], [], {}, {}, [])
-        store.count_nodes.return_value = 5
+        store.count_nodes_many.side_effect = lambda ids: {s: 5 for s in ids}
         mock_store_cls.return_value = store
 
         rag = _make_retriever()
@@ -449,3 +453,296 @@ class TestGraphRAGTopK:
         docs = rag._get_data()
 
         assert len(docs) == 2
+
+
+# ── Embeddings resolution ─────────────────────────────────────────────────────
+
+
+@pytest.mark.unit
+class TestEmbedQueryResolution:
+    def test_embed_query_uses_shared_resolver(self):
+        """Query embedding must go through ``get_embeddings``.
+
+        Only the resolver knows the bundled local-model path, so calling the
+        singleton directly loads a second copy of the model (or crashes on the
+        positional key).
+        """
+        fake = Mock()
+        fake.embed_query.return_value = [0.1, 0.2, 0.3]
+
+        with patch(
+            "application.retriever.graph_rag.get_embeddings", return_value=fake
+        ) as mock_resolver:
+            result = GraphRAGRetriever._embed_query(object(), "a question")
+
+        mock_resolver.assert_called_once_with()
+        fake.embed_query.assert_called_once_with("a question")
+        assert result == [0.1, 0.2, 0.3]
+
+
+# ── Batched retrieval across sources ─────────────────────────────────────────
+
+
+def _multi_source_retriever(sources, **overrides):
+    """Retriever over several attached sources."""
+    return _make_retriever(
+        source={"question": "q", "active_docs": list(sources)}, **overrides
+    )
+
+
+def _recording_classic(rag, docs):
+    """Stub ``ClassicRAG._get_data`` that records the sources it was handed."""
+    seen = []
+
+    def _run():
+        seen.append(list(rag._classic.vectorstores))
+        return [dict(doc) for doc in docs]
+
+    rag._classic._get_data = Mock(side_effect=_run)
+    return seen
+
+
+def _single_node_store(counts):
+    """Graph store whose every source yields one chunk, with ``counts`` shape."""
+    store = _store_with_graph(
+        [{"id": "n1", "doc_freq": 1}],
+        [],
+        {"n1": ["c1"]},
+        {"c1": "graph text"},
+        [{"id": "n1", "distance": 0.0}],
+    )
+    store.count_nodes_many.side_effect = lambda ids: {
+        source_id: counts[source_id] for source_id in ids
+    }
+    return store
+
+
+_CLASSIC_DOC = {"title": "cl", "text": "classic", "source": "a", "filename": "cl"}
+
+
+class _SourceConfig:
+    """Minimal stand-in for the Dispatcher's per-source RetrievalConfig."""
+
+    def __init__(self, chunks: int):
+        self.chunks = chunks
+
+
+@pytest.mark.unit
+class TestGraphRAGBatching:
+    """N attached sources cost one count query and one classic run, not N of each."""
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_node_counts_fetched_in_one_query(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        store = MagicMock()
+        store.count_nodes_many.return_value = {"a": 0, "b": 0, "c": 0}
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b", "c"])
+        _recording_classic(rag, [])
+
+        rag._get_data()
+
+        store.count_nodes_many.assert_called_once_with(["a", "b", "c"])
+        store.count_nodes.assert_not_called()
+
+    @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10)
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_graphless_sources_share_one_classic_call(
+        self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
+    ):
+        store = _single_node_store({"a": 0, "b": 3, "c": 0})
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b", "c"], chunks=3)
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        docs = rag._get_data()
+
+        assert rag._classic._get_data.call_count == 1
+        assert seen == [["a", "c"]]
+        # The classic batch occupies the slot of the first graphless source, so
+        # the graph source's docs still follow it in attachment order.
+        assert [doc["text"] for doc in docs] == ["classic", "graph text"]
+
+    @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10)
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_only_the_batched_sources_keep_their_overrides(
+        self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
+    ):
+        store = _single_node_store({"a": 0, "b": 3, "c": 0})
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b", "c"], chunks=3)
+        configs = {sid: _SourceConfig(2) for sid in ("a", "b", "c")}
+        rag.per_source_retrieval = dict(configs)
+        captured = {}
+
+        def _run():
+            captured["overrides"] = dict(rag._classic.per_source_retrieval)
+            return []
+
+        rag._classic._get_data = Mock(side_effect=_run)
+
+        rag._get_data()
+
+        assert captured["overrides"] == {"a": configs["a"], "c": configs["c"]}
+        # Restored afterwards, exactly as the per-source path did.
+        assert rag._classic.per_source_retrieval == {}
+
+    @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10)
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_query_is_embedded_once_for_several_graph_sources(
+        self, _avail, mock_store_cls, _tok, _patch_llm_creator
+    ):
+        store = _single_node_store({"a": 3, "b": 3})
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b"], chunks=4)
+        rag._embed_query = Mock(return_value=[0.1, 0.2, 0.3])
+
+        docs = rag._get_data()
+
+        assert rag._embed_query.call_count == 1
+        assert store.search_nodes_by_embedding.call_count == 2
+        # The one vector is what every source searches with.
+        for call in store.search_nodes_by_embedding.call_args_list:
+            assert call.args[1] == [0.1, 0.2, 0.3]
+        assert [doc["text"] for doc in docs] == ["graph text", "graph text"]
+
+    @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10)
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_failed_graph_sources_land_in_one_batched_fallback(
+        self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
+    ):
+        store = _single_node_store({"a": 3, "b": 3, "c": 3})
+
+        def _seed(source_id, embedding, k=10):
+            if source_id in ("b", "c"):
+                raise RuntimeError("graph exploded")
+            return [{"id": "n1", "distance": 0.0}]
+
+        store.search_nodes_by_embedding.side_effect = _seed
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b", "c"], chunks=6)
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        docs = rag._get_data()
+
+        assert rag._classic._get_data.call_count == 1
+        assert seen == [["b", "c"]]
+        # The retried batch is appended after the graph results.
+        assert [doc["text"] for doc in docs] == ["graph text", "classic"]
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_embedding_failure_falls_back_for_every_graph_source(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        store = _single_node_store({"a": 3, "b": 3})
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b"])
+        rag._embed_query = Mock(side_effect=RuntimeError("no embeddings"))
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        docs = rag._get_data()
+
+        assert seen == [["a", "b"]]
+        assert [doc["text"] for doc in docs] == ["classic"]
+        store.search_nodes_by_embedding.assert_not_called()
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_count_failure_falls_back_in_one_call(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        store = MagicMock()
+        store.count_nodes_many.side_effect = RuntimeError("no graph tables")
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a", "b"])
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        docs = rag._get_data()
+
+        assert seen == [["a", "b"]]
+        assert [doc["text"] for doc in docs] == ["classic"]
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_unbuildable_store_falls_back_in_one_call(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        mock_store_cls.side_effect = RuntimeError("no connection string")
+
+        rag = _multi_source_retriever(["a", "b"])
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        rag._get_data()
+
+        assert seen == [["a", "b"]]
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=False)
+    def test_unavailable_graphrag_makes_one_batched_classic_call(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        rag = _multi_source_retriever(["a", "b", "c"])
+        seen = _recording_classic(rag, [_CLASSIC_DOC])
+
+        rag._get_data()
+
+        assert rag._classic._get_data.call_count == 1
+        assert seen == [["a", "b", "c"]]
+        mock_store_cls.assert_not_called()
+
+    @patch("application.retriever.graph_rag.num_tokens_from_string", return_value=10)
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_store_is_closed_on_the_success_path(
+        self, _avail, mock_store_cls, _tok, _patch_llm_creator, _patch_embed
+    ):
+        store = _single_node_store({"a": 3})
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a"], chunks=2)
+        rag._get_data()
+
+        store.close.assert_called_once()
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_store_is_closed_when_retrieval_raises(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        store = MagicMock()
+        store.count_nodes_many.return_value = {"a": 0}
+        mock_store_cls.return_value = store
+
+        rag = _multi_source_retriever(["a"])
+        rag._classic._get_data = Mock(side_effect=RuntimeError("boom"))
+
+        with pytest.raises(RuntimeError):
+            rag._get_data()
+
+        store.close.assert_called_once()
+
+    @patch("application.retriever.graph_rag.GraphStore")
+    @patch("application.retriever.graph_rag.graphrag_available", return_value=True)
+    def test_empty_source_list_never_builds_a_store(
+        self, _avail, mock_store_cls, _patch_llm_creator
+    ):
+        rag = _make_retriever(source={"question": "q", "active_docs": []})
+        rag._classic._get_data = Mock(return_value=[])
+
+        assert rag._get_data() == []
+        mock_store_cls.assert_not_called()
+        rag._classic._get_data.assert_not_called()

tests/retriever/test_parallel_sources.py

@@ -0,0 +1,374 @@
+"""Multi-source retrieval: one query embedding, concurrent per-source search.
+
+Six attached sources used to cost six query embeddings and six serial round
+trips. These tests pin the two halves of the fix: the query is embedded once
+and the vector is handed to every store, and the per-source searches run
+concurrently while the merged output keeps the original source order.
+"""
+
+import threading
+import time
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+
+from application.retriever.classic_rag import ClassicRAG
+from application.retriever.hybrid_rag import HybridRetriever
+
+
+@pytest.fixture
+def _patch_llm_creator(mock_llm, monkeypatch):
+    monkeypatch.setattr(
+        "application.retriever.classic_rag.LLMCreator.create_llm",
+        Mock(return_value=mock_llm),
+    )
+    return mock_llm
+
+
+def _make_doc(page_content, source="s", title="t"):
+    doc = Mock()
+    doc.page_content = page_content
+    doc.metadata = {"title": title, "source": source}
+    return doc
+
+
+def _make_embedder(vector=(0.1, 0.2, 0.3)):
+    embedder = Mock()
+    embedder.embed_query = Mock(return_value=list(vector))
+    return embedder
+
+
+def _make_store(embedder, hits=None, keyword_hits=None):
+    store = Mock()
+    store._embedding = embedder
+    store.score_kind = None
+    store.search.return_value = list(hits or [])
+    store.search_with_scores.return_value = [(doc, 0.5) for doc in (hits or [])]
+    store.keyword_search.return_value = list(keyword_hits or [])
+    return store
+
+
+def _make_rag(cls=ClassicRAG, source=None, **overrides):
+    kwargs = dict(
+        source=source or {"question": "q", "active_docs": ["a", "b", "c"]},
+        chat_history=None,
+        prompt="",
+        chunks=6,
+        doc_token_limit=50000,
+        model_id="test-model",
+        llm_name="openai",
+        api_key="fake",
+        decoded_token={"sub": "user1"},
+    )
+    kwargs.update(overrides)
+    return cls(**kwargs)
+
+
+def _run(rag, stores):
+    """Run ``_get_data`` with ``VectorCreator`` handing out ``stores`` by id."""
+    with patch(
+        "application.retriever.classic_rag.VectorCreator.create_vectorstore",
+        side_effect=lambda _type, source_id, _key: stores[source_id],
+    ), patch(
+        "application.retriever.classic_rag.num_tokens_from_string", return_value=10
+    ):
+        return rag._get_data()
+
+
+@pytest.mark.unit
+class TestQueryEmbeddedOnce:
+    def test_three_sources_embed_the_query_once(self, _patch_llm_creator):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b", "c")
+        }
+
+        docs = _run(_make_rag(), stores)
+
+        assert len(docs) == 3
+        assert embedder.embed_query.call_count == 1
+        embedder.embed_query.assert_called_once_with("q")
+
+    def test_every_store_receives_the_same_vector(self, _patch_llm_creator):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b", "c")
+        }
+
+        _run(_make_rag(), stores)
+
+        vectors = [
+            store.search.call_args.kwargs["query_vector"] for store in stores.values()
+        ]
+        assert vectors == [[0.1, 0.2, 0.3]] * 3
+
+    def test_include_scores_path_passes_the_vector_too(self, _patch_llm_creator):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b")
+        }
+
+        _run(_make_rag(include_scores=True, source={
+            "question": "q", "active_docs": ["a", "b"]
+        }), stores)
+
+        assert embedder.embed_query.call_count == 1
+        for store in stores.values():
+            kwargs = store.search_with_scores.call_args.kwargs
+            assert kwargs["query_vector"] == [0.1, 0.2, 0.3]
+
+    def test_hybrid_reuses_the_shared_vector(self, _patch_llm_creator):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b")
+        }
+
+        _run(
+            _make_rag(
+                cls=HybridRetriever,
+                source={"question": "q", "active_docs": ["a", "b"]},
+            ),
+            stores,
+        )
+
+        assert embedder.embed_query.call_count == 1
+        for store in stores.values():
+            assert store.search.call_args.kwargs["query_vector"] == [0.1, 0.2, 0.3]
+
+    def test_distinct_questions_are_embedded_once_each(
+        self, _patch_llm_creator, mock_llm
+    ):
+        """Per-source rephrase means two queries — one embedding each, no more."""
+        from application.storage.db.source_config import RetrievalConfig
+
+        mock_llm.gen = Mock(return_value="REPHRASED")
+        embedder = Mock()
+        embedder.embed_query = Mock(side_effect=lambda q: [len(q)])
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b", "c")
+        }
+
+        rag = _make_rag(
+            source={"question": "original", "active_docs": ["a", "b", "c"]},
+            chat_history=[{"prompt": "hi", "response": "yo"}],
+            defer_rephrase=True,
+        )
+        rag.per_source_retrieval = {
+            "a": RetrievalConfig(rephrase_query=False),
+            "b": RetrievalConfig(rephrase_query=True),
+            "c": RetrievalConfig(rephrase_query=True),
+        }
+        _run(rag, stores)
+
+        assert embedder.embed_query.call_count == 2
+        assert stores["a"].search.call_args.kwargs["query_vector"] == [len("original")]
+        assert stores["b"].search.call_args.kwargs["query_vector"] == [len("REPHRASED")]
+        assert stores["c"].search.call_args.kwargs["query_vector"] == [len("REPHRASED")]
+
+    def test_embedding_failure_falls_back_to_per_store_embedding(
+        self, _patch_llm_creator
+    ):
+        """A broken embedder must not break retrieval — stores embed their own."""
+        embedder = Mock()
+        embedder.embed_query = Mock(side_effect=RuntimeError("model down"))
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b")
+        }
+
+        docs = _run(
+            _make_rag(source={"question": "q", "active_docs": ["a", "b"]}), stores
+        )
+
+        assert len(docs) == 2
+        for store in stores.values():
+            assert "query_vector" not in store.search.call_args.kwargs
+
+    def test_store_without_embedder_still_searches(self, _patch_llm_creator):
+        """A store exposing no embeddings object keeps the old call shape."""
+        store = _make_store(None, [_make_doc("hit")])
+        store._embedding = None
+        store._embeddings = None
+        store.embeddings = None
+        store._get_embeddings = Mock(side_effect=AttributeError("nope"))
+
+        docs = _run(
+            _make_rag(source={"question": "q", "active_docs": ["a"]}), {"a": store}
+        )
+
+        assert len(docs) == 1
+        assert "query_vector" not in store.search.call_args.kwargs
+
+
+@pytest.mark.unit
+class TestConcurrentSourceSearch:
+    def test_sources_are_searched_concurrently(self, _patch_llm_creator):
+        """Three sources must be in flight at once, not one after another.
+
+        The barrier only clears if three searches overlap; a serial loop
+        would trip its timeout and lose every source.
+        """
+        barrier = threading.Barrier(3, timeout=10)
+        embedder = _make_embedder()
+
+        def _search_factory(src):
+            def _search(*_args, **_kwargs):
+                barrier.wait()
+                return [_make_doc(f"hit {src}", source=src)]
+
+            return _search
+
+        stores = {}
+        for src in ("a", "b", "c"):
+            store = _make_store(embedder)
+            store.search.side_effect = _search_factory(src)
+            stores[src] = store
+
+        docs = _run(_make_rag(), stores)
+
+        assert [d["text"] for d in docs] == ["hit a", "hit b", "hit c"]
+
+    def test_results_keep_source_order_regardless_of_completion_order(
+        self, _patch_llm_creator
+    ):
+        """The slowest source is first — it still leads the merged output."""
+        embedder = _make_embedder()
+        delays = {"a": 0.06, "b": 0.0, "c": 0.0}
+
+        def _search_factory(src):
+            def _search(*_args, **_kwargs):
+                time.sleep(delays[src])
+                return [_make_doc(f"hit {src}", source=src)]
+
+            return _search
+
+        stores = {}
+        for src in ("a", "b", "c"):
+            store = _make_store(embedder)
+            store.search.side_effect = _search_factory(src)
+            stores[src] = store
+
+        docs = _run(_make_rag(), stores)
+
+        assert [d["text"] for d in docs] == ["hit a", "hit b", "hit c"]
+        assert [d["source"] for d in docs] == ["a", "b", "c"]
+
+    def test_output_matches_the_serial_path_exactly(self, _patch_llm_creator):
+        """Parity: one worker (serial) and four workers agree, doc for doc."""
+
+        def _stores():
+            embedder = _make_embedder()
+            return {
+                src: _make_store(
+                    embedder,
+                    [
+                        _make_doc(f"{src}-0", source=src),
+                        _make_doc(f"{src}-1", source=src),
+                    ],
+                )
+                for src in ("a", "b", "c")
+            }
+
+        with patch(
+            "application.retriever.classic_rag._max_parallel_sources",
+            side_effect=lambda n: 1,
+        ):
+            serial = _run(_make_rag(), _stores())
+        with patch(
+            "application.retriever.classic_rag._max_parallel_sources",
+            side_effect=lambda n: 4,
+        ):
+            parallel = _run(_make_rag(), _stores())
+
+        assert serial == parallel
+        assert len(serial) == 6
+
+    def test_one_failing_source_does_not_kill_the_others(self, _patch_llm_creator):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b", "c")
+        }
+        stores["b"].search.side_effect = RuntimeError("store b is down")
+
+        docs = _run(_make_rag(), stores)
+
+        assert [d["text"] for d in docs] == ["hit a", "hit c"]
+
+    def test_failing_store_construction_does_not_kill_the_others(
+        self, _patch_llm_creator
+    ):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [_make_doc(f"hit {src}", source=src)])
+            for src in ("a", "b", "c")
+        }
+
+        def _create(_type, source_id, _key):
+            if source_id == "a":
+                raise RuntimeError("connection failed")
+            return stores[source_id]
+
+        rag = _make_rag()
+        with patch(
+            "application.retriever.classic_rag.VectorCreator.create_vectorstore",
+            side_effect=_create,
+        ), patch(
+            "application.retriever.classic_rag.num_tokens_from_string", return_value=10
+        ):
+            docs = rag._get_data()
+
+        assert [d["text"] for d in docs] == ["hit b", "hit c"]
+
+    def test_single_source_skips_the_pool(self, _patch_llm_creator):
+        """One source must not pay for a thread pool."""
+        embedder = _make_embedder()
+        store = _make_store(embedder, [_make_doc("only hit")])
+        main_thread = threading.current_thread().name
+        seen = []
+        store.search.side_effect = lambda *_a, **_k: (
+            seen.append(threading.current_thread().name),
+            [_make_doc("only hit")],
+        )[1]
+
+        docs = _run(
+            _make_rag(source={"question": "q", "active_docs": ["a"]}), {"a": store}
+        )
+
+        assert len(docs) == 1
+        assert seen == [main_thread]
+
+
+@pytest.mark.unit
+class TestWorkerCap:
+    """``RETRIEVAL_MAX_PARALLEL_SOURCES`` bounds the fan-out."""
+
+    def _cap(self, monkeypatch, n_sources, configured=None):
+        import application.retriever.classic_rag as classic_rag
+
+        stub = SimpleNamespace()
+        if configured is not None:
+            stub.RETRIEVAL_MAX_PARALLEL_SOURCES = configured
+        monkeypatch.setattr(classic_rag, "settings", stub)
+        return classic_rag._max_parallel_sources(n_sources)
+
+    def test_default_is_four(self, monkeypatch):
+        assert self._cap(monkeypatch, 10) == 4
+
+    def test_never_more_workers_than_sources(self, monkeypatch):
+        assert self._cap(monkeypatch, 2) == 2
+
+    def test_setting_is_honoured(self, monkeypatch):
+        assert self._cap(monkeypatch, 10, configured=8) == 8
+
+    def test_floor_is_one(self, monkeypatch):
+        assert self._cap(monkeypatch, 10, configured=0) == 1
+
+    def test_garbage_setting_falls_back_to_default(self, monkeypatch):
+        assert self._cap(monkeypatch, 10, configured="lots") == 4

tests/services/test_search_service.py

@@ -287,3 +287,340 @@ class TestSourceAuthorization:
 
         monkeypatch.setattr(ts, "can_access", _boom)
         assert _authorized_source_ids(None, self._agent(), ["s1"]) == []
+
+
+def _serial_search_sources(query, source_ids, chunks):
+    """The pre-fan-out ``_search_sources``, kept as a parity oracle."""
+    from application.services.search_service import VectorCreator, settings
+
+    if chunks <= 0 or not source_ids:
+        return []
+
+    results = []
+    chunks_per_source = max(1, chunks // len(source_ids))
+    seen_texts = set()
+
+    for source_id in source_ids:
+        if not source_id or not source_id.strip():
+            continue
+        try:
+            docsearch = VectorCreator.create_vectorstore(
+                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
+            )
+            docs = docsearch.search(query, k=chunks_per_source * 2)
+            for doc in docs:
+                if len(results) >= chunks:
+                    break
+                if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
+                    page_content = doc.page_content
+                    metadata = doc.metadata
+                else:
+                    page_content = doc.get("text", doc.get("page_content", ""))
+                    metadata = doc.get("metadata", {})
+                text_hash = hash(page_content[:200])
+                if text_hash in seen_texts:
+                    continue
+                seen_texts.add(text_hash)
+                title = metadata.get("title", metadata.get("post_title", ""))
+                if not isinstance(title, str):
+                    title = str(title) if title else ""
+                if title:
+                    title = title.split("/")[-1]
+                else:
+                    title = metadata.get("filename", page_content[:50] + "...")
+                source = metadata.get("source", source_id)
+                results.append(
+                    {"text": page_content, "title": title, "source": source}
+                )
+            if len(results) >= chunks:
+                break
+        except Exception:
+            continue
+
+    return results[:chunks]
+
+
+def _make_embedder(vector=(0.1, 0.2, 0.3)):
+    embedder = MagicMock()
+    embedder.embed_query = MagicMock(return_value=list(vector))
+    return embedder
+
+
+def _make_store(embedder, docs=None):
+    store = MagicMock()
+    store._embedding = embedder
+    store.search.return_value = list(docs or [])
+    return store
+
+
+def _no_embedder_store(docs=None):
+    """A store exposing no reachable embeddings object."""
+    store = MagicMock()
+    store._embedding = None
+    store._embeddings = None
+    store.embeddings = None
+    store._get_embeddings = MagicMock(side_effect=AttributeError("nope"))
+    store.search.return_value = list(docs or [])
+    return store
+
+
+def _run_search_sources(query, source_ids, chunks, stores, serial=False):
+    """Run ``_search_sources`` (or its serial oracle) against ``stores``."""
+    from application.services import search_service
+
+    impl = _serial_search_sources if serial else search_service._search_sources
+    with patch(
+        "application.services.search_service.VectorCreator.create_vectorstore",
+        side_effect=lambda _type, source_id, _key: stores[source_id],
+    ):
+        return impl(query, source_ids, chunks)
+
+
+@pytest.mark.unit
+class TestSearchSourcesFanOut:
+    """Multi-source search: one query embedding, one bounded fan-out."""
+
+    def test_query_is_embedded_once_across_three_sources(self):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b", "c")
+        }
+
+        results = _run_search_sources("q", ["a", "b", "c"], 5, stores)
+
+        assert len(results) == 3
+        embedder.embed_query.assert_called_once_with("q")
+
+    def test_every_store_receives_the_same_vector(self):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b", "c")
+        }
+
+        _run_search_sources("q", ["a", "b", "c"], 5, stores)
+
+        vectors = [
+            store.search.call_args.kwargs["query_vector"] for store in stores.values()
+        ]
+        assert vectors == [[0.1, 0.2, 0.3]] * 3
+
+    def test_k_per_source_is_unchanged(self):
+        embedder = _make_embedder()
+        stores = {src: _make_store(embedder) for src in ("a", "b", "c")}
+
+        _run_search_sources("q", ["a", "b", "c"], 6, stores)
+
+        for store in stores.values():
+            assert store.search.call_args.kwargs["k"] == 4
+
+    def test_output_matches_the_serial_implementation(self):
+        """Parity: order, dedupe and the ``chunks`` cap are byte-identical."""
+        duplicate = "shared chunk " * 30
+
+        def _stores():
+            embedder = _make_embedder()
+            return {
+                "a": _make_store(
+                    embedder,
+                    [
+                        {"text": "a-0", "metadata": {"title": "A0", "source": "/a0"}},
+                        {"text": duplicate, "metadata": {"title": "dup"}},
+                        {"text": "a-2", "metadata": {"filename": "a2.pdf"}},
+                    ],
+                ),
+                "b": _make_store(
+                    embedder,
+                    [
+                        {"text": duplicate, "metadata": {"title": "dup again"}},
+                        {"text": "b-1", "metadata": {}},
+                    ],
+                ),
+                "c": _make_store(
+                    embedder,
+                    [
+                        {"text": "c-0", "metadata": {"post_title": "x/y/C0"}},
+                        {"text": "c-1", "metadata": {"title": 42}},
+                    ],
+                ),
+            }
+
+        for chunks in (1, 3, 5, 6, 12):
+            serial = _run_search_sources("q", ["a", "b", "c"], chunks, _stores(), serial=True)
+            parallel = _run_search_sources("q", ["a", "b", "c"], chunks, _stores())
+            assert parallel == serial, f"mismatch at chunks={chunks}"
+
+    def test_langchain_documents_match_the_serial_implementation(self):
+        def _stores():
+            embedder = _make_embedder()
+            stores = {}
+            for src in ("a", "b"):
+                doc = MagicMock()
+                doc.page_content = f"content {src}"
+                doc.metadata = {"title": f"path/to/{src}", "source": f"/{src}"}
+                stores[src] = _make_store(embedder, [doc])
+            return stores
+
+        serial = _run_search_sources("q", ["a", "b"], 4, _stores(), serial=True)
+        parallel = _run_search_sources("q", ["a", "b"], 4, _stores())
+        assert parallel == serial
+        assert [r["title"] for r in parallel] == ["a", "b"]
+
+    def test_results_keep_source_order_when_a_slow_source_is_first(self):
+        import time
+
+        embedder = _make_embedder()
+        delays = {"a": 0.06, "b": 0.0, "c": 0.0}
+
+        def _factory(src):
+            def _search(*_args, **_kwargs):
+                time.sleep(delays[src])
+                return [{"text": f"hit {src}", "metadata": {}}]
+
+            return _search
+
+        stores = {}
+        for src in ("a", "b", "c"):
+            store = _make_store(embedder)
+            store.search.side_effect = _factory(src)
+            stores[src] = store
+
+        results = _run_search_sources("q", ["a", "b", "c"], 6, stores)
+
+        assert [r["text"] for r in results] == ["hit a", "hit b", "hit c"]
+
+    def test_sources_are_searched_concurrently(self):
+        """The barrier only clears if three searches overlap."""
+        import threading
+
+        barrier = threading.Barrier(3, timeout=10)
+        embedder = _make_embedder()
+
+        def _factory(src):
+            def _search(*_args, **_kwargs):
+                barrier.wait()
+                return [{"text": f"hit {src}", "metadata": {}}]
+
+            return _search
+
+        stores = {}
+        for src in ("a", "b", "c"):
+            store = _make_store(embedder)
+            store.search.side_effect = _factory(src)
+            stores[src] = store
+
+        results = _run_search_sources("q", ["a", "b", "c"], 6, stores)
+
+        assert [r["text"] for r in results] == ["hit a", "hit b", "hit c"]
+
+    def test_single_source_skips_the_pool(self):
+        import threading
+
+        main_thread = threading.current_thread().name
+        seen = []
+        embedder = _make_embedder()
+        store = _make_store(embedder)
+        store.search.side_effect = lambda *_a, **_k: (
+            seen.append(threading.current_thread().name),
+            [{"text": "only", "metadata": {}}],
+        )[1]
+
+        results = _run_search_sources("q", ["a"], 5, {"a": store})
+
+        assert [r["text"] for r in results] == ["only"]
+        assert seen == [main_thread]
+
+    def test_one_failing_search_does_not_kill_the_others(self):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b", "c")
+        }
+        stores["b"].search.side_effect = RuntimeError("index missing")
+
+        results = _run_search_sources("q", ["a", "b", "c"], 6, stores)
+
+        assert [r["text"] for r in results] == ["hit a", "hit c"]
+
+    def test_failing_first_store_does_not_kill_the_others(self):
+        embedder = _make_embedder()
+        stores = {
+            src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b", "c")
+        }
+
+        def _create(_type, source_id, _key):
+            if source_id == "a":
+                raise RuntimeError("connection failed")
+            return stores[source_id]
+
+        from application.services import search_service
+
+        with patch(
+            "application.services.search_service.VectorCreator.create_vectorstore",
+            side_effect=_create,
+        ):
+            results = search_service._search_sources("q", ["a", "b", "c"], 6)
+
+        assert [r["text"] for r in results] == ["hit b", "hit c"]
+
+    def test_no_embedder_means_no_query_vector_kwarg(self):
+        stores = {
+            src: _no_embedder_store([{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b")
+        }
+
+        results = _run_search_sources("q", ["a", "b"], 6, stores)
+
+        assert [r["text"] for r in results] == ["hit a", "hit b"]
+        for store in stores.values():
+            assert "query_vector" not in store.search.call_args.kwargs
+
+    def test_embedding_failure_falls_back_to_per_store_embedding(self):
+        embedder = MagicMock()
+        embedder.embed_query = MagicMock(side_effect=RuntimeError("model down"))
+        stores = {
+            src: _make_store(embedder, [{"text": f"hit {src}", "metadata": {}}])
+            for src in ("a", "b")
+        }
+
+        results = _run_search_sources("q", ["a", "b"], 6, stores)
+
+        assert [r["text"] for r in results] == ["hit a", "hit b"]
+        for store in stores.values():
+            assert "query_vector" not in store.search.call_args.kwargs
+
+    def test_blank_source_ids_are_skipped_but_still_split_the_budget(self):
+        """Blank ids build no store, yet ``chunks_per_source`` counts them."""
+        embedder = _make_embedder()
+        stores = {"a": _make_store(embedder, [{"text": "hit a", "metadata": {}}])}
+
+        from application.services import search_service
+
+        with patch(
+            "application.services.search_service.VectorCreator.create_vectorstore",
+            side_effect=lambda _t, source_id, _k: stores[source_id],
+        ):
+            results = search_service._search_sources("q", ["  ", "a", ""], 6)
+
+        assert [r["text"] for r in results] == ["hit a"]
+        assert stores["a"].search.call_args.kwargs["k"] == 4
+
+    def test_all_blank_source_ids_build_no_store(self):
+        from application.services import search_service
+
+        with patch(
+            "application.services.search_service.VectorCreator.create_vectorstore"
+        ) as mock_create:
+            assert search_service._search_sources("q", ["  ", ""], 5) == []
+        mock_create.assert_not_called()
+
+    def test_zero_chunks_short_circuits(self):
+        from application.services import search_service
+
+        with patch(
+            "application.services.search_service.VectorCreator.create_vectorstore"
+        ) as mock_create:
+            assert search_service._search_sources("q", ["a"], 0) == []
+        mock_create.assert_not_called()

tests/storage/db/test_bootstrap_vector_schema.py

@@ -0,0 +1,200 @@
+"""Boot-time ownership of the pgvector schema.
+
+``PGVectorStore`` used to run its DDL on every instantiation — once per source
+per request. The schema is created here instead, once per process, and the boot
+hook is also the only place that can catch an embedding-dimension mismatch:
+a ``documents`` table built for a different model silently retrieves garbage.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from application.storage.db.bootstrap import ensure_vector_schema
+
+_APP_PY = Path(__file__).resolve().parents[3] / "application" / "app.py"
+
+
+@pytest.fixture
+def vector_settings(monkeypatch):
+    """Settings configured for a pgvector deployment."""
+    from application.core import settings as settings_module
+
+    settings = settings_module.settings
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
+    monkeypatch.setattr(
+        settings,
+        "PGVECTOR_CONNECTION_STRING",
+        "postgresql://user:pass@localhost/db",
+        raising=False,
+    )
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
+    monkeypatch.setattr(settings, "EMBEDDINGS_NAME", "test-model", raising=False)
+    return settings
+
+
+def _embeddings(dimension):
+    stub = MagicMock()
+    stub.dimension = dimension
+    return stub
+
+
+@pytest.mark.unit
+class TestEnsureVectorSchemaSkips:
+    def test_skips_when_the_store_is_not_pgvector(self, vector_settings, monkeypatch):
+        monkeypatch.setattr(vector_settings, "VECTOR_STORE", "faiss", raising=False)
+
+        with patch("psycopg.connect") as connect:
+            ensure_vector_schema()
+
+        connect.assert_not_called()
+
+    def test_skips_when_no_connection_string_is_configured(
+        self, vector_settings, monkeypatch
+    ):
+        monkeypatch.setattr(
+            vector_settings, "PGVECTOR_CONNECTION_STRING", None, raising=False
+        )
+        monkeypatch.setattr(vector_settings, "POSTGRES_URI", None, raising=False)
+
+        with patch("psycopg.connect") as connect:
+            ensure_vector_schema()
+
+        connect.assert_not_called()
+
+
+@pytest.mark.unit
+class TestEnsureVectorSchemaCreates:
+    def _run(self, dimension=768, table_dimension=768):
+        conn = MagicMock()
+        cursor = MagicMock()
+        conn.cursor.return_value = cursor
+        with patch("psycopg.connect", return_value=conn) as connect, patch(
+            "application.vectorstore.base.get_embeddings",
+            return_value=_embeddings(dimension),
+        ), patch(
+            "application.vectorstore.pgvector.PGVectorStore.create_schema"
+        ) as vector_schema, patch(
+            "application.vectorstore.pgvector.PGVectorStore.table_dimension",
+            return_value=table_dimension,
+        ), patch(
+            "application.graphrag.store.GraphStore.create_schema"
+        ) as graph_schema, patch(
+            "application.vectorstore.pgvector._pool_for"
+        ) as pool_for:
+            ensure_vector_schema()
+        return connect, conn, cursor, vector_schema, graph_schema, pool_for
+
+    def test_creates_the_vector_schema_under_the_advisory_lock(self, vector_settings):
+        connect, conn, cursor, vector_schema, graph_schema, _ = self._run()
+
+        connect.assert_called_once_with("postgresql://user:pass@localhost/db")
+        statements = " ".join(str(c) for c in cursor.execute.call_args_list)
+        assert "pg_advisory_xact_lock" in statements
+        vector_schema.assert_called_once()
+        assert vector_schema.call_args.kwargs["dimension"] == 768
+        graph_schema.assert_not_called()
+        conn.commit.assert_called_once()
+        conn.close.assert_called_once()
+
+    def test_creates_the_graph_schema_only_when_graphrag_is_enabled(
+        self, vector_settings, monkeypatch
+    ):
+        monkeypatch.setattr(vector_settings, "GRAPHRAG_ENABLED", True, raising=False)
+
+        _, _, _, _, graph_schema, _ = self._run()
+
+        graph_schema.assert_called_once()
+        assert graph_schema.call_args.kwargs["dimension"] == 768
+
+    def test_uses_a_direct_connection_never_the_store_pool(self, vector_settings):
+        # The hook can run pre-fork under ``gunicorn --preload``; a pooled
+        # socket inherited by a worker is a corrupted connection.
+        _, _, _, _, _, pool_for = self._run()
+
+        pool_for.assert_not_called()
+
+
+@pytest.mark.unit
+class TestEnsureVectorSchemaDimensionCheck:
+    def test_raises_when_the_table_width_disagrees_with_the_model(
+        self, vector_settings
+    ):
+        conn = MagicMock()
+        with patch("psycopg.connect", return_value=conn), patch(
+            "application.vectorstore.base.get_embeddings",
+            return_value=_embeddings(1536),
+        ), patch("application.vectorstore.pgvector.PGVectorStore.create_schema"), patch(
+            "application.vectorstore.pgvector.PGVectorStore.table_dimension",
+            return_value=768,
+        ):
+            with pytest.raises(RuntimeError) as excinfo:
+                ensure_vector_schema()
+
+        message = str(excinfo.value)
+        assert "768" in message and "1536" in message
+        assert "test-model" in message
+        conn.close.assert_called_once()
+
+    def test_skips_the_check_when_the_embeddings_expose_no_dimension(
+        self, vector_settings
+    ):
+        conn = MagicMock()
+        stub = MagicMock()
+        del stub.dimension
+        with patch("psycopg.connect", return_value=conn), patch(
+            "application.vectorstore.base.get_embeddings", return_value=stub
+        ), patch(
+            "application.vectorstore.pgvector.PGVectorStore.create_schema"
+        ) as vector_schema, patch(
+            "application.vectorstore.pgvector.PGVectorStore.table_dimension",
+            return_value=1536,
+        ):
+            ensure_vector_schema()  # must not raise
+
+        assert vector_schema.call_args.kwargs["dimension"] == 768
+
+    def test_skips_the_check_when_the_model_cannot_be_loaded(self, vector_settings):
+        conn = MagicMock()
+        with patch("psycopg.connect", return_value=conn), patch(
+            "application.vectorstore.base.get_embeddings",
+            side_effect=RuntimeError("no model"),
+        ), patch(
+            "application.vectorstore.pgvector.PGVectorStore.create_schema"
+        ) as vector_schema, patch(
+            "application.vectorstore.pgvector.PGVectorStore.table_dimension",
+            return_value=1536,
+        ):
+            ensure_vector_schema()  # must not raise
+
+        assert vector_schema.call_args.kwargs["dimension"] == 768
+
+
+@pytest.mark.unit
+class TestBootGating:
+    """The boot hook must stay behind AUTO_VECTOR_SCHEMA.
+
+    Asserted on the source because importing ``application.app`` runs the hook,
+    and a test that imports it cannot observe the gate it just executed.
+    """
+
+    def test_app_calls_the_hook_behind_the_setting(self):
+        source = _APP_PY.read_text()
+
+        assert "if settings.AUTO_VECTOR_SCHEMA:" in source
+        assert "ensure_vector_schema(" in source
+
+    def test_setting_defaults_on(self):
+        from application.core.settings import Settings
+
+        assert Settings.model_fields["AUTO_VECTOR_SCHEMA"].default is True
+
+    def test_test_suite_opts_out_by_default(self):
+        conftest = Path(__file__).resolve().parents[2] / "conftest.py"
+
+        assert 'os.environ.setdefault("AUTO_VECTOR_SCHEMA", "false")' in (
+            conftest.read_text()
+        )

tests/test_cache.py

@@ -658,3 +658,254 @@ def test_gen_cache_key_handles_bytearray_and_memoryview():
     # All three should hash the same content to the same key.
     assert gen_cache_key(msgs_ba, "m") == gen_cache_key(msgs_b, "m")
     assert gen_cache_key(msgs_mv, "m") == gen_cache_key(msgs_b, "m")
+
+
+# =====================================================================
+# Generation kwargs are part of the cache key
+#
+# The decorators wrap ``_raw_gen``/``_raw_gen_stream``, whose extra kwargs
+# carry ``response_format`` (OpenAI structured output) and
+# ``response_schema`` (Google). Before these were hashed, changing a
+# workflow node's JSON schema replayed the previous schema's answer for
+# the whole 30-minute TTL.
+# =====================================================================
+
+
+class _FakeRedis:
+    """Dict-backed stand-in for the redis client used by the decorators."""
+
+    def __init__(self):
+        self.store = {}
+
+    def get(self, key):
+        return self.store.get(key)
+
+    def set(self, key, value, ex=None):
+        self.store[key] = value.encode("utf-8") if isinstance(value, str) else value
+
+    def delete(self, key):
+        self.store.pop(key, None)
+
+
+_SCHEMA_A = {
+    "type": "json_schema",
+    "json_schema": {"name": "r", "schema": {"properties": {"a": {"type": "string"}}}},
+}
+_SCHEMA_B = {
+    "type": "json_schema",
+    "json_schema": {"name": "r", "schema": {"properties": {"b": {"type": "number"}}}},
+}
+
+
+@pytest.mark.unit
+def test_gen_cache_key_differs_for_different_response_format():
+    messages = [{"role": "user", "content": "test"}]
+    key_a = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_A})
+    key_b = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_B})
+    assert key_a != key_b
+
+
+@pytest.mark.unit
+def test_gen_cache_key_stable_for_same_response_format():
+    messages = [{"role": "user", "content": "test"}]
+    key_a = gen_cache_key(messages, "m", None, extra={"response_format": _SCHEMA_A})
+    key_b = gen_cache_key(messages, "m", None, extra={"response_format": dict(_SCHEMA_A)})
+    assert key_a == key_b
+
+
+@pytest.mark.unit
+def test_gen_cache_key_response_format_differs_from_no_format():
+    messages = [{"role": "user", "content": "test"}]
+    assert gen_cache_key(messages, "m") != gen_cache_key(
+        messages, "m", None, extra={"response_format": _SCHEMA_A}
+    )
+
+
+@pytest.mark.unit
+def test_gen_cache_key_differs_for_different_response_schema():
+    """Google's structured-output kwarg is keyed just like OpenAI's."""
+    messages = [{"role": "user", "content": "test"}]
+    key_a = gen_cache_key(messages, "m", None, extra={"response_schema": _SCHEMA_A})
+    key_b = gen_cache_key(messages, "m", None, extra={"response_schema": _SCHEMA_B})
+    assert key_a != key_b
+
+
+@pytest.mark.unit
+def test_gen_cache_key_covers_other_generation_kwargs():
+    messages = [{"role": "user", "content": "test"}]
+    assert gen_cache_key(
+        messages, "m", None, extra={"temperature": 0.1}
+    ) != gen_cache_key(messages, "m", None, extra={"temperature": 0.9})
+    assert gen_cache_key(
+        messages, "m", None, extra={"reasoning_effort": "low"}
+    ) != gen_cache_key(messages, "m", None, extra={"reasoning_effort": "high"})
+
+
+@pytest.mark.unit
+def test_gen_cache_key_ignores_usage_attachments():
+    """``_usage_attachments`` is a token-accounting side channel that never
+    reaches the provider — and the gen/stream decorator stacks disagree on
+    whether it is still in kwargs — so it must not move the key."""
+    messages = [{"role": "user", "content": "test"}]
+    plain = gen_cache_key(messages, "m")
+    with_attachments = gen_cache_key(
+        messages, "m", None, extra={"_usage_attachments": [{"id": "att1"}]}
+    )
+    assert plain == with_attachments
+
+
+@pytest.mark.unit
+def test_gen_cache_key_ignores_none_valued_kwargs():
+    """``response_format=None`` is the default, not a distinct request."""
+    messages = [{"role": "user", "content": "test"}]
+    assert gen_cache_key(messages, "m") == gen_cache_key(
+        messages, "m", None, extra={"response_format": None, "response_schema": None}
+    )
+
+
+@pytest.mark.unit
+def test_gen_cache_key_extra_ordering_is_irrelevant():
+    messages = [{"role": "user", "content": "test"}]
+    first = gen_cache_key(messages, "m", None, extra={"a": 1, "b": 2})
+    second = gen_cache_key(messages, "m", None, extra={"b": 2, "a": 1})
+    assert first == second
+
+
+@pytest.mark.unit
+def test_gen_cache_key_unserializable_extra_raises_value_error():
+    """A key we cannot compute must raise so the decorators bypass the
+    cache rather than reuse a wrong entry."""
+    messages = [{"role": "user", "content": "test"}]
+    circular = {}
+    circular["self"] = circular
+    with pytest.raises(ValueError):
+        gen_cache_key(messages, "m", None, extra={"response_format": circular})
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_gen_cache_does_not_serve_entry_from_other_response_format(mock_make_redis):
+    fake = _FakeRedis()
+    mock_make_redis.return_value = fake
+    calls = []
+
+    @gen_cache
+    def mock_function(self, model, messages, stream, tools, **kwargs):
+        calls.append(kwargs.get("response_format"))
+        return f"answer-{len(calls)}"
+
+    messages = [{"role": "user", "content": "test"}]
+    first = mock_function(
+        None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_A
+    )
+    second = mock_function(
+        None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_B
+    )
+    cached = mock_function(
+        None, "m", messages, stream=False, tools=None, response_format=_SCHEMA_A
+    )
+
+    assert first == "answer-1"
+    assert second == "answer-2"
+    assert cached == "answer-1"
+    assert calls == [_SCHEMA_A, _SCHEMA_B]
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_stream_cache_does_not_replay_entry_from_other_response_format(mock_make_redis):
+    """The reported bug: a workflow node whose schema changed replayed the
+    old schema's cached stream for the rest of the TTL."""
+    fake = _FakeRedis()
+    mock_make_redis.return_value = fake
+    calls = []
+
+    @stream_cache
+    def mock_function(self, model, messages, stream, tools, **kwargs):
+        calls.append(kwargs.get("response_format"))
+        yield f"chunk-{len(calls)}"
+
+    messages = [{"role": "user", "content": "test"}]
+    first = list(
+        mock_function(
+            None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_A
+        )
+    )
+    second = list(
+        mock_function(
+            None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_B
+        )
+    )
+    replay = list(
+        mock_function(
+            None, "m", messages, stream=True, tools=None, response_format=_SCHEMA_A
+        )
+    )
+
+    assert first == ["chunk-1"]
+    assert second == ["chunk-2"]
+    assert replay == ["chunk-1"]
+    assert calls == [_SCHEMA_A, _SCHEMA_B]
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_gen_cache_bypassed_for_previous_response_id(mock_make_redis):
+    """A Responses API turn chained to a server-held id depends on state no
+    key can capture, so it must not read or write the cache."""
+    fake = MagicMock()
+    mock_make_redis.return_value = fake
+
+    @gen_cache
+    def mock_function(self, model, messages, stream, tools, **kwargs):
+        return "fresh"
+
+    messages = [{"role": "user", "content": "test"}]
+    result = mock_function(
+        None, "m", messages, stream=False, tools=None, previous_response_id="resp_1"
+    )
+
+    assert result == "fresh"
+    fake.get.assert_not_called()
+    fake.set.assert_not_called()
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_stream_cache_bypassed_for_previous_response_id(mock_make_redis):
+    fake = MagicMock()
+    mock_make_redis.return_value = fake
+
+    @stream_cache
+    def mock_function(self, model, messages, stream, tools, **kwargs):
+        yield "fresh"
+
+    messages = [{"role": "user", "content": "test"}]
+    result = list(
+        mock_function(
+            None, "m", messages, stream=True, tools=None, previous_response_id="resp_1"
+        )
+    )
+
+    assert result == ["fresh"]
+    fake.get.assert_not_called()
+    fake.set.assert_not_called()
+
+
+@pytest.mark.unit
+@patch("application.cache.get_redis_instance")
+def test_gen_cache_ignores_stream_payload_stored_under_same_key(mock_make_redis):
+    """``gen_cache`` and ``stream_cache`` share a key space; a stream
+    envelope must never be handed back as a non-streaming answer."""
+    fake = _FakeRedis()
+    mock_make_redis.return_value = fake
+    messages = [{"role": "user", "content": "test"}]
+    key = gen_cache_key(messages, "m", None)
+    fake.set(key, json.dumps({"version": 1, "chunks": ["a", "b"]}))
+
+    @gen_cache
+    def mock_function(self, model, messages, stream, tools, **kwargs):
+        return "real answer"
+
+    result = mock_function(None, "m", messages, stream=False, tools=None)
+    assert result == "real answer"

tests/vectorstore/test_base.py

@@ -6,8 +6,12 @@ from application.vectorstore.base import (
     BaseVectorStore,
     EmbeddingsSingleton,
     RemoteEmbeddings,
+    get_embeddings,
 )
 
+HF_MPNET = "huggingface_sentence-transformers/all-mpnet-base-v2"
+LOCAL_MPNET = "/app/models/all-mpnet-base-v2"
+
 
 # --- RemoteEmbeddings ---
 
@@ -254,6 +258,46 @@ class TestEmbeddingsSingleton:
         assert result.headers["Authorization"] == "Bearer sk-from-settings"
 
 
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base._get_embeddings_wrapper")
+    def test_get_instance_hf_ignores_positional_key(
+        self, mock_get_wrapper, mock_settings
+    ):
+        """A stray key must not reach the zero-arg HuggingFace factory.
+
+        The factories are ``lambda: EmbeddingsWrapper(...)``, so a caller that
+        passed ``settings.EMBEDDINGS_KEY`` positionally used to blow up with
+        ``TypeError: <lambda>() takes 0 positional arguments``.
+        """
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_wrapper_cls = Mock()
+        mock_instance = Mock()
+        mock_wrapper_cls.return_value = mock_instance
+        mock_get_wrapper.return_value = mock_wrapper_cls
+
+        result = EmbeddingsSingleton.get_instance(HF_MPNET, None)
+
+        assert result is mock_instance
+        mock_wrapper_cls.assert_called_once_with(
+            "sentence-transformers/all-mpnet-base-v2"
+        )
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base._get_embeddings_wrapper")
+    def test_get_instance_hf_ignores_keyword_args(
+        self, mock_get_wrapper, mock_settings
+    ):
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_wrapper_cls = Mock()
+        mock_get_wrapper.return_value = mock_wrapper_cls
+
+        EmbeddingsSingleton.get_instance(HF_MPNET, openai_api_key="sk-nope")
+
+        mock_wrapper_cls.assert_called_once_with(
+            "sentence-transformers/all-mpnet-base-v2"
+        )
+
+
 # --- BaseVectorStore ---
 
 
@@ -415,3 +459,159 @@ class TestSearchWithScoresDefault:
                 return []
 
         assert _Store().search_with_scores("q") == []
+
+
+# --- get_embeddings (the single resolver) ---
+
+
+@pytest.mark.unit
+class TestGetEmbeddingsResolver:
+    """``get_embeddings`` is the one entry point every caller must use.
+
+    Calling ``EmbeddingsSingleton.get_instance`` directly reproduces neither the
+    bundled local-model path nor the OpenAI/Azure key handling.
+    """
+
+    def setup_method(self):
+        EmbeddingsSingleton._instances = {}
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base._get_embeddings_wrapper")
+    @patch("os.path.exists", return_value=False)
+    def test_defaults_from_settings_do_not_raise(
+        self, _mock_exists, mock_get_wrapper, mock_settings
+    ):
+        """The default config (HF mpnet name, no key) must resolve, not crash."""
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.EMBEDDINGS_NAME = HF_MPNET
+        mock_settings.EMBEDDINGS_KEY = None
+        mock_wrapper_cls = Mock()
+        mock_instance = Mock()
+        mock_wrapper_cls.return_value = mock_instance
+        mock_get_wrapper.return_value = mock_wrapper_cls
+
+        result = get_embeddings()
+
+        assert result is mock_instance
+        assert set(EmbeddingsSingleton._instances) == {HF_MPNET}
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base._get_embeddings_wrapper")
+    @patch("os.path.exists", return_value=False)
+    def test_shares_cache_entry_with_vectorstore_helper(
+        self, _mock_exists, mock_get_wrapper, mock_settings
+    ):
+        """Same object, same cache key as the vector stores get — one model."""
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.EMBEDDINGS_NAME = HF_MPNET
+        mock_settings.EMBEDDINGS_KEY = None
+        mock_wrapper_cls = Mock()
+        mock_wrapper_cls.return_value = Mock()
+        mock_get_wrapper.return_value = mock_wrapper_cls
+
+        store_result = ConcreteVectorStore()._get_embeddings(HF_MPNET, None)
+        resolver_result = get_embeddings()
+
+        assert resolver_result is store_result
+        assert set(EmbeddingsSingleton._instances) == {HF_MPNET}
+        mock_wrapper_cls.assert_called_once()
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base._get_embeddings_wrapper")
+    @patch("os.path.exists")
+    def test_uses_local_model_path_and_caches_it(
+        self, mock_exists, mock_get_wrapper, mock_settings
+    ):
+        """With the bundled model present, the instance is keyed by its path.
+
+        A second call must not load a second copy of the model.
+        """
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.EMBEDDINGS_NAME = HF_MPNET
+        mock_settings.EMBEDDINGS_KEY = None
+        mock_exists.side_effect = lambda path: path == LOCAL_MPNET
+        mock_wrapper_cls = Mock()
+        mock_wrapper_cls.return_value = Mock()
+        mock_get_wrapper.return_value = mock_wrapper_cls
+
+        first = get_embeddings()
+        second = get_embeddings()
+
+        assert first is second
+        assert set(EmbeddingsSingleton._instances) == {LOCAL_MPNET}
+        mock_wrapper_cls.assert_called_once_with(LOCAL_MPNET)
+
+    @patch("application.vectorstore.base.settings")
+    def test_remote_when_base_url_configured(self, mock_settings):
+        mock_settings.EMBEDDINGS_BASE_URL = "http://remote:8080"
+        mock_settings.EMBEDDINGS_NAME = HF_MPNET
+        mock_settings.EMBEDDINGS_KEY = "sk-remote"
+
+        result = get_embeddings()
+
+        assert isinstance(result, RemoteEmbeddings)
+        assert result.api_url == "http://remote:8080"
+        assert result.model_name == HF_MPNET
+        assert result.headers["Authorization"] == "Bearer sk-remote"
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
+    def test_openai_passes_key(self, mock_get_instance, mock_settings):
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.OPENAI_API_BASE = None
+        mock_settings.OPENAI_API_VERSION = None
+        mock_settings.AZURE_DEPLOYMENT_NAME = None
+        mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002"
+        mock_settings.EMBEDDINGS_KEY = "sk-from-settings"
+
+        get_embeddings()
+
+        mock_get_instance.assert_called_once_with(
+            "openai_text-embedding-ada-002", openai_api_key="sk-from-settings"
+        )
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
+    def test_openai_azure_uses_deployment_name(
+        self, mock_get_instance, mock_settings
+    ):
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.OPENAI_API_BASE = "https://azure.openai.com"
+        mock_settings.OPENAI_API_VERSION = "2023-05-15"
+        mock_settings.AZURE_DEPLOYMENT_NAME = "deploy"
+        mock_settings.AZURE_EMBEDDINGS_DEPLOYMENT_NAME = "embed-deploy"
+        mock_settings.EMBEDDINGS_NAME = "openai_text-embedding-ada-002"
+        mock_settings.EMBEDDINGS_KEY = "sk-key"
+
+        get_embeddings()
+
+        mock_get_instance.assert_called_once_with(
+            "openai_text-embedding-ada-002", model="embed-deploy"
+        )
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base.EmbeddingsSingleton.get_instance")
+    def test_explicit_arguments_win_over_settings(
+        self, mock_get_instance, mock_settings
+    ):
+        mock_settings.EMBEDDINGS_BASE_URL = None
+        mock_settings.EMBEDDINGS_NAME = HF_MPNET
+        mock_settings.EMBEDDINGS_KEY = "sk-from-settings"
+
+        get_embeddings("some_custom_embedding", "sk-explicit")
+
+        mock_get_instance.assert_called_once_with("some_custom_embedding")
+
+    @patch("application.vectorstore.base.settings")
+    @patch("application.vectorstore.base.get_embeddings")
+    def test_vectorstore_helper_delegates_to_resolver(
+        self, mock_resolver, _mock_settings
+    ):
+        """``BaseVectorStore._get_embeddings`` is a thin delegate now."""
+        sentinel = Mock()
+        mock_resolver.return_value = sentinel
+
+        result = ConcreteVectorStore()._get_embeddings("a-name", "a-key")
+
+        assert result is sentinel
+        mock_resolver.assert_called_once_with("a-name", "a-key")

tests/vectorstore/test_pgvector.py

@@ -31,13 +31,17 @@ def _make_store(
 
         from application.vectorstore.pgvector import PGVectorStore
 
-        # Patch _ensure_table_exists to avoid DB calls during init
-        with patch.object(PGVectorStore, "_ensure_table_exists"):
-            store = PGVectorStore(
-                source_id=source_id,
-                embeddings_key=embeddings_key,
-                connection_string=connection_string,
-            )
+        store = PGVectorStore(
+            source_id=source_id,
+            embeddings_key=embeddings_key,
+            connection_string=connection_string,
+        )
+        # Boot owns the schema, so these unit stores start already ensured;
+        # the write-path safety net has its own tests.
+        store._schema_ensured = True
+        # Legacy direct-connection mode: pooling would build a real
+        # psycopg_pool against the fake DSN. The pooled path has its own tests.
+        store._pool_max_size = 0
         # Provide a mock connection
         mock_conn = MagicMock()
         mock_cursor = MagicMock()
@@ -438,3 +442,47 @@ class TestPGVectorSearchWithScores:
         mock_cursor.execute.side_effect = Exception("connection lost")
 
         assert store.search_with_scores("query") == []
+
+
+@pytest.mark.unit
+class TestPGVectorPrecomputedQueryVector:
+    """A caller that already embedded the query must not pay for it again.
+
+    Multi-source retrieval embeds the query once and hands the same vector to
+    every source's store; without this the query is embedded once per source.
+    """
+
+    def test_search_with_scores_skips_embedding(self):
+        store, _, mock_cursor, mock_emb = _make_store()
+        mock_cursor.fetchall.return_value = [("hello world", {}, 0.1)]
+
+        results = store.search_with_scores("query", k=2, query_vector=[0.9, 0.8, 0.7])
+
+        mock_emb.embed_query.assert_not_called()
+        assert [doc.page_content for doc, _ in results] == ["hello world"]
+
+    def test_search_skips_embedding(self):
+        store, _, mock_cursor, mock_emb = _make_store()
+        mock_cursor.fetchall.return_value = [("hello world", {}, 0.1)]
+
+        results = store.search("query", k=2, query_vector=[0.9, 0.8, 0.7])
+
+        mock_emb.embed_query.assert_not_called()
+        assert [doc.page_content for doc in results] == ["hello world"]
+
+    def test_supplied_vector_is_what_the_query_binds(self):
+        store, _, mock_cursor, _ = _make_store(source_id="src1")
+        mock_cursor.fetchall.return_value = [("a", {}, 0.1)]
+
+        store.search_with_scores("query", k=1, query_vector=[0.9, 0.8, 0.7])
+
+        _sql, params = mock_cursor.execute.call_args_list[0].args
+        assert params == ([0.9, 0.8, 0.7], "src1", [0.9, 0.8, 0.7], 1)
+
+    def test_none_vector_still_embeds(self):
+        store, _, mock_cursor, mock_emb = _make_store()
+        mock_cursor.fetchall.return_value = []
+
+        store.search_with_scores("query", k=2, query_vector=None)
+
+        mock_emb.embed_query.assert_called_once_with("query")

tests/vectorstore/test_pgvector_live_schema.py

@@ -0,0 +1,261 @@
+"""Live pgvector end-to-end: boot-owned schema, pure-read queries, pooling.
+
+Runs against the ephemeral pytest-postgresql cluster (never the operator's dev
+database) with an 8-dimension stub embeddings model, so a real ``CREATE
+EXTENSION vector`` / ``vector(8)`` round trip is exercised in a few seconds.
+Skips when the cluster's Postgres has no pgvector build installed.
+"""
+
+from __future__ import annotations
+
+import threading
+from unittest.mock import patch
+
+import pytest
+
+from application.storage.db.bootstrap import ensure_vector_schema
+from application.vectorstore import pgvector as pgvector_module
+from application.vectorstore.pgvector import PGVectorStore
+
+pytestmark = pytest.mark.integration
+
+STUB_DIM = 8
+
+
+class _StubEmbeddings:
+    """Deterministic tiny embeddings; no model download, no network."""
+
+    dimension = STUB_DIM
+
+    def embed_documents(self, texts):
+        return [self._vector(float(i + 1)) for i, _ in enumerate(texts)]
+
+    def embed_query(self, query):
+        return self._vector(1.0)
+
+    @staticmethod
+    def _vector(seed: float):
+        vector = [0.0] * STUB_DIM
+        vector[0] = seed
+        return vector
+
+
+class _WideStubEmbeddings(_StubEmbeddings):
+    dimension = 16
+
+
+def _dsn(info) -> str:
+    password = f":{info.password}" if info.password else ""
+    return (
+        f"postgresql://{info.user}{password}@{info.host}:{info.port}/{info.dbname}"
+    )
+
+
+def _regclass(conn, name: str):
+    """Look a relation up on ``conn``, in a transaction that starts fresh.
+
+    ``to_regclass`` takes no lock, so a negative catalog-cache entry from an
+    earlier lookup in the same transaction survives another backend's DDL. A
+    new transaction accepts the invalidation.
+    """
+    conn.rollback()
+    with conn.cursor() as cursor:
+        cursor.execute("SELECT to_regclass(%s);", (name,))
+        return cursor.fetchone()[0]
+
+
+@pytest.fixture(autouse=True)
+def _close_pools():
+    """Never leak a pool into another test; the DSN dies with the test DB."""
+    yield
+    for dsn, pool in list(pgvector_module._POOLS.items()):
+        try:
+            pool.close()
+        except Exception:
+            pass
+        pgvector_module._POOLS.pop(dsn, None)
+
+
+@pytest.fixture
+def live_dsn(postgresql, monkeypatch):
+    """DSN of an empty ephemeral DB, with pgvector proven installable."""
+    try:
+        with postgresql.cursor() as cursor:
+            cursor.execute("CREATE EXTENSION vector;")
+        postgresql.rollback()  # leave the database pristine for the hook
+    except Exception as exc:
+        postgresql.rollback()
+        pytest.skip(f"pgvector extension unavailable: {exc}")
+
+    dsn = _dsn(postgresql.info)
+    from application.core import settings as settings_module
+
+    settings = settings_module.settings
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector", raising=False)
+    monkeypatch.setattr(
+        settings, "PGVECTOR_CONNECTION_STRING", dsn, raising=False
+    )
+    monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", False, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_IVFFLAT_PROBES", None, raising=False)
+    monkeypatch.setattr(settings, "PGVECTOR_POOL_MAX_SIZE", 4, raising=False)
+    return dsn
+
+
+@pytest.fixture
+def stub_embeddings():
+    stub = _StubEmbeddings()
+    with patch(
+        "application.vectorstore.base.get_embeddings", return_value=stub
+    ), patch(
+        "application.vectorstore.base.BaseVectorStore._get_embeddings",
+        return_value=stub,
+    ):
+        yield stub
+
+
+def _store(dsn, source_id="live-source"):
+    return PGVectorStore(source_id=source_id, connection_string=dsn)
+
+
+class TestBootHookCreatesTheSchema:
+    def test_creates_the_table_and_indexes_and_is_idempotent(
+        self, live_dsn, postgresql, stub_embeddings
+    ):
+        ensure_vector_schema()
+        ensure_vector_schema()  # second boot must be a no-op
+
+        assert _regclass(postgresql, "documents") is not None
+        with postgresql.cursor() as cursor:
+            cursor.execute(
+                "SELECT indexname FROM pg_indexes WHERE tablename = 'documents';"
+            )
+            indexes = {row[0] for row in cursor.fetchall()}
+        assert "documents_source_id_idx" in indexes
+        assert "documents_text_fts_idx" in indexes
+        assert PGVectorStore.table_dimension(postgresql) == STUB_DIM
+
+    def test_mismatched_model_width_fails_loudly(self, live_dsn, stub_embeddings):
+        ensure_vector_schema()
+
+        wide = _WideStubEmbeddings()
+        with patch(
+            "application.vectorstore.base.get_embeddings", return_value=wide
+        ):
+            with pytest.raises(RuntimeError) as excinfo:
+                ensure_vector_schema()
+
+        message = str(excinfo.value)
+        assert f"vector({STUB_DIM})" in message
+        assert "16-dim" in message
+
+    def test_creates_graph_tables_when_graphrag_is_enabled(
+        self, live_dsn, postgresql, stub_embeddings, monkeypatch
+    ):
+        from application.core import settings as settings_module
+
+        monkeypatch.setattr(
+            settings_module.settings, "GRAPHRAG_ENABLED", True, raising=False
+        )
+
+        ensure_vector_schema()
+
+        for table in (
+            "graph_nodes",
+            "graph_edges",
+            "graph_node_chunks",
+            "graph_ingest_progress",
+        ):
+            assert _regclass(postgresql, table) is not None
+
+
+class TestStoreRoundTrip:
+    def test_construction_is_free_and_writes_then_reads_survive(
+        self, live_dsn, stub_embeddings
+    ):
+        ensure_vector_schema()
+
+        store = _store(live_dsn)
+        # Construction must not open a connection or build a pool: the
+        # retriever makes one of these per source, per request.
+        assert store._connection is None
+        assert pgvector_module._POOLS == {}
+
+        ids = store.add_texts(
+            ["alpha chunk", "beta chunk"],
+            [{"source": "a.txt"}, {"source": "b.txt"}],
+        )
+        assert len(ids) == 2
+
+        results = store.search_with_scores("alpha", k=2)
+        assert [doc.page_content for doc, _ in results] == [
+            "alpha chunk",
+            "beta chunk",
+        ]
+        assert results[0][1] == pytest.approx(1.0)
+
+        keyword_hits = store.keyword_search("alpha", k=5)
+        assert [doc.page_content for doc in keyword_hits] == ["alpha chunk"]
+
+        # The read paths must not leave the pooled connection in a transaction.
+        assert store._connection.info.transaction_status.name == "IDLE"
+        store.close()
+        assert store._connection is None
+
+
+class TestWritePathSafetyNet:
+    """A process that never ran the boot hook must still be able to ingest.
+
+    ``register_vector`` raises when the ``vector`` type is absent, which is the
+    state of a brand-new database — so the store has to tolerate that, create
+    the extension, and re-register before it inserts a vector.
+    """
+
+    def test_first_write_bootstraps_a_brand_new_database(
+        self, live_dsn, postgresql, stub_embeddings
+    ):
+        assert _regclass(postgresql, "documents") is None
+
+        store = _store(live_dsn)
+        store.add_texts(["alpha chunk"], [{"source": "a.txt"}])
+
+
+        assert _regclass(postgresql, "documents") is not None
+        assert [doc.page_content for doc in store.search("alpha", k=2)] == [
+            "alpha chunk"
+        ]
+        store.close()
+
+
+class TestPooling:
+    def test_concurrent_stores_share_one_pool_and_return_connections(
+        self, live_dsn, stub_embeddings
+    ):
+        ensure_vector_schema()
+        seed = _store(live_dsn)
+        seed.add_texts(["alpha chunk"], [{"source": "a.txt"}])
+        seed.close()
+
+        errors: list[Exception] = []
+        hits: list[int] = []
+
+        def _search():
+            store = _store(live_dsn)
+            try:
+                hits.append(len(store.search("alpha", k=2)))
+            except Exception as exc:  # surfaced below
+                errors.append(exc)
+            finally:
+                store.close()
+
+        threads = [threading.Thread(target=_search) for _ in range(2)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(timeout=20)
+
+        assert errors == []
+        assert hits == [1, 1]
+        assert list(pgvector_module._POOLS) == [live_dsn]
+        stats = pgvector_module._POOLS[live_dsn].get_stats()
+        assert stats["pool_size"] <= 4
+        assert stats["pool_available"] == stats["pool_size"]

tests/vectorstore/test_pgvector_probes.py

@@ -109,10 +109,10 @@ class TestNoVectorIndexOnEmptyTable:
     sources with hundreds of chunks, silently.
     """
 
-    def test_ensure_table_exists_creates_no_vector_index(self):
+    def test_create_schema_creates_no_vector_index(self):
         import inspect
 
-        source = inspect.getsource(PGVectorStore._ensure_table_exists)
+        source = inspect.getsource(PGVectorStore.create_schema)
         assert "USING ivfflat" not in source
         assert "USING hnsw" not in source
         # the non-vector indexes are still expected

tests/vectorstore/test_pgvector_schema_and_pool.py

@@ -0,0 +1,312 @@
+"""Schema ownership and connection pooling for ``PGVectorStore``.
+
+Constructing a store used to run six DDL statements plus a fresh unpooled
+connect, on every instantiation — and the retriever builds one store per source
+per request. Schema is now owned at boot (``ensure_vector_schema``); the query
+path is pure reads over a pooled connection, and the write path keeps a
+one-shot safety net for processes that never ran the boot hook.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+
+from application.vectorstore import pgvector as pgvector_module
+
+CONNECTION_STRING = "postgresql://user:pass@localhost/db"
+
+
+def _make_store(source_id="test-source", connection_string=CONNECTION_STRING):
+    """Build a store with every external dependency mocked, no DDL patching."""
+    with patch(
+        "application.vectorstore.base.BaseVectorStore._get_embeddings"
+    ) as mock_get_emb, patch(
+        "application.vectorstore.pgvector.settings"
+    ) as mock_settings, patch.dict(
+        "sys.modules",
+        {
+            "psycopg": MagicMock(),
+            "pgvector": MagicMock(),
+            "pgvector.psycopg": MagicMock(),
+        },
+    ):
+        mock_emb = Mock()
+        mock_emb.embed_query = Mock(return_value=[0.1, 0.2, 0.3])
+        mock_emb.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]])
+        mock_emb.dimension = 768
+        mock_get_emb.return_value = mock_emb
+        mock_settings.EMBEDDINGS_NAME = "test_model"
+        mock_settings.PGVECTOR_CONNECTION_STRING = connection_string
+
+        from application.vectorstore.pgvector import PGVectorStore
+
+        store = PGVectorStore(
+            source_id=source_id,
+            embeddings_key="key",
+            connection_string=connection_string,
+        )
+
+    mock_conn = MagicMock()
+    mock_cursor = MagicMock()
+    mock_conn.cursor.return_value = mock_cursor
+    mock_conn.closed = False
+    store._connection = mock_conn
+    store._pooled = False
+    store._ensure_table_exists = Mock()
+    return store, mock_conn, mock_cursor, mock_emb
+
+
+@pytest.mark.unit
+class TestConstructionTouchesNothing:
+    def test_init_opens_no_connection_and_runs_no_ddl(self):
+        with patch(
+            "application.vectorstore.base.BaseVectorStore._get_embeddings"
+        ) as mock_get_emb, patch(
+            "application.vectorstore.pgvector.settings"
+        ) as mock_settings, patch.dict(
+            "sys.modules",
+            {
+                "psycopg": MagicMock(),
+                "pgvector": MagicMock(),
+                "pgvector.psycopg": MagicMock(),
+            },
+        ):
+            mock_get_emb.return_value = Mock(dimension=768)
+            mock_settings.EMBEDDINGS_NAME = "test_model"
+            mock_settings.PGVECTOR_CONNECTION_STRING = CONNECTION_STRING
+
+            from application.vectorstore.pgvector import PGVectorStore
+
+            with patch.object(
+                PGVectorStore, "_get_connection"
+            ) as mock_get_conn, patch.object(
+                PGVectorStore, "_ensure_table_exists"
+            ) as mock_ensure:
+                store = PGVectorStore(
+                    source_id="s1", connection_string=CONNECTION_STRING
+                )
+
+            mock_get_conn.assert_not_called()
+            mock_ensure.assert_not_called()
+            assert store._schema_ensured is False
+
+    def test_pool_size_falls_back_when_setting_is_not_an_int(self):
+        # The unit suite patches ``settings`` with a MagicMock; a MagicMock
+        # attribute must not become the pool size.
+        store, _, _, _ = _make_store()
+        assert store._pool_max_size == 8
+
+
+@pytest.mark.unit
+class TestWritePathEnsuresSchemaOnce:
+    def test_add_texts_ensures_once_across_two_calls(self):
+        store, _, mock_cursor, _ = _make_store()
+        mock_cursor.fetchone.return_value = (1,)
+
+        store.add_texts(["a"])
+        store.add_texts(["b"])
+
+        assert store._ensure_table_exists.call_count == 1
+
+    def test_add_chunk_ensures_once_across_two_calls(self):
+        store, _, mock_cursor, _ = _make_store()
+        mock_cursor.fetchone.return_value = (1,)
+
+        store.add_chunk("a")
+        store.add_chunk("b")
+
+        assert store._ensure_table_exists.call_count == 1
+
+    def test_reads_and_deletes_never_ensure_schema(self):
+        store, _, mock_cursor, _ = _make_store()
+        mock_cursor.fetchall.return_value = []
+        mock_cursor.fetchone.return_value = (0,)
+        mock_cursor.rowcount = 0
+
+        store.search("q")
+        store.keyword_search("q")
+        store.get_chunks()
+        store.delete_index()
+        store.delete_chunk("1")
+
+        store._ensure_table_exists.assert_not_called()
+
+
+@pytest.mark.unit
+class TestCreateSchema:
+    def test_emits_the_ddl_and_leaves_the_commit_to_the_caller(self):
+        from application.vectorstore.pgvector import PGVectorStore
+
+        conn, cursor = MagicMock(), MagicMock()
+        conn.cursor.return_value = cursor
+
+        PGVectorStore.create_schema(conn, dimension=8)
+
+        statements = " ".join(str(c) for c in cursor.execute.call_args_list)
+        assert "CREATE EXTENSION IF NOT EXISTS vector" in statements
+        assert "CREATE TABLE IF NOT EXISTS documents" in statements
+        assert "vector(8)" in statements
+        assert "documents_source_id_idx" in statements
+        assert "documents_text_fts_idx" in statements
+        assert cursor.execute.call_count == 4
+        conn.commit.assert_not_called()
+
+    def test_ensure_table_exists_locks_then_commits(self):
+        store, mock_conn, mock_cursor, _ = _make_store()
+        del store._ensure_table_exists  # exercise the real method
+
+        store._ensure_table_exists()
+
+        statements = " ".join(str(c) for c in mock_cursor.execute.call_args_list)
+        assert "pg_advisory_xact_lock" in statements
+        assert "CREATE TABLE IF NOT EXISTS documents" in statements
+        mock_conn.commit.assert_called_once()
+
+
+@pytest.mark.unit
+class TestTableDimension:
+    def _conn(self, rows):
+        conn, cursor = MagicMock(), MagicMock()
+        conn.cursor.return_value = cursor
+        cursor.fetchone.side_effect = rows
+        return conn
+
+    def test_parses_the_declared_vector_width(self):
+        from application.vectorstore.pgvector import PGVectorStore
+
+        conn = self._conn([("documents",), ("vector(768)",)])
+        assert PGVectorStore.table_dimension(conn) == 768
+
+    def test_returns_none_when_the_table_is_absent(self):
+        from application.vectorstore.pgvector import PGVectorStore
+
+        conn = self._conn([(None,)])
+        assert PGVectorStore.table_dimension(conn) is None
+
+    def test_returns_none_when_the_column_type_is_not_a_vector(self):
+        from application.vectorstore.pgvector import PGVectorStore
+
+        conn = self._conn([("documents",), ("text",)])
+        assert PGVectorStore.table_dimension(conn) is None
+
+
+@pytest.mark.unit
+class TestReadPathsCloseTheirTransaction:
+    """A read that never commits leaves the connection idle-in-transaction.
+
+    On a persistent (now pooled) connection that holds a snapshot open for the
+    life of the process and makes psycopg_pool roll back every returned
+    connection.
+    """
+
+    def test_search_commits(self):
+        store, mock_conn, mock_cursor, _ = _make_store()
+        mock_cursor.fetchall.return_value = [("hello", {}, 0.1), ("bye", {}, 0.2)]
+
+        store.search_with_scores("q", k=2)
+
+        mock_conn.commit.assert_called_once()
+
+    def test_keyword_search_commits(self):
+        store, mock_conn, mock_cursor, _ = _make_store()
+        mock_cursor.fetchall.return_value = []
+
+        store.keyword_search("q")
+
+        mock_conn.commit.assert_called_once()
+
+    def test_get_chunks_commits(self):
+        store, mock_conn, mock_cursor, _ = _make_store()
+        mock_cursor.fetchall.return_value = []
+
+        store.get_chunks()
+
+        mock_conn.commit.assert_called_once()
+
+
+@pytest.mark.unit
+class TestConnectionPooling:
+    def _pooled_store(self, monkeypatch):
+        store, _, _, _ = _make_store()
+        store._connection = None
+        store._pool_max_size = 4
+        pooled_conn = MagicMock()
+        pooled_conn.closed = False
+        fake_pool = MagicMock()
+        fake_pool.getconn.return_value = pooled_conn
+        monkeypatch.setattr(
+            pgvector_module, "_pool_for", lambda dsn, size: fake_pool
+        )
+        return store, fake_pool, pooled_conn
+
+    def test_get_connection_checks_out_of_the_pool(self, monkeypatch):
+        store, fake_pool, pooled_conn = self._pooled_store(monkeypatch)
+
+        conn = store._get_connection()
+
+        assert conn is pooled_conn
+        assert store._pooled is True
+        fake_pool.getconn.assert_called_once()
+        store._psycopg.connect.assert_not_called()
+
+    def test_probes_are_applied_once_per_physical_connection(self, monkeypatch):
+        store, _, pooled_conn = self._pooled_store(monkeypatch)
+        applied = Mock()
+        monkeypatch.setattr(store, "_apply_ivfflat_probes", applied)
+
+        store._get_connection()
+        store._connection = None  # simulate a second checkout of the same conn
+        store._get_connection()
+
+        assert applied.call_count == 1
+
+    def test_close_rolls_back_and_returns_the_connection(self, monkeypatch):
+        store, fake_pool, pooled_conn = self._pooled_store(monkeypatch)
+        store._get_connection()
+        pooled_conn.info.transaction_status.name = "INTRANS"
+        monkeypatch.setitem(
+            pgvector_module._POOLS, store._connection_string, fake_pool
+        )
+
+        store.close()
+
+        pooled_conn.rollback.assert_called_once()
+        fake_pool.putconn.assert_called_once_with(pooled_conn)
+        pooled_conn.close.assert_not_called()
+        assert store._connection is None
+
+    def test_close_does_not_roll_back_an_idle_connection(self, monkeypatch):
+        store, fake_pool, pooled_conn = self._pooled_store(monkeypatch)
+        store._get_connection()
+        pooled_conn.info.transaction_status.name = "IDLE"
+        monkeypatch.setitem(
+            pgvector_module._POOLS, store._connection_string, fake_pool
+        )
+
+        store.close()
+
+        pooled_conn.rollback.assert_not_called()
+        fake_pool.putconn.assert_called_once_with(pooled_conn)
+
+    def test_legacy_path_connects_directly_and_closes(self):
+        store, mock_conn, _, _ = _make_store()
+        store._pool_max_size = 0
+        mock_conn.closed = True
+        new_conn = MagicMock()
+        new_conn.closed = False
+        store._psycopg.connect.return_value = new_conn
+
+        conn = store._get_connection()
+        assert conn is new_conn
+        assert store._pooled is False
+
+        store.close()
+        new_conn.close.assert_called_once()
+
+    def test_del_never_raises(self):
+        store, mock_conn, _, _ = _make_store()
+        mock_conn.close.side_effect = RuntimeError("already gone")
+
+        store.__del__()  # must not propagate