From 804d99e07cfbee55fb5a600dfa4188c4c0a23411 Mon Sep 17 00:00:00 2001 From: ningguoqiang Date: Thu, 18 Jun 2026 16:36:55 +0800 Subject: [PATCH] feat(serve): add read_source MCP tool to read original source via shared server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shared HTTP MCP server exposed only the graph.json index — graph queries returned a source_file + source_location pointer but no file content, so clients had to fall back to a local checkout. read_source makes a single shared deployment serve both query and read-original. - read_source tool: node label (resolves source_file + source_location, context window) or repo-relative file + optional start/end range. - Core logic in module-level _read_source_slice (unit-testable without mcp). - New --source-root flag (env GRAPHIFY_SOURCE_ROOT) anchors repo-relative source_file paths for containerised deployments. - Security (F-010): reject absolute paths; enforce resolve() containment; sanitize_label every line; hard cap 500 lines / 64 KiB. - 10 new tests in tests/test_serve.py. --- graphify/serve.py | 191 ++++++++++++++++++++++++++++++++++++++++++-- tests/test_serve.py | 99 +++++++++++++++++++++++ 2 files changed, 285 insertions(+), 5 deletions(-) diff --git a/graphify/serve.py b/graphify/serve.py index e3204ffca..80b82c926 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -480,6 +480,99 @@ def _find_node(G: nx.Graph, label: str) -> list[str]: return exact + prefix + substring +# Caps for read_source output — source content is corpus-derived and may be +# huge or attacker-crafted, so every response is line- and byte-bounded. +_READ_SOURCE_MAX_LINES = 500 +_READ_SOURCE_MAX_BYTES = 64 * 1024 +_READ_SOURCE_DEFAULT_WINDOW = 50 + + +def _read_source_slice( + source_root: Path, + source_file: str, + start: int | None = None, + end: int | None = None, + *, + max_lines: int = _READ_SOURCE_MAX_LINES, + max_bytes: int = _READ_SOURCE_MAX_BYTES, + default_window: int = _READ_SOURCE_DEFAULT_WINDOW, +) -> str: + """Read and format a line slice of an original source file. + + ``source_file`` is repo-root-relative (POSIX, as stored in graph.json); + ``source_root`` anchors it. The resolved path must stay inside + ``source_root`` (blocks ``../`` and symlink escapes). Every emitted line is + passed through :func:`sanitize_label` because source content is + corpus-derived / attacker-controllable (F-010). Returns a friendly message + string for expected misses (absolute path, escape, missing file, empty), + never raises on those. + """ + if not source_file: + return "No source_file to read." + + # source_file is meant to be repo-relative; reject absolute client input. + # (resolve() below still enforces containment regardless.) + if Path(source_file).is_absolute(): + return "read_source requires a repo-relative path, not an absolute path." + + try: + start = max(1, int(start)) if start is not None else 1 + except (TypeError, ValueError): + start = 1 + if end is None: + end = start + default_window + try: + end = int(end) + except (TypeError, ValueError): + end = start + default_window + if end < start: + end = start + + full = (source_root / source_file).resolve() + try: + full.relative_to(source_root) + except ValueError: + return ( + f"Path '{source_file}' escapes the server's source-root " + f"({source_root}). read_source is confined to the repo." + ) + if not full.is_file(): + return ( + f"Source file not found on server: {source_file}. The server's " + f"source-root is {source_root}; ensure the repo is mounted or " + f"checked out there (e.g. --source-root /repo in a container)." + ) + + text = full.read_text(encoding="utf-8", errors="replace") + lines = text.splitlines() + total = len(lines) + if total == 0: + return f"File: {source_file} (empty)" + + slice_start = max(1, start) + slice_end = min(end, total) + out_lines: list[str] = [] + total_bytes = 0 + rendered = 0 + truncated = False + for n in range(slice_start, slice_end + 1): + if rendered >= max_lines: + truncated = True + break + line = f"{n:>6}: {sanitize_label(lines[n - 1])}" + if total_bytes + len(line) + 1 > max_bytes: + truncated = True + break + out_lines.append(line) + total_bytes += len(line) + 1 + rendered += 1 + + header = f"File: {source_file} (lines {slice_start}-{slice_end} of {total})" + if truncated: + header += f" — truncated at {max_lines} lines / {max_bytes // 1024} KiB cap" + return header + "\n" + "\n".join(out_lines) + + def _filter_blank_stdin() -> None: """Filter blank lines from stdin before MCP reads it. @@ -510,13 +603,19 @@ def _relay() -> None: sys.stdin = open(0, "r", closefd=False) -def _build_server(graph_path: str): +def _build_server(graph_path: str, source_root: str | None = None): """Build the configured low-level MCP Server (shared by every transport). All graph query tools and resources are registered here over a single ``mcp.server.Server`` instance; the caller picks the transport (stdio or Streamable HTTP) and runs it. Hot-reload of graph.json works the same way regardless of transport, since reloads happen inside the tool handlers. + + ``source_root`` is the directory ``read_source`` resolves repo-relative + ``source_file`` paths against. When omitted it is inferred: the parent of + ``graphify-out/`` if the graph lives there, else the current working + directory. Pass an explicit path for containerised deployments where the + repo is mounted somewhere other than next to ``graphify-out``. """ import threading @@ -530,6 +629,17 @@ def _build_server(graph_path: str): G = _load_graph(graph_path) communities = _communities_from_graph(G) + # Directory read_source resolves source_file paths against. source_file is + # stored POSIX repo-root-relative (extract.py relativises it), and the repo + # root is not persisted in graph.json, so we infer it here unless told. + _graphify_out = Path(graph_path).resolve().parent + if source_root: + _source_root = Path(source_root).resolve() + elif _graphify_out.name == "graphify-out": + _source_root = _graphify_out.parent + else: + _source_root = Path.cwd().resolve() + # Hot-reload state: mtime+size key lets us detect graph.json changes without # polling. Initialised from the file stat at startup so the first tool call # never triggers a redundant reload. @@ -599,6 +709,27 @@ async def list_tools() -> list[types.Tool]: "required": ["label"], }, ), + types.Tool( + name="read_source", + description=( + "Read a slice of an original source file (text files only). " + "Pass either `label` (a graph node name — resolves its source_file " + "+ source_location and returns a window of context lines around " + "that location) or `file` (a repo-relative source_file path, e.g. " + "'src/auth/x.py') with optional start/end line numbers. Source " + "files must be reachable from the server's source-root." + ), + inputSchema={ + "type": "object", + "properties": { + "label": {"type": "string", "description": "Node label — resolves source_file + source_location from the graph"}, + "file": {"type": "string", "description": "Repo-relative source_file path (POSIX, e.g. src/auth/x.py)"}, + "start": {"type": "integer", "description": "1-indexed start line (inclusive). Defaults to 1."}, + "end": {"type": "integer", "description": "1-indexed end line (inclusive). Defaults to start + window."}, + "context": {"type": "integer", "default": 20, "description": "Lines of context around a label-resolved location"}, + }, + }, + ), types.Tool( name="get_neighbors", description="Get all direct neighbors of a node with edge details.", @@ -737,6 +868,45 @@ def _tool_get_node(arguments: dict) -> str: f" Degree: {G.degree(nid)}", ]) + def _tool_read_source(arguments: dict) -> str: + # Read a slice of an original source file. source_file is stored + # repo-root-relative (POSIX); _source_root (inferred or --source-root) + # anchors it. The path-resolve / slice / sanitise logic lives in the + # module-level _read_source_slice helper so it can be tested without + # the mcp package; this closure just resolves a node label to its + # source_file + source_location first. + label = arguments.get("label") + file_arg = arguments.get("file") + + if not label and not file_arg: + return "read_source requires 'label' or 'file'." + + if label: + matches = _find_node(G, label) + if not matches: + return f"No node matching '{label}' found." + d = G.nodes[matches[0]] + source_file = d.get("source_file") or "" + if not source_file: + return f"Node '{label}' has no source_file to read." + # source_location is "L{n}"; centre the window on it. + loc = str(d.get("source_location") or "") + m = re.search(r"\d+", loc) + center = int(m.group()) if m else 1 + context = arguments.get("context", 20) or 20 + try: + context = max(1, int(context)) + except (TypeError, ValueError): + context = 20 + start = max(1, center - context) + end = center + context + else: + source_file = str(file_arg) + start = arguments.get("start") + end = arguments.get("end") + + return _read_source_slice(_source_root, source_file, start, end) + def _tool_get_neighbors(arguments: dict) -> str: label = arguments["label"].lower() rel_filter = arguments.get("relation_filter", "").lower() @@ -942,6 +1112,7 @@ def _tool_triage_prs(arguments: dict) -> str: _handlers = { "query_graph": _tool_query_graph, "get_node": _tool_get_node, + "read_source": _tool_read_source, "get_neighbors": _tool_get_neighbors, "get_community": _tool_get_community, "god_nodes": _tool_god_nodes, @@ -1038,7 +1209,7 @@ async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: return server -def serve(graph_path: str = "graphify-out/graph.json") -> None: +def serve(graph_path: str = "graphify-out/graph.json", source_root: str | None = None) -> None: """Start the MCP server over stdio (the default, per-developer transport).""" try: from mcp.server.stdio import stdio_server @@ -1046,7 +1217,7 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None: raise ImportError('mcp not installed. Run: pip install "graphifyy[mcp]"') from e import asyncio - server = _build_server(graph_path) + server = _build_server(graph_path, source_root) async def main() -> None: async with stdio_server() as streams: @@ -1123,6 +1294,7 @@ def _build_http_app( json_response: bool = False, stateless: bool = False, session_timeout: float | None = 3600.0, + source_root: str | None = None, ): """Build the Starlette ASGI app for the Streamable HTTP transport. @@ -1153,7 +1325,7 @@ def _build_http_app( # mistaken for "auth on" — normalize it to None so the gate is unambiguous. api_key = (api_key or "").strip() or None - server = _build_server(graph_path) + server = _build_server(graph_path, source_root) # DNS-rebinding protection. When the operator binds a wildcard address they # are intentionally exposing the server, so accept any Host header; for a @@ -1205,6 +1377,7 @@ def serve_http( json_response: bool = False, stateless: bool = False, session_timeout: float | None = 3600.0, + source_root: str | None = None, ) -> None: """Start the MCP server over Streamable HTTP (MCP spec 2025-03-26). @@ -1236,6 +1409,7 @@ def serve_http( json_response=json_response, stateless=stateless, session_timeout=session_timeout, + source_root=source_root, ) auth_note = "api-key required" if api_key else "no auth (set --api-key to require one)" @@ -1303,6 +1477,12 @@ def _main(argv: list[str] | None = None) -> None: default=3600.0, help="Reap stateful sessions idle this many seconds (default: 3600; 0 disables)", ) + parser.add_argument( + "--source-root", + default=os.environ.get("GRAPHIFY_SOURCE_ROOT"), + help="Root directory read_source resolves source_file paths against " + "(default: parent of graphify-out, or CWD). env: GRAPHIFY_SOURCE_ROOT", + ) args = parser.parse_args(argv) graph_path = args.graph_flag or args.graph_path or "graphify-out/graph.json" @@ -1316,9 +1496,10 @@ def _main(argv: list[str] | None = None) -> None: json_response=args.json_response, stateless=args.stateless, session_timeout=args.session_timeout, + source_root=args.source_root, ) else: - serve(graph_path) + serve(graph_path, source_root=args.source_root) if __name__ == "__main__": diff --git a/tests/test_serve.py b/tests/test_serve.py index 5b5926b96..2d8897346 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -19,6 +19,7 @@ _resolve_context_filters, _subgraph_to_text, _load_graph, + _read_source_slice, ) @@ -577,3 +578,101 @@ def test_query_text_chinese_finds_routing_nodes(): text = _query_graph_text(G, "页面路由", mode="bfs", depth=2) assert "No matching nodes found." not in text assert "路由" in text + + +# --- _read_source_slice (read_source tool core) --- + +def _write_source(tmp_path, rel, lines): + """Write a file under tmp_path and return its POSIX repo-relative path.""" + p = tmp_path / rel + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text("\n".join(lines) + "\n", encoding="utf-8") + return rel.replace("\\", "/") + + +def test_read_source_slice_file_range(tmp_path): + rel = _write_source(tmp_path, "src/auth/x.py", [f"line {i}" for i in range(1, 11)]) + out = _read_source_slice(tmp_path, rel, start=2, end=4) + assert "File: src/auth/x.py (lines 2-4 of 10)" in out + assert " 2: line 2" in out + assert " 4: line 4" in out + assert "line 1" not in out # before the range + assert "line 5" not in out # after the range + + +def test_read_source_slice_default_window(tmp_path): + rel = _write_source(tmp_path, "doc.md", [f"l{i}" for i in range(1, 200)]) + out = _read_source_slice(tmp_path, rel, start=10) # end defaults to start + 50 + assert "(lines 10-60 of 199)" in out + + +def test_read_source_slice_path_traversal_rejected(tmp_path): + _write_source(tmp_path, "escape.py", ["secret"]) + # ../escape.py resolves outside source_root -> must be refused. + out = _read_source_slice(tmp_path, "../escape.py") + assert "escapes the server's source-root" in out + assert "secret" not in out + + +def test_read_source_slice_absolute_path_rejected(tmp_path): + _write_source(tmp_path, "x.py", ["content"]) + out = _read_source_slice(tmp_path, str((tmp_path / "x.py").resolve())) + assert "repo-relative path" in out + assert "content" not in out + + +def test_read_source_slice_truncation_cap(tmp_path): + rel = _write_source(tmp_path, "big.txt", [f"line {i}" for i in range(1, 1001)]) + out = _read_source_slice(tmp_path, rel, start=1, end=1000) + assert "truncated" in out + # The cap is 500 lines; line 600 must not appear. + assert "line 600" not in out + assert "line 1" in out + + +def test_read_source_slice_missing_file_message(tmp_path): + out = _read_source_slice(tmp_path, "does/not/exist.py") + assert "Source file not found on server" in out + assert "--source-root" in out # the Docker-misconfig hint + + +def test_read_source_slice_empty_file(tmp_path): + (tmp_path / "empty.txt").write_text("", encoding="utf-8") + out = _read_source_slice(tmp_path, "empty.txt") + assert out == "File: empty.txt (empty)" + + +def test_read_source_slice_sanitises_control_chars(tmp_path): + # Source content is attacker-controllable (F-010); control chars must be + # stripped from emitted lines so they cannot forge log lines / escapes. + rel = _write_source(tmp_path, "evil.py", ["clean\x1b[31mbad\x00"]) + out = _read_source_slice(tmp_path, rel, start=1, end=1) + assert "\x1b" not in out + assert "\x00" not in out + assert "clean" in out and "bad" in out + + +def test_read_source_slice_start_clamped(tmp_path): + rel = _write_source(tmp_path, "x.py", ["only line"]) + out = _read_source_slice(tmp_path, rel, start=-5, end=1) + assert "(lines 1-1 of 1)" in out + assert "only line" in out + + +def test_read_source_label_resolution_via_find_node(tmp_path): + # End-to-end label resolution: a graph node points at a real file; the + # read_source wrapper resolves label -> source_file + source_location and + # centres a context window. _find_node is the same matcher the wrapper uses. + rel = _write_source(tmp_path, "src/svc.py", [f"line {i}" for i in range(1, 31)]) + G = nx.Graph() + G.add_node("svc", label="UserService", source_file=rel, source_location="L15", community=0) + matches = _find_node(G, "UserService") + assert matches # label resolves + d = G.nodes[matches[0]] + center = 15 + context = 2 + out = _read_source_slice(tmp_path, d["source_file"], center - context, center + context) + assert "(lines 13-17 of 30)" in out + assert "line 15" in out + assert "line 12" not in out # outside the context window +