From dd37718181658cedd1c34d40ff65b45b9202919e Mon Sep 17 00:00:00 2001
From: Yurii214 <216080096+Yurii214@users.noreply.github.com>
Date: Thu, 30 Jul 2026 16:03:05 +0000
Subject: [PATCH] fix(server): thread graph_rel_types through the mcp
 kb_context tool
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

the mcp kb_context tool exposes expand_graph/graph_depth/graph_limit but
drops graph_rel_types, which the jsonl _h_context handler threads and
build_context_pack supports. so an mcp agent — the primary vouch client —
can graph-expand a context pack but cannot scope the walk to specific
relation types, unlike a jsonl/http client (and unlike the kb_neighbors
tool, which does expose rel_types on mcp). thread it, mirroring _h_context.
---
 src/vouch/server.py |  4 ++++
 tests/test_graph.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/vouch/server.py b/src/vouch/server.py
index e74235b0..b12fbe11 100644
--- a/src/vouch/server.py
+++ b/src/vouch/server.py
@@ -331,6 +331,9 @@ def kb_context(
     expand_graph: bool = False,
     graph_depth: int = 1,
     graph_limit: int = 20,
+    # restrict which relation types the graph walk follows, mirroring the jsonl
+    # _h_context handler and the kb_neighbors tool; None follows every type.
+    graph_rel_types: list[str] | None = None,
 ) -> dict[str, Any]:
     """Build a ContextPack ready to inject into an agent prompt.
 
@@ -347,6 +350,7 @@ def kb_context(
         min_items=min_items, require_citations=require_citations,
         project=project, agent=agent,
         expand_graph=expand_graph, graph_depth=graph_depth, graph_limit=graph_limit,
+        graph_rel_types=graph_rel_types,
     )
     result = salience_mod.attach_salience(result, store, session_id, cfg)
     pack_items = result.get("items") if isinstance(result, dict) else None
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 8130dc5e..a423b422 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -126,6 +126,38 @@ def test_context_expand_graph_adds_neighbors(store: KBStore) -> None:
     assert any("graph expansion" in w for w in pack["warnings"])
 
 
+def test_kb_context_mcp_honors_graph_rel_types(
+    store: KBStore, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # the mcp kb_context tool must accept graph_rel_types like the jsonl handler
+    # and kb_neighbors do, so an agent can scope the graph walk to specific
+    # relation types. this surface silently dropped the filter.
+    from vouch.server import kb_context
+
+    src = store.put_source(b"e")
+    store.put_entity(Entity(id="auth", name="Auth", type=EntityType.SYSTEM))
+    store.put_entity(Entity(id="risk", name="Risk", type=EntityType.SYSTEM))
+    store.put_claim(Claim(id="jwt-claim", text="JWT tokens secure the API", evidence=[src.id]))
+    store.put_relation(
+        Relation(
+            id="c-ref-auth", source="jwt-claim", relation=RelationType.REFERENCES, target="auth"
+        )
+    )
+    store.put_relation(
+        Relation(
+            id="c-blk-risk", source="jwt-claim", relation=RelationType.BLOCKS, target="risk"
+        )
+    )
+    health.rebuild_index(store)
+    monkeypatch.chdir(store.root)
+
+    pack = kb_context("JWT tokens", limit=5, expand_graph=True, graph_rel_types=["references"])
+
+    ids = {it["id"] for it in pack["items"]}
+    assert "auth" in ids  # the references edge is followed
+    assert "risk" not in ids  # the blocks edge is filtered out by graph_rel_types
+
+
 def test_context_expand_graph_includes_page_claims(store: KBStore) -> None:
     src = store.put_source(b"e")
     store.put_claim(Claim(id="c1", text="detail fact", evidence=[src.id]))