diff --git a/graphify/serve.py b/graphify/serve.py index 1c9b979b6..23c140427 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -752,6 +752,99 @@ def _find_node(G: nx.Graph, label: str) -> list[str]: return source_exact + 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. @@ -797,13 +890,19 @@ def _community_header(cid: int, community_name) -> str: return base -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 @@ -816,6 +915,17 @@ def _build_server(graph_path: str): from graphify import paths as _paths +# 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() + # Per-graph context cache: resolved graph.json path -> {key, G, communities}. # The server's default graph is just the first entry; a tool call carrying a # project_path adds its own. Routing every graph through one cache means the @@ -918,6 +1028,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.", @@ -1070,6 +1201,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() @@ -1276,6 +1446,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, @@ -1374,7 +1545,7 @@ async def call_tool(name: str, arguments: dict) -> list[types.TextContent]: return server -def serve(graph_path: str | None = None) -> None: +def serve(graph_path: str | None = None, source_root: str | None = None) -> None: """Start the MCP server over stdio (the default, per-developer transport).""" graph_path = graph_path or _default_graph_json() try: @@ -1383,7 +1554,7 @@ def serve(graph_path: str | None = None) -> 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: @@ -1460,6 +1631,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. @@ -1490,7 +1662,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 @@ -1542,6 +1714,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). @@ -1574,6 +1747,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)" @@ -1641,6 +1815,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 _default_graph_json() @@ -1654,9 +1834,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 9ef7ac0fc..02811bfd7 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -25,7 +25,8 @@ _resolve_context_filters, _subgraph_to_text, _load_graph, - _community_header, +_community_header, + _read_source_slice, ) @@ -915,3 +916,100 @@ def test_community_header_sanitizes_name(): out = _community_header(3, "Pay\x00ments\x1b[31m") assert out.startswith("Community 3 — ") assert "\x00" not in out and "\x1b" not in out + + +# --- _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