fix(agents): hand the graph tool's connection back, and gate it like the rest

Three faults in the graph tool, all on the agent's path:

The store was cached on the tool, and the executor caches the tool for the
whole agent run -- so one pgvector pooled connection stayed checked out across
every LLM round trip of that run, minutes at a time, and enough concurrent
runs exhaust the pool. GraphRAGRetriever releases its store before falling
back for this reason. The tool now releases it at the end of each action.

It gated on GRAPHRAG_ENABLED where everything else asks graphrag_available(),
which also requires the pgvector store. Under any other vector store the graph
tables are not the ones the sources were ingested into, but the tool was still
offered and still queried Postgres.

Pages were labelled by hand rather than through labels_from_metadata, which
exists so citation labels match across retrievers. A page read by the tool and
the same chunk retrieved by internal_search are one document, and citations
key on (source, title) -- so the research agent gave that document two
citation numbers. The recorded doc also keeps the full chunk text now, so it
dedupes against the retriever's copy; only what the model reads is truncated.

docsgpt/agents/tools/graph_search.py

@@ -28,7 +28,8 @@ import logging
 from typing import Any, Dict, List, Optional
 
 from docsgpt.agents.tools.base import Tool
-from docsgpt.core.settings import settings
+from docsgpt.graphrag import graphrag_available
+from docsgpt.retriever.labels import labels_from_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -61,6 +62,24 @@ class GraphSearchTool(Tool):
             self._store = GraphStore()
         return self._store
 
+    def _release_store(self) -> None:
+        """Hand the pooled connection back at the end of an action.
+
+        The executor caches this tool for the whole agent run, so a store kept
+        between actions pins one connection of the shared pgvector pool across
+        every LLM round trip of that run -- minutes at a time, and enough
+        concurrent runs exhaust the pool. ``GraphRAGRetriever`` releases its
+        store before falling back for the same reason. Checking one back out
+        costs a pool acquire.
+        """
+        store, self._store = self._store, None
+        if store is None:
+            return
+        try:
+            store.close()
+        except Exception as exc:  # noqa: BLE001 -- releasing must not fail an action
+            logger.debug(f"Graph tool could not release its store: {exc}")
+
     def _embed(self, text: str) -> Optional[List[float]]:
         try:
             from docsgpt.vectorstore.base import get_embeddings
@@ -72,7 +91,9 @@ class GraphSearchTool(Tool):
 
     # -- actions -------------------------------------------------------------
     def execute_action(self, action_name: str, **kwargs):
-        if not settings.GRAPHRAG_ENABLED:
+        # The graph lives in the pgvector store, so the flag alone is not
+        # enough: under another vector store there is no graph to read.
+        if not graphrag_available():
             return "The knowledge graph is not enabled for this deployment."
         if not self._sources():
             return "No graph-backed sources are configured."
@@ -86,6 +107,8 @@ class GraphSearchTool(Tool):
         except Exception as e:  # noqa: BLE001
             logger.error(f"Graph tool action {action_name} failed: {e}", exc_info=True)
             return "The graph lookup failed."
+        finally:
+            self._release_store()
         return f"Unknown action: {action_name}"
 
     def _search_entities(self, **kwargs) -> str:
@@ -137,18 +160,17 @@ class GraphSearchTool(Tool):
         parts: List[str] = []
         for source_id in self._sources():
             for page in store.entity_pages(source_id, entity):
-                metadata = page.get("metadata") or {}
-                title = (
-                    metadata.get("file_path")
-                    or metadata.get("title")
-                    or metadata.get("source")
-                    or "document"
-                )
-                text = (page.get("text") or "")[:MAX_PAGE_CHARS]
-                doc = {"title": title, "text": text, "source": metadata.get("source", "")}
+                text = page.get("text") or ""
+                # The retrievers' own labelling: a page read here and the same
+                # chunk retrieved by internal_search are one document, and
+                # citations key on (source, title). Labelling it differently
+                # gives that document two citation numbers.
+                labels = labels_from_metadata(page.get("metadata"), text, source_id)
+                doc = {**labels, "text": text}
                 if doc not in self.retrieved_docs:
                     self.retrieved_docs.append(doc)
-                parts.append(f"--- {title} ---\n{text}")
+                header = labels["filename"] or labels["title"]
+                parts.append(f"--- {header} ---\n{text[:MAX_PAGE_CHARS]}")
         if not parts:
             return f"No documents mention {entity!r}."
         return "\n\n".join(parts)
@@ -259,7 +281,7 @@ def add_graph_search_tool(tools_dict: Dict, retriever_config: Dict) -> None:
     tool follows that same per-source exposure choice. A graph source left at
     ``prefetch`` in a classic agent is used for ranking only.
     """
-    if not settings.GRAPHRAG_ENABLED:
+    if not graphrag_available():
         return
     source = retriever_config.get("source") or {}
     if not source.get("active_docs") or not sources_have_graph(source):

tests/graphrag/test_graph_search_tool.py

@@ -9,6 +9,8 @@ must refuse clearly rather than silently when it has nothing to offer.
 
 from __future__ import annotations
 
+import pytest
+
 from docsgpt.agents.tools.graph_search import (
     GRAPH_TOOL_ID,
     GraphSearchTool,
@@ -38,6 +40,7 @@ class _StubStore:
 
 def _tool(monkeypatch, store, enabled=True):
     monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", enabled)
+    monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector")
     tool = GraphSearchTool({"source": SOURCE})
     tool._store = store
     monkeypatch.setattr(tool, "_embed", lambda text: [0.0, 0.1])
@@ -52,6 +55,7 @@ class TestGating:
 
     def test_reports_when_no_sources_are_configured(self, monkeypatch):
         monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector")
         tool = GraphSearchTool({"source": {"active_docs": []}})
 
         assert "No graph-backed sources" in tool.execute_action(
@@ -109,16 +113,49 @@ class TestActions:
         tool = _tool(
             monkeypatch,
             _StubStore(
-                pages=[{"metadata": {"file_path": "quill-store.md"}, "text": "x" * 5000}]
+                pages=[
+                    {
+                        "metadata": {"title": "quill-store.md", "source": "quill-store.md"},
+                        "text": "x" * 5000,
+                    }
+                ]
             ),
         )
 
         result = tool.execute_action("read_entity_pages", entity="Quill")
 
         assert "--- quill-store.md ---" in result
+        # Only what the model reads is truncated.
         assert len(result) < 3000
         # Accumulated so the answer can cite what the walk actually read.
         assert tool.retrieved_docs[0]["title"] == "quill-store.md"
+        assert len(tool.retrieved_docs[0]["text"]) == 5000
+
+    def test_page_labels_match_what_the_retrievers_record(self, monkeypatch):
+        """A page read here and the same chunk retrieved by internal_search are
+        one document. The citation manager keys on (source, title), so labels
+        derived differently give the same document two citation numbers."""
+        from docsgpt.retriever.labels import labels_from_metadata
+
+        metadata = {"title": "Quill Store", "source": "quill-store.md"}
+        text = "Quill is a write-ahead store."
+        tool = _tool(monkeypatch, _StubStore(pages=[{"metadata": metadata, "text": text}]))
+
+        tool.execute_action("read_entity_pages", entity="Quill")
+
+        expected = labels_from_metadata(metadata, text, "src-1")
+        doc = tool.retrieved_docs[0]
+        assert {k: doc[k] for k in ("title", "source", "filename")} == expected
+        # The full chunk text, so the doc dedupes against the retriever's copy;
+        # only what the model reads is truncated.
+        assert doc["text"] == text
+
+    def test_a_page_with_no_metadata_falls_back_to_its_source_id(self, monkeypatch):
+        tool = _tool(monkeypatch, _StubStore(pages=[{"metadata": {}, "text": "body"}]))
+
+        tool.execute_action("read_entity_pages", entity="Quill")
+
+        assert tool.retrieved_docs[0]["source"] == "src-1"
 
     def test_pages_absent_is_stated_plainly(self, monkeypatch):
         tool = _tool(monkeypatch, _StubStore(pages=[]))
@@ -150,6 +187,7 @@ class TestWiring:
 
     def test_not_added_when_the_sources_have_no_graph(self, monkeypatch):
         monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector")
         monkeypatch.setattr(
             "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: False
         )
@@ -161,6 +199,7 @@ class TestWiring:
 
     def test_added_with_its_sentinel_id_and_source_config(self, monkeypatch):
         monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(settings, "VECTOR_STORE", "pgvector")
         monkeypatch.setattr(
             "docsgpt.agents.tools.graph_search.sources_have_graph", lambda source: True
         )
@@ -246,3 +285,66 @@ class TestSourcesHaveGraph:
         assert sources_have_graph({"active_docs": []}) is False
         self._patch_counts(monkeypatch, error=RuntimeError("no pgvector"))
         assert sources_have_graph({"active_docs": ["a"]}) is False
+
+
+class TestGraphsMustBeAvailable:
+    """The graph lives in the pgvector store, so the flag alone is not enough.
+
+    With another vector store configured the graph tables are not the ones the
+    sources were ingested into; everything else in the app asks
+    ``graphrag_available()``, which requires both.
+    """
+
+    def test_the_tool_is_not_offered_without_pgvector(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(settings, "VECTOR_STORE", "faiss")
+        monkeypatch.setattr(
+            "docsgpt.agents.tools.graph_search.sources_have_graph",
+            lambda source: pytest.fail("must not reach the database"),
+        )
+        tools = {}
+        add_graph_search_tool(tools, {"source": SOURCE})
+        assert tools == {}
+
+    def test_actions_report_it_rather_than_querying(self, monkeypatch):
+        monkeypatch.setattr(settings, "GRAPHRAG_ENABLED", True)
+        monkeypatch.setattr(settings, "VECTOR_STORE", "faiss")
+        tool = GraphSearchTool({"source": SOURCE})
+        tool._store = _StubStore(nodes=[{"name": "Quill", "distance": 0.1}])
+
+        assert "not enabled" in tool.execute_action("search_entities", query="quill")
+
+
+class TestPooledConnection:
+    """The tool is cached for the whole agent run; its connection must not be."""
+
+    class _ClosingStore(_StubStore):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self.closed = 0
+
+        def close(self):
+            self.closed += 1
+
+    def test_the_connection_goes_back_after_each_action(self, monkeypatch):
+        store = self._ClosingStore(relationships=[{"source": "A", "target": "B", "type": "r"}])
+        tool = _tool(monkeypatch, store)
+
+        tool.execute_action("get_relationships", entity="A")
+
+        # Held open, one pooled connection would be pinned across every LLM
+        # round trip of the run.
+        assert store.closed == 1
+        assert tool._store is None
+
+    def test_a_failing_action_still_releases_it(self, monkeypatch):
+        class _Broken(self._ClosingStore):
+            def entity_relationships(self, source_id, name, limit=25):
+                raise RuntimeError("connection lost")
+
+        store = _Broken()
+        tool = _tool(monkeypatch, store)
+
+        assert tool.execute_action("get_relationships", entity="A") == "The graph lookup failed."
+        assert store.closed == 1
+        assert tool._store is None