heapspy-mcp 0.1.1__py3-none-any.whl

heapspy_mcp/__init__.py

@@ -0,0 +1,3 @@
+"""MCP server that gives coding agents access to heapspy."""
+
+__version__ = "0.1.1"

heapspy_mcp/__main__.py

@@ -0,0 +1,167 @@
+"""Command-line entry point: ``heapspy-mcp``."""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+
+from . import __version__
+from . import config as cpu_config  # CPU profiling (py-spy)
+from .config import (
+    DEFAULT_CAPTURE_TIMEOUT,
+    DEFAULT_TIMEOUT,
+    ENV_ALLOW,
+    ENV_BIN,
+    ENV_CAPTURE_TIMEOUT,
+    ENV_MAX_RSS,
+    ENV_MAX_STR,
+    ENV_SNAPSHOT_DIR,
+    ENV_TIMEOUT,
+    Config,
+    ConfigError,
+    add_diagnostics_arguments,
+    default_snapshot_dir,
+    diagnostics_from_args,
+)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="heapspy-mcp",
+        description="MCP server exposing heapspy heap snapshots and analysis to coding agents.",
+    )
+    p.add_argument("--version", action="version", version=f"heapspy-mcp {__version__}")
+    p.add_argument(
+        "--transport",
+        choices=["stdio", "streamable-http"],
+        default="stdio",
+        help="MCP transport (default: stdio)",
+    )
+    p.add_argument("--host", default="127.0.0.1", help="streamable-http bind address")
+    p.add_argument("--port", type=int, default=8000, help="streamable-http port")
+    p.add_argument(
+        "--heapspy-bin",
+        metavar="PATH",
+        help=f"heapspy executable (else ${ENV_BIN}, else 'heapspy' on PATH)",
+    )
+    p.add_argument(
+        "--allow",
+        action="append",
+        metavar="PATTERN",
+        help=(
+            "allow live capture of processes matching PATTERN: a glob on the process name, "
+            "executable name or command line, or pid:N. Repeatable; comma-separated lists "
+            f"accepted. Merged with ${ENV_ALLOW}. Default: none (live capture disabled)."
+        ),
+    )
+    p.add_argument(
+        "--max-str",
+        type=int,
+        metavar="N",
+        help=(
+            "default per-string cap, in characters, for capture_contents=true when the "
+            f"tool call gives no max_str (also ${ENV_MAX_STR}; default 1024)"
+        ),
+    )
+    p.add_argument(
+        "--max-rss",
+        metavar="SIZE",
+        help=f"refuse live capture above this RSS, e.g. 8G (also ${ENV_MAX_RSS}; default 8G)",
+    )
+    p.add_argument(
+        "--snapshot-dir",
+        metavar="DIR",
+        help=f"where snapshots are written (also ${ENV_SNAPSHOT_DIR}; "
+        f"default {default_snapshot_dir()})",
+    )
+    p.add_argument(
+        "--timeout",
+        type=float,
+        metavar="SECONDS",
+        help=f"timeout for analysis commands (also ${ENV_TIMEOUT}; default {DEFAULT_TIMEOUT:g})",
+    )
+    p.add_argument(
+        "--capture-timeout",
+        type=float,
+        metavar="SECONDS",
+        help=f"timeout for heapspy dump (also ${ENV_CAPTURE_TIMEOUT}; "
+        f"default {DEFAULT_CAPTURE_TIMEOUT:g})",
+    )
+    add_diagnostics_arguments(p)
+    # --- CPU profiling (py-spy) ---
+    p.add_argument(
+        "--py-spy-bin",
+        metavar="PATH",
+        help=f"py-spy executable for the CPU tools (else ${cpu_config.ENV_PY_SPY_BIN}, "
+        "else 'py-spy' on PATH or next to this Python)",
+    )
+    p.add_argument(
+        "--cpu-max-duration",
+        type=int,
+        metavar="SECONDS",
+        help=f"longest CPU profile a tool may record (also ${cpu_config.ENV_CPU_MAX_DURATION}; "
+        f"default {cpu_config.DEFAULT_CPU_MAX_DURATION})",
+    )
+    # --- end CPU profiling ---
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(logging.Formatter("heapspy-mcp: %(levelname)s %(message)s"))
+    log = logging.getLogger("heapspy_mcp")
+    log.addHandler(handler)
+    log.setLevel(logging.INFO)
+    log.propagate = False
+
+    try:
+        config = Config.from_sources(
+            heapspy_bin=args.heapspy_bin,
+            snapshot_dir=args.snapshot_dir,
+            allow=args.allow,
+            max_str=args.max_str,
+            max_rss=args.max_rss,
+            timeout=args.timeout,
+            capture_timeout=args.capture_timeout,
+        )
+        config.diagnostics = diagnostics_from_args(args)
+        cpu_config.load_cpu_settings(
+            config, py_spy_bin=args.py_spy_bin, cpu_max_duration=args.cpu_max_duration
+        )
+    except ConfigError as e:
+        print(f"heapspy-mcp: error: {e}", file=sys.stderr)
+        return 2
+
+    if config.heapspy_bin is None:
+        log.warning(
+            "heapspy executable not found (no --heapspy-bin, no $%s, not on PATH); "
+            "analysis and capture tools will return binary_not_found",
+            ENV_BIN,
+        )
+    log.info(
+        "heapspy=%s snapshot_dir=%s allow=%s default_max_str=%s max_rss=%d",
+        config.heapspy_bin,
+        config.snapshot_dir,
+        config.allow or "(none: live capture disabled)",
+        config.max_str,
+        config.max_rss_bytes,
+    )
+
+    if config.py_spy_bin is None:
+        log.warning("py-spy not found; the CPU tools will return py_spy_not_found")
+    log.info("py_spy=%s cpu_max_duration=%ds", config.py_spy_bin, config.cpu_max_duration)
+
+    from .server import build_server
+
+    server = build_server(config)
+    if args.transport == "stdio":
+        server.run("stdio")
+    else:
+        server.run("streamable-http", host=args.host, port=args.port)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

heapspy_mcp/allocations.py

@@ -0,0 +1,347 @@
+"""Bounded summaries of heapspy allocation data (``.hpya`` profiles).
+
+Two sources:
+
+- ``analyze SNAP --allocations FILE --json``: the ``allocations`` block
+  (recording header, coverage, ``top_stacks``) and ``allocated_at`` on each
+  top retainer; ``query --type ... --allocations`` fills ``allocated_at`` on
+  instances. See docs/json-schema.md, "Allocations".
+- ``flamegraph --allocations FILE --folded``: the only output heapspy has for
+  a profile without a snapshot. It gives estimated bytes per stack and
+  nothing else (no sample counts, event kinds or recording header).
+
+Frames are innermost first everywhere in the output, as in the JSON report.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+from .summarize import MAX_NAME, clip, notes, with_human
+
+MAX_FRAMES = 6
+MAX_STACK_HINTS = 3
+MAX_ALLOCATED_AT_STACKS = 3
+MAX_SUMMARY_STACKS = 5
+MAX_FOLDED_LINES = 100_000
+MAX_LOCATION = 160
+
+CAVEATS = (
+    "Sampled, not complete: only arena, pool and large-object events made while recording "
+    "are sampled. Blocks reused from freed memory, and objects allocated before recording "
+    "started, have no stack.",
+    "Arena and pool links are approximate: they name the most recent growth event that "
+    "created the memory, not necessarily the call that allocated the object. "
+    "exact_linked_bytes counts links through large-object samples, which are exact: the "
+    "object is the sampled allocation itself. Objects on memory a sampled large block "
+    "left behind, which they do not match, are not linked (coverage.unlinked_reused_*).",
+    "estimated_bytes is the sum of sample weights over the recording window: an estimate "
+    "of bytes allocated through sampled events, not of bytes still alive. linked_bytes is "
+    "what is still alive in the snapshot.",
+)
+
+_IDENT_PATH_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(/[A-Za-z_][A-Za-z0-9_]*)*\.py$")
+_FOLDED_FRAME_RE = re.compile(r"^(?P<function>.*) \((?P<location>[^()]*)\)$")
+
+
+# -- frames and hints ----------------------------------------------------------------
+
+
+def _module_guess(file: str) -> str | None:
+    """``app/session.py`` -> ``app.session`` for a relative path of identifiers."""
+    path = file.replace("\\", "/").removeprefix("./")
+    if not _IDENT_PATH_RE.match(path):
+        return None
+    module = path.removesuffix(".py").replace("/", ".")
+    return module.removesuffix(".__init__") or None
+
+
+def qualname_guess(function: str | None, file: str | None) -> str | None:
+    """A ``module.qualname`` for ``map_to_source`` from a frame, when one can be guessed.
+
+    Only for relative ``co_filename`` values made of identifiers; the module
+    is then the path. For an absolute filename the repository root is not
+    known, so there is no guess: read the file at the line instead.
+    """
+    if not isinstance(file, str) or not isinstance(function, str):
+        return None
+    module = _module_guess(file)
+    if module is None:
+        return None
+    if function == "<module>":
+        return module
+    parts = function.split(".")
+    if not all(p.isidentifier() or p == "<locals>" for p in parts) or parts[0] == "<locals>":
+        return None
+    return f"{module}.{function}"
+
+
+def clip_left(text: Any, limit: int = MAX_LOCATION) -> Any:
+    """Clip a path from the left: the file name and line are the useful end."""
+    if not isinstance(text, str) or len(text) <= limit:
+        return text
+    return "..." + text[-(limit - 3) :]
+
+
+def _frame(frame: dict[str, Any]) -> dict[str, Any]:
+    line = frame.get("line")
+    approximate = frame.get("approximate")
+    return {
+        "function": clip(frame.get("function"), MAX_NAME),
+        "file": clip_left(frame.get("file")),
+        "line": line if isinstance(line, int) else None,
+        "location": clip_left(frame.get("location")),
+        # None in folded output, which does not carry the flag.
+        "approximate": approximate if isinstance(approximate, bool) else None,
+    }
+
+
+def stack_source_hints(frames: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Where to look, innermost frame first: ``file:line`` and a qualname guess.
+
+    The first hint is the allocating function itself; the next ones are its
+    callers, which are often where the decision to keep the object is made.
+    """
+    hints: list[dict[str, Any]] = []
+    for f in frames:
+        if len(hints) >= MAX_STACK_HINTS:
+            break
+        file, line = f.get("file"), f.get("line")
+        if not isinstance(file, str) or not file or file.startswith("<"):
+            continue
+        hints.append(
+            {
+                "location": f.get("location") or f"{file}:{line if line is not None else '?'}",
+                "file": file,
+                "line": line,
+                "function": f.get("function"),
+                "qualname_guess": qualname_guess(f.get("function"), file),
+            }
+        )
+    return hints
+
+
+def _frames(raw: Any) -> tuple[list[dict[str, Any]], int]:
+    frames = [f for f in raw or [] if isinstance(f, dict)]
+    return [_frame(f) for f in frames[:MAX_FRAMES]], len(frames)
+
+
+def _stack_common(frames_raw: Any) -> dict[str, Any]:
+    frames, total = _frames(frames_raw)
+    out: dict[str, Any] = {
+        "allocated_in": frames[0]["location"] if frames else None,
+        "frames": frames,
+        "frames_total": total,
+        "source_hints": stack_source_hints(frames),
+    }
+    return out
+
+
+# -- analyze / query reports ---------------------------------------------------------
+
+
+def summarize_top_stack(stack: dict[str, Any]) -> dict[str, Any]:
+    item: dict[str, Any] = {
+        "rank": stack.get("rank"),
+        "estimated_bytes": stack.get("estimated_bytes"),
+        "samples": stack.get("samples"),
+        "events": list(stack.get("events") or [])[:8],
+        "linked_bytes": stack.get("linked_bytes"),
+        "exact_linked_bytes": stack.get("exact_linked_bytes"),
+        "approximate": stack.get("approximate"),
+        **_stack_common(stack.get("frames")),
+    }
+    return with_human(item, "estimated_bytes", "linked_bytes", "exact_linked_bytes")
+
+
+def summarize_allocated_at(
+    at: dict[str, Any] | None, max_stacks: int = MAX_ALLOCATED_AT_STACKS
+) -> dict[str, Any] | None:
+    """``allocated_at`` of a retainer or instance, with at most ``max_stacks`` stacks."""
+    if not isinstance(at, dict):
+        return None
+    raw = [s for s in at.get("stacks") or [] if isinstance(s, dict)]
+    stacks = []
+    for s in raw[:max_stacks]:
+        item = {
+            "linked_objects": s.get("linked_objects"),
+            "linked_bytes": s.get("linked_bytes"),
+            "exact_linked_bytes": s.get("exact_linked_bytes"),
+            "top_stack_rank": s.get("top_stack_rank"),
+            **_stack_common(s.get("frames")),
+        }
+        stacks.append(with_human(item, "linked_bytes"))
+    out = {
+        "linked_objects": at.get("linked_objects"),
+        "linked_bytes": at.get("linked_bytes"),
+        "exact_linked_bytes": at.get("exact_linked_bytes"),
+        "approximate": at.get("approximate"),
+        "stacks": stacks,
+        "stacks_total": len(raw),
+    }
+    return with_human(out, "linked_bytes", "exact_linked_bytes")
+
+
+def allocation_notes(items: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
+    """The notes about the profile (``allocations_*`` codes)."""
+    return [
+        n
+        for n in notes(items)
+        if isinstance(n.get("code"), str) and n["code"].startswith("allocations_")
+    ]
+
+
+def _not_linked_warning(note_list: list[dict[str, Any]]) -> str | None:
+    if any(n.get("code") == "allocations_not_linked" for n in note_list):
+        return (
+            "No object was linked to a sample: the profile does not match the snapshot (see "
+            "the allocations_* notes: another pid, another CPython version, or a snapshot "
+            "taken before recording started). top_stacks still shows estimated bytes, but "
+            "linked_bytes is 0 everywhere."
+        )
+    return None
+
+
+def summarize_allocations(
+    block: dict[str, Any] | None, doc_notes: list[dict[str, Any]] | None, top: int
+) -> dict[str, Any] | None:
+    """The ``allocations`` block of an analyze report, bounded to ``top`` stacks."""
+    if not isinstance(block, dict):
+        return None
+    sampling = block.get("sampling") if isinstance(block.get("sampling"), dict) else {}
+    overhead = block.get("overhead") if isinstance(block.get("overhead"), dict) else {}
+    coverage = dict(block.get("coverage") or {})
+    with_human(coverage, "linked_bytes", "exact_linked_bytes", "unlinked_reused_bytes")
+    recording = {
+        "path": clip_left(block.get("path"), 400),
+        "pid": block.get("pid"),
+        "python_version": block.get("python_version"),
+        "started_at": block.get("started_at"),
+        "duration_ms": block.get("duration_ms"),
+        "mechanisms": list(block.get("mechanisms") or [])[:8],
+        "sampling": {
+            "events": list(sampling.get("events") or [])[:8],
+            "byte_interval": sampling.get("byte_interval"),
+            "max_frames": sampling.get("max_frames"),
+        },
+        "overhead": {
+            k: overhead.get(k)
+            for k in ("samples_taken", "samples_dropped", "probe_hits", "total_pause_ns")
+        },
+        "truncated": block.get("truncated"),
+        "approximate": block.get("approximate"),
+        "samples": block.get("samples"),
+        "estimated_bytes": block.get("estimated_bytes"),
+    }
+    with_human(recording, "estimated_bytes")
+    stacks = [s for s in block.get("top_stacks") or [] if isinstance(s, dict)]
+    note_list = allocation_notes(doc_notes)
+    out: dict[str, Any] = {
+        "recording": recording,
+        "coverage": coverage,
+        "top_stacks": [summarize_top_stack(s) for s in stacks[:top]],
+        "top_stacks_listed": len(stacks),
+        "notes": note_list,
+        "caveats": list(CAVEATS),
+    }
+    warning = _not_linked_warning(note_list)
+    if warning:
+        out["warning"] = warning
+    return out
+
+
+def add_to_analyze(summary: dict[str, Any], doc: dict[str, Any], top: int) -> dict[str, Any]:
+    """Add the allocation join of an analyze report to its summary (in place)."""
+    block = doc.get("allocations")
+    if not isinstance(block, dict):
+        return summary
+    raw = [r for r in doc.get("top_retainers") or [] if isinstance(r, dict)]
+    for item, r in zip(summary.get("top_retainers") or [], raw, strict=False):
+        item["allocated_at"] = summarize_allocated_at(r.get("allocated_at"))
+    summary["allocations"] = summarize_allocations(
+        block, doc.get("notes"), min(top, MAX_SUMMARY_STACKS)
+    )
+    return summary
+
+
+def add_to_type_query(summary: dict[str, Any], doc: dict[str, Any]) -> dict[str, Any]:
+    """Add ``allocated_at`` of each instance of a ``query --type --allocations`` report."""
+    raw = [i for i in doc.get("top_instances") or [] if isinstance(i, dict)]
+    for item, inst in zip(summary.get("top_instances") or [], raw, strict=False):
+        item["allocated_at"] = summarize_allocated_at(inst.get("allocated_at"))
+    note_list = allocation_notes(doc.get("notes"))
+    out: dict[str, Any] = {"notes": note_list, "caveats": list(CAVEATS)}
+    warning = _not_linked_warning(note_list)
+    if warning:
+        out["warning"] = warning
+    summary["allocations"] = out
+    return summary
+
+
+# -- profile only: folded stacks -----------------------------------------------------
+
+
+def _parse_folded_frame(text: str) -> dict[str, Any]:
+    m = _FOLDED_FRAME_RE.match(text)
+    if not m:
+        # "[no Python stack]" and anything unexpected.
+        return {"function": text, "file": None, "line": None, "location": None}
+    function, location = m.group("function"), m.group("location")
+    file: str | None = None
+    line: int | None = None
+    if location != "<unknown>":
+        head, sep, tail = location.rpartition(":")
+        if sep and (tail.isdigit() or tail == "?"):
+            file = head
+            line = int(tail) if tail.isdigit() else None
+        else:
+            file = location
+    return {"function": function, "file": file, "line": line, "location": location}
+
+
+def parse_folded(text: str) -> list[tuple[list[dict[str, Any]], int]]:
+    """``outer;...;inner bytes`` lines as (frames innermost first, bytes)."""
+    out: list[tuple[list[dict[str, Any]], int]] = []
+    for raw in text.splitlines()[:MAX_FOLDED_LINES]:
+        line = raw.rstrip()
+        stack, sep, weight = line.rpartition(" ")
+        if not sep or not weight.isdigit() or not stack:
+            continue
+        frames = [_parse_folded_frame(f) for f in stack.split(";")]
+        frames.reverse()
+        out.append((frames, int(weight)))
+    return out
+
+
+def summarize_folded(text: str, top: int) -> dict[str, Any]:
+    """Top stacks of a profile from its folded stacks, merged by identical frames."""
+    merged: dict[tuple[tuple[Any, ...], ...], int] = {}
+    order: dict[tuple[tuple[Any, ...], ...], list[dict[str, Any]]] = {}
+    for frames, weight in parse_folded(text):
+        key = tuple((f["function"], f["location"]) for f in frames)
+        merged[key] = merged.get(key, 0) + weight
+        order.setdefault(key, frames)
+    ranked = sorted(merged.items(), key=lambda kv: -kv[1])
+    total = sum(merged.values())
+    stacks = []
+    for rank, (key, weight) in enumerate(ranked[:top], start=1):
+        frames = [{**f, "approximate": None} for f in order[key]]
+        item: dict[str, Any] = {
+            "rank": rank,
+            "estimated_bytes": weight,
+            "estimated_pct": round(weight / total * 100, 2) if total else None,
+            # Not in folded output; available with a snapshot.
+            "samples": None,
+            "events": None,
+            "linked_bytes": None,
+            "exact_linked_bytes": None,
+            **_stack_common(frames),
+        }
+        stacks.append(with_human(item, "estimated_bytes"))
+    out = {
+        "estimated_bytes": total,
+        "stacks_total": len(merged),
+        "top_stacks": stacks,
+    }
+    return with_human(out, "estimated_bytes")