From fef9dbb5e74580371326050cca7d2b649b556ff4 Mon Sep 17 00:00:00 2001
From: B <8027448+ojmucianski@users.noreply.github.com>
Date: Tue, 21 Jul 2026 01:18:31 +0100
Subject: [PATCH] feat(serve): honor a token_budget on
 get_neighbors/get_community
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

get_neighbors and get_community render every edge/member line unbounded — on a
god node or a large community that floods an MCP client's context window with
100KB+ of text in one tool result. query_graph already solves this with the
~3-chars/token cut in _subgraph_to_text; this applies the same budget rule to
the two line-list tools via a shared helper (_cut_lines_to_budget): cut at a
line boundary, report how many lines were dropped, and point at the narrowing
path (relation_filter / get_node). Default 2000 like query_graph; output under
budget is byte-identical to today.
---
 graphify/serve.py | 36 +++++++++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/graphify/serve.py b/graphify/serve.py
index 20e96472..75798f5b 100644
--- a/graphify/serve.py
+++ b/graphify/serve.py
@@ -908,6 +908,26 @@ def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_bu
     return output
 
 
+def _cut_lines_to_budget(lines: list[str], token_budget: int, narrow_hint: str) -> str:
+    """Render pre-built lines under the same ~3-chars/token budget rule as
+    _subgraph_to_text; over-budget output is cut at a line boundary with a count and a
+    narrowing hint instead of flooding the caller's context window."""
+    output = "\n".join(lines)
+    char_budget = token_budget * 3
+    if len(output) <= char_budget:
+        return output
+    cut_at = output[:char_budget].rfind("\n")
+    cut_at = cut_at if cut_at > 0 else char_budget
+    kept = output[:cut_at]
+    cut_count = len(lines) - kept.count("\n") - 1
+    return (
+        kept
+        + f"\n... (truncated — {cut_count} more lines cut by ~{token_budget}-token budget. "
+        + narrow_hint
+        + ")"
+    )
+
+
 def _query_graph_text(
     G: nx.Graph,
     question: str,
@@ -1189,6 +1209,7 @@ def _build_server(graph_path: str):
                     "properties": {
                         "label": {"type": "string"},
                         "relation_filter": {"type": "string", "description": "Optional: filter by relation type"},
+                        "token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"},
                     },
                     "required": ["label"],
                 },
@@ -1198,7 +1219,10 @@ def _build_server(graph_path: str):
                 description="Get all nodes in a community by community ID.",
                 inputSchema={
                     "type": "object",
-                    "properties": {"community_id": {"type": "integer", "description": "Community ID (0-indexed by size)"}},
+                    "properties": {
+                        "community_id": {"type": "integer", "description": "Community ID (0-indexed by size)"},
+                        "token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"},
+                    },
                     "required": ["community_id"],
                 },
             ),
@@ -1367,7 +1391,10 @@ def _build_server(graph_path: str):
                 f"  <-- {sanitize_label(G.nodes[nb].get('label', nb))} "
                 f"[{sanitize_label(str(rel))}] [{sanitize_label(str(d.get('confidence', '')))}]{_edge_at(d)}"
             )
-        return "\n".join(lines)
+        budget = int(arguments.get("token_budget", 2000))
+        return _cut_lines_to_budget(
+            lines, budget, "Narrow with relation_filter or use get_node for a specific symbol"
+        )
 
     def _tool_get_community(arguments: dict) -> str:
         cid = int(arguments["community_id"])
@@ -1383,7 +1410,10 @@ def _build_server(graph_path: str):
                 f"  {sanitize_label(d.get('label', n))} "
                 f"[{sanitize_label(str(d.get('source_file', '')))}]"
             )
-        return "\n".join(lines)
+        budget = int(arguments.get("token_budget", 2000))
+        return _cut_lines_to_budget(
+            lines, budget, "Raise token_budget or use get_node for specific members"
+        )
 
     def _tool_god_nodes(arguments: dict) -> str:
         from graphify.analyze import god_nodes as _god_nodes