tracy/extra/mcp/tracy_mcp.py

# -*- coding: utf-8 -*-
from __future__ import annotations

import asyncio
import atexit
import builtins
import concurrent.futures
import faulthandler
import glob
import io
import os
import logging
import re
import socket
import struct
import sys
import time
import urllib.error
import urllib.request
import uuid
from contextlib import asynccontextmanager, redirect_stdout

from mcp.server.mcpserver import MCPServer

# Suppress noisy ASGI shutdown errors known to occur with SSE and Control-C.
# These occur when Starlette attempts to send a 500 error after the loop is cancelled
# but after the SSE 200 OK headers have already been sent. Global level suppression
# is used because surgical filtering of ASGI exceptions is unreliable in this stack.
logging.getLogger("uvicorn.error").setLevel(logging.CRITICAL)
logging.getLogger("starlette").setLevel(logging.CRITICAL)

_HERE = os.path.dirname(os.path.abspath(__file__))
_PORT_FILE = os.path.join(_HERE, "tracy_mcp.port")
_PID_FILE  = os.path.join(_HERE, "tracy_mcp.pid")
_CRASH_LOG_FILE = os.path.join(_HERE, "tracy_mcp.crash.log")
_PREFERRED_PORT = int(os.environ.get("TRACY_MCP_PORT", "47380"))
_TRANSPORT = os.environ.get("TRACY_MCP_TRANSPORT", "streamable-http").strip().lower()
# MCPServer.run()'s own defaults -- tracked here too since v2's Settings no
# longer exposes them for us to read back (SDK v1 -> v2 migration).
_SSE_PATH = "/sse"
_STREAMABLE_HTTP_PATH = "/mcp"

# Shared documentation surfaces. system.prompt.md is Tracy Assist's source
# system prompt; exposing it as an MCP resource keeps analysis guidance in
# sync across both surfaces with no plumbing. eval_guide.md covers
# bindings-layer detail (ctx object model, units, source-location ID joins).
_LLM_DIR = os.path.normpath(os.path.join(_HERE, "..", "..", "profiler", "src", "llm"))
_PROMPT_PATH = os.path.join(_LLM_DIR, "system.prompt.md")
_EVAL_GUIDE_PATH = os.path.join(_HERE, "eval_guide.md")


def _read_text(path: str) -> str:
    try:
        with open(path, encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"(unavailable: {e})"


# Tracy UDP broadcast packet support. Tracy clients announce themselves on
# port 8086 with a BroadcastMessage (see public/common/TracyProtocol.hpp).
# The dev GUI reads protocolVersion from the broadcast and refuses connection
# on mismatch instead of hitting an opaque TCP timeout. We do the same.
_PROTOCOL_HPP = os.path.normpath(
    os.path.join(_HERE, "..", "..", "public", "common", "TracyProtocol.hpp")
)
_BROADCAST_PORT = 8086
_PROGRAM_NAME_SIZE = 64


def _read_bindings_protocol_version() -> int | None:
    """Parse ProtocolVersion from TracyProtocol.hpp at startup so our 'expected'
    version stays in sync with the bindings build without extra C++ wiring."""
    try:
        with open(_PROTOCOL_HPP, encoding="utf-8") as f:
            for line in f:
                m = re.search(r"constexpr\s+uint32_t\s+ProtocolVersion\s*=\s*(\d+)", line)
                if m:
                    return int(m.group(1))
    except Exception:
        pass
    return None


_OUR_PROTOCOL_VERSION = _read_bindings_protocol_version()


def _parse_broadcast(data: bytes) -> dict | None:
    """Parse a Tracy BroadcastMessage. Handles broadcast versions 0-3.

    Fixed-field sizes (from TracyProtocol.hpp, packed):
      v3: u16 bv, u16 lp, u32 pv, u64 pid, i32 at, char[<=64] name  (>=20 + name)
      v2: u16 bv, u16 lp, u32 pv, i32 at, char[<=64] name           (>=12 + name)
      v1: u32 bv, u32 pv, u32 lp, u32 at, char[<=64] name           (>=16 + name)
      v0: u32 bv, u32 pv, u32 at, char[<=64] name                   (>=12 + name)

    The programName field is variable-length on the wire — the sender writes
    only the actual name plus null terminator, not the full 64-byte buffer.
    """
    if len(data) < 4:
        return None

    def _name(buf: bytes) -> str:
        return buf[:_PROGRAM_NAME_SIZE].split(b"\0", 1)[0].decode("utf-8", "replace")

    bv16 = struct.unpack_from("<H", data, 0)[0]
    if bv16 == 3 and len(data) >= 21:
        bv, lp, pv, pid, at = struct.unpack_from("<HHIQi", data, 0)
        return {"broadcast_version": bv, "listen_port": lp,
                "protocol_version": pv, "pid": pid,
                "active_seconds": at, "program": _name(data[20:])}
    if bv16 == 2 and len(data) >= 13:
        bv, lp, pv, at = struct.unpack_from("<HHIi", data, 0)
        return {"broadcast_version": bv, "listen_port": lp,
                "protocol_version": pv, "active_seconds": at,
                "program": _name(data[12:])}
    bv32 = struct.unpack_from("<I", data, 0)[0]
    if bv32 == 1 and len(data) >= 17:
        bv, pv, lp, at = struct.unpack_from("<IIII", data, 0)
        return {"broadcast_version": bv, "listen_port": lp,
                "protocol_version": pv, "active_seconds": at,
                "program": _name(data[16:])}
    if bv32 == 0 and len(data) >= 13:
        bv, pv, at = struct.unpack_from("<III", data, 0)
        return {"broadcast_version": bv, "listen_port": None,
                "protocol_version": pv, "active_seconds": at,
                "program": _name(data[12:])}
    return None


async def _listen_broadcasts(timeout_s: float = 1.5) -> list[dict]:
    """Listen briefly on UDP 8086 for Tracy client announcements.

    Returns a list of parsed broadcasts (deduplicated by listen_port). Empty
    list means no broadcast received — the target may use TRACY_ON_DEMAND,
    a non-default broadcast port, or simply isn't running.
    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        s.bind(("", _BROADCAST_PORT))
    except OSError:
        s.close()
        return []
    s.setblocking(False)
    loop = asyncio.get_running_loop()
    seen: dict[int | None, dict] = {}
    deadline = loop.time() + timeout_s
    try:
        while loop.time() < deadline:
            remaining = deadline - loop.time()
            if remaining <= 0:
                break
            try:
                fut = loop.sock_recvfrom(s, 2048)
                data, _addr = await asyncio.wait_for(fut, timeout=remaining)
            except (asyncio.TimeoutError, BlockingIOError):
                break
            parsed = _parse_broadcast(data)
            if parsed:
                seen.setdefault(parsed.get("listen_port"), parsed)
    finally:
        s.close()
    return list(seen.values())


async def _watch_broadcasts_during(
    port: int, deadline_s: float, log_prefix: str
) -> None:
    """Log broadcasts heard on `port` for `deadline_s`. activeTime==-1 means
    the client's Accept() fired for our connection; if it never appears, the
    client's accept loop itself is dead, not just slow to respond.
    """
    start = time.monotonic()
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(("", _BROADCAST_PORT))
    except OSError as e:
        print(f"{log_prefix} broadcast watch: bind failed ({e})", file=sys.stderr, flush=True)
        return
    s.setblocking(False)
    loop = asyncio.get_running_loop()
    deadline = loop.time() + deadline_s
    last_active: int | None = None
    try:
        while loop.time() < deadline:
            remaining = deadline - loop.time()
            if remaining <= 0:
                break
            try:
                fut = loop.sock_recvfrom(s, 2048)
                data, _addr = await asyncio.wait_for(fut, timeout=remaining)
            except (asyncio.TimeoutError, BlockingIOError):
                continue
            parsed = _parse_broadcast(data)
            if not parsed or parsed.get("listen_port") != port:
                continue
            elapsed = time.monotonic() - start
            active = parsed.get("active_seconds")
            transition = ""
            if active == -1 and last_active != -1:
                transition = "  <-- Accept() fired, entering handshake"
            elif last_active == -1 and active != -1:
                transition = "  <-- back in accept-wait loop (handshake attempt ended)"
            last_active = active
            print(
                f"{log_prefix} broadcast @t={elapsed:.2f}s activeTime={active}{transition}",
                file=sys.stderr,
                flush=True,
            )
    finally:
        s.close()
        if last_active is None:
            print(
                f"{log_prefix} broadcast watch: no broadcasts heard in {deadline_s:.1f}s "
                "on this port -- worker thread's accept-wait loop is not ticking at all.",
                file=sys.stderr,
                flush=True,
            )


def _http_ping(port: int, timeout_s: float = 2.0) -> bool:
    """Confirms the server answers HTTP, not just that its PID exists — a
    deadlocked process still passes os.kill(pid, 0). Any response, even an
    error one, proves the transport is alive; only a timeout or refused
    connection means it's not.
    """
    path = _SSE_PATH if _TRANSPORT == "sse" else _STREAMABLE_HTTP_PATH
    try:
        urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=timeout_s)
        return True
    except urllib.error.HTTPError:
        return True
    except Exception:
        return False


def _is_our_server_running() -> tuple[bool, int]:
    """Check whether our server is running: PID alive AND responding to
    an HTTP self-ping (a hung process still passes os.kill(pid, 0)).
    Returns (running, port).
    """
    try:
        with open(_PID_FILE) as f:
            pid = int(f.read().strip())
        with open(_PORT_FILE) as f:
            port = int(f.read().strip())
        os.kill(pid, 0)   # raises OSError if process is gone
    except Exception:
        return False, 0
    if not _http_ping(port):
        print(
            f"Tracy MCP process {pid} is alive (PID check passed) but isn't "
            f"responding on port {port} within {2.0:.0f}s -- likely hung/deadlocked, "
            f"not just busy. Starting a fresh server on a new port; kill PID {pid} "
            f"manually once you're able to (it's still holding the port and the "
            f"TracyServerBindings.pyd file lock).",
            file=sys.stderr,
        )
        return False, 0
    return True, port


def _find_free_port() -> int:
    """Scan from preferred port upward; fall back to OS-assigned if the range is exhausted."""
    for port in range(_PREFERRED_PORT, _PREFERRED_PORT + 16):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind(("127.0.0.1", port))
            s.close()
            return port
        except OSError:
            s.close()
    # Let OS assign any free port
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port


def _write_pid_and_port(port: int) -> None:
    try:
        with open(_PID_FILE, "w") as f:
            f.write(str(os.getpid()))
        with open(_PORT_FILE, "w") as f:
            f.write(str(port))
    except Exception:
        pass


def _cleanup_pid_files() -> None:
    for path in (_PID_FILE, _PORT_FILE):
        try:
            os.unlink(path)
        except Exception:
            pass


# Attempt to import Tracy Server bindings
try:
    import TracyServerBindings as tracy_server
except ImportError:
    _bindings_dir = os.path.join(os.path.dirname(__file__), "../../build/python")
    sys.path.append(_bindings_dir)
    # Multi-config generators (VS) put the .pyd under a per-config subdir
    # (Release/, Debug/, ...) instead of directly in build/python.
    for _config_dir in glob.glob(os.path.join(_bindings_dir, "*")):
        if os.path.isdir(_config_dir) and glob.glob(
            os.path.join(_config_dir, "TracyServerBindings*.pyd")
        ):
            sys.path.append(_config_dir)
    try:
        import TracyServerBindings as tracy_server
    except ImportError:
        tracy_server = None

_MAX_INSTANCES = int(os.environ.get("TRACY_MCP_MAX_INSTANCES", "4"))
_DISCONNECTED_TTL_S = float(os.environ.get("TRACY_MCP_DISCONNECTED_TTL_S", "1800"))
_FILE_IDLE_TTL_S = float(os.environ.get("TRACY_MCP_FILE_IDLE_TTL_S", "1800"))
_SWEEP_INTERVAL_S = float(os.environ.get("TRACY_MCP_SWEEP_INTERVAL_S", "300"))
_DEFAULT_LIVE_MEMORY_LIMIT_MB = int(os.environ.get("TRACY_MCP_LIVE_MEMORY_LIMIT_MB", "8192"))


class Task:
    def __init__(self, task_id: str, code: str):
        self.id = task_id
        self.code = code
        self.status = "pending"
        self.result = None
        self.error = None
        self.start_time = time.time()
        self.end_time = None


class TracyInstance:
    def __init__(self, name: str, worker: object | None = None):
        self.name = name
        self.worker = worker
        self.path = None
        self.mtime = None
        self.last_used = time.time()
        # Wall-clock time the instance was first observed disconnected (live
        # instances only). None while connected or for file-loaded captures.
        self.disconnected_since: float | None = None

    def touch(self) -> None:
        self.last_used = time.time()

    def is_evictable(self) -> bool:
        """True if this instance is safe to drop without losing in-progress data.

        A live instance still connected is actively recording — never evict it
        automatically. Everything else (file captures, disconnected live
        instances) is fair game for the LRU cap.
        """
        if self.worker is not None and self.path is None:
            try:
                return not self.worker.is_connected()
            except Exception:
                return True
        return True


def _shutdown_worker(worker: object | None) -> None:
    if worker is None:
        return
    try:
        worker.shutdown()
    except Exception:
        pass


def _evict_idle() -> list[str]:
    """Unload instances idle past their TTL: disconnected live instances
    past _DISCONNECTED_TTL_S, file-loaded captures past _FILE_IDLE_TTL_S
    (safe — already durably on disk). Connected live instances are never
    touched here.
    """
    now = time.time()
    evicted = []
    for name, inst in list(instances.items()):
        if inst.path is not None:
            if now - inst.last_used >= _FILE_IDLE_TTL_S:
                del instances[name]
                _shutdown_worker(inst.worker)
                evicted.append(name)
            continue
        if inst.worker is None:
            continue
        try:
            connected = inst.worker.is_connected()
        except Exception:
            connected = False
        if connected:
            inst.disconnected_since = None
            continue
        if inst.disconnected_since is None:
            inst.disconnected_since = now
            continue
        if now - inst.disconnected_since >= _DISCONNECTED_TTL_S:
            del instances[name]
            _shutdown_worker(inst.worker)
            evicted.append(name)
    return evicted


def _evict_for_capacity(exclude: str | None = None) -> str | None:
    """If at/over _MAX_INSTANCES, drop the least-recently-used evictable
    instance to make room for a new one. Returns the evicted name, if any."""
    if len(instances) < _MAX_INSTANCES:
        return None
    candidates = [
        inst for name, inst in instances.items()
        if name != exclude and inst.is_evictable()
    ]
    if not candidates:
        return None
    victim = min(candidates, key=lambda inst: inst.last_used)
    del instances[victim.name]
    _shutdown_worker(victim.worker)
    return victim.name


async def _sweep_loop() -> None:
    """Periodic eviction sweep, doubling as a heartbeat — distinguishes a
    crashed process (see faulthandler) from a hung one (stops producing
    these lines). flush=True since a hang is exactly when buffering would
    hide the last known-good timestamp.
    """
    start = time.time()
    while True:
        await asyncio.sleep(_SWEEP_INTERVAL_S)
        try:
            evicted = _evict_idle()
        except Exception:
            evicted = None
        print(
            f"[heartbeat] uptime={int(time.time() - start)}s "
            f"instances={len(instances)} tasks={len(tasks)} "
            f"evicted={evicted or 'none'}",
            file=sys.stderr,
            flush=True,
        )


@asynccontextmanager
async def _lifespan(_server: MCPServer):
    sweep_task = asyncio.create_task(_sweep_loop())
    try:
        yield
    finally:
        sweep_task.cancel()
        try:
            await sweep_task
        except asyncio.CancelledError:
            pass


mcp_server = MCPServer("Tracy Profiler", lifespan=_lifespan)
executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)

instances: dict[str, TracyInstance] = {}
tasks: dict[str, Task] = {}
captures_dir: str | None = os.environ.get("TRACY_CAPTURES_DIR")


@mcp_server.resource("tracy://prompt")
def _prompt_resource() -> str:
    """Tracy Assist's analysis guidance (system.prompt.md). Contains workflows
    for optimization, callstack inspection, and privacy rules. %TIME%, %USER%,
    and %PROGRAMNAME% are placeholders filled by the in-app chat — ignore them
    when reading from MCP."""
    return _read_text(_PROMPT_PATH)


@mcp_server.resource("tracy://eval-guide")
def _eval_guide_resource() -> str:
    """Bindings-layer guide for the eval tool: ctx object model, time units,
    source-location ID semantics, and worked examples of common ctx queries."""
    return _read_text(_EVAL_GUIDE_PATH)


@mcp_server.tool()
async def list_captures() -> list[str]:
    """List .tracy capture files in the TRACY_CAPTURES_DIR directory (non-recursive)."""
    if not captures_dir:
        return []
    return sorted(glob.glob(os.path.join(captures_dir, "*.tracy")))


@mcp_server.tool()
async def list_instances() -> list[dict]:
    """List all loaded Tracy instances and captures with metadata.

    Each full trace stays resident in memory for as long as this server
    process runs (it's a singleton shared across all MCP clients) until
    unloaded. `idle_seconds` is how long since this instance was last used
    via `eval`/`save_trace` — call `unload_capture` on ones you no longer
    need instead of waiting for automatic eviction. Automatic eviction kicks
    in at the `TRACY_MCP_MAX_INSTANCES` cap (LRU, connected live instances
    are never evicted), after a disconnected live instance sits idle past
    `TRACY_MCP_DISCONNECTED_TTL_S`, or after a file-loaded capture sits idle
    past `TRACY_MCP_FILE_IDLE_TTL_S` — it's already durably on disk, so
    nothing is lost; just `load_capture` it again.

    `background_done: false` means a just-loaded capture is still building
    zone/symbol statistics on a background thread — `get_all_zone_stats`
    and similar stats-based `eval` queries can return empty or partial
    results until this flips to true.
    """
    now = time.time()
    return [
        {
            "id": name,
            "path": inst.path,
            "mtime": inst.mtime,
            "live": inst.path is None,
            "connected": inst.worker.is_connected() if inst.path is None and inst.worker else None,
            "background_done": inst.worker.is_background_done() if inst.worker else None,
            "idle_seconds": now - inst.last_used,
        }
        for name, inst in instances.items()
    ]


@mcp_server.tool()
async def discover_instances(port_range: str = "8086-8095") -> list[dict]:
    """
    Scan for running Tracy-instrumented applications on local ports.

    Returns a list of discovered ports that are listening.
    """
    start_port, end_port = map(int, port_range.split("-"))
    discovered = []

    async def check_port(port: int) -> None:
        try:
            _, writer = await asyncio.wait_for(
                asyncio.open_connection("127.0.0.1", port), timeout=0.1
            )
            writer.close()
            await writer.wait_closed()
            discovered.append({"port": port, "address": "127.0.0.1"})
        except (OSError, asyncio.TimeoutError, ConnectionRefusedError):
            pass

    await asyncio.gather(*(check_port(p) for p in range(start_port, end_port + 1)))
    return discovered


@mcp_server.tool()
async def live_connect(
    address: str = "127.0.0.1",
    port: int = 8086,
    alias: str | None = None,
    memory_limit_mb: int | None = None,
) -> str:
    """
    Connect to a live running Tracy-instrumented application.

    Wraps Worker(addr, port, memoryLimit). A long-lived session on a busy
    target grows unbounded otherwise -- every zone/message/memory event
    stays resident until disconnect, which can OOM-kill this whole process
    rather than just the one instance. memory_limit_mb defaults to
    TRACY_MCP_LIVE_MEMORY_LIMIT_MB (8192 if unset); pass 0 to disable.
    Hitting the limit disconnects that Worker cleanly (Tracy's own
    QueryTerminate path) -- already-collected data stays queryable and
    save_trace-able, it just stops growing.

    Returns the instance_id.
    """
    if not tracy_server:
        return "Error: Tracy Server bindings not found."

    # Pre-flight: read Tracy's UDP broadcast on port 8086 and compare protocol
    # versions before attempting TCP. Mirrors what the Tracy GUI does so a
    # version mismatch produces a precise error instead of an opaque timeout.
    # Tracy clients broadcast every ~3s (TracyProfiler.cpp), so we listen a
    # little longer to guarantee catching at least one beat.
    broadcasts = await _listen_broadcasts(timeout_s=3.5)
    match = next((b for b in broadcasts if b.get("listen_port") == port), None)
    if match and _OUR_PROTOCOL_VERSION is not None:
        if match["protocol_version"] != _OUR_PROTOCOL_VERSION:
            return (
                f"Protocol mismatch: target program '{match['program']}' "
                f"announces Tracy protocol v{match['protocol_version']} on "
                f"{address}:{port}, but these server bindings are built "
                f"against v{_OUR_PROTOCOL_VERSION}. Rebuild the bindings or "
                f"the target against a matching Tracy version."
            )

    limit_mb = _DEFAULT_LIVE_MEMORY_LIMIT_MB if memory_limit_mb is None else memory_limit_mb
    memory_limit = limit_mb * 1024 * 1024 if limit_mb > 0 else -1
    print(f"[live_connect] connecting to {address}:{port} ...", file=sys.stderr, flush=True)
    try:
        w = tracy_server.Worker(address, port, memory_limit)
    except Exception as e:
        return f"Failed to connect: {str(e)}"

    # Worker construction returns immediately even on protocol failure (the
    # bindings expose no error state — is_connected() is the only signal).
    # Probe briefly so silent failures (e.g. TRACY_ON_DEMAND with no profiler
    # request yet, or a target broadcasting on a non-default port) surface
    # cleanly even when broadcast pre-flight didn't catch them.
    deadline_s = 2.0
    step_s = 0.1
    elapsed = 0.0
    watch_task = asyncio.create_task(
        _watch_broadcasts_during(port, deadline_s, "[live_connect]")
    )
    while elapsed < deadline_s and not w.is_connected():
        await asyncio.sleep(step_s)
        elapsed += step_s
    if not watch_task.done():
        watch_task.cancel()
    try:
        await watch_task
    except (asyncio.CancelledError, Exception):
        pass

    if not w.is_connected():
        print(
            f"[live_connect] handshake timed out after {elapsed:.1f}s",
            file=sys.stderr,
            flush=True,
        )
        try:
            w.shutdown()
        except Exception:
            pass
        hint = ""
        if broadcasts and not match:
            seen = ", ".join(
                f"'{b['program']}' on port {b.get('listen_port')} (protocol v{b['protocol_version']})"
                for b in broadcasts
            )
            hint = f" Detected other Tracy broadcasts: {seen}."
        elif not broadcasts:
            hint = (
                " No Tracy broadcasts were received on port 8086 in 3.5s — "
                "the target may use TRACY_ON_DEMAND, a non-default broadcast "
                "port, or isn't running."
            )
        try:
            status_code = w.get_handshake_status()
        except Exception:
            status_code = None
        status_names = {0: "none", 1: "Welcome", 2: "ProtocolMismatch", 3: "NotAvailable", 4: "Dropped"}
        status_hint = ""
        if status_code is not None:
            status_name = status_names.get(status_code, str(status_code))
            print(
                f"[live_connect] final handshake status: {status_name} ({status_code})",
                file=sys.stderr,
                flush=True,
            )
            if status_code == 0:
                status_hint = " Handshake never reached a Welcome/error response at all."
            elif status_code == 1:
                status_hint = (
                    " Handshake reached Welcome; the connection was then lost "
                    "or is still stuck reading the next protocol message -- "
                    "check failure_hint below for the actual reason if there is one."
                )
            else:
                status_hint = f" Handshake ended in state {status_name}."

        # A status==1 drop often means the server itself disconnected the
        # client over a corrupt event stream; get_failure_type() says why.
        failure_hint = ""
        try:
            failure_names = {
                0: "None", 1: "ZoneStack", 2: "ZoneDoubleEnd", 3: "ZoneText",
                4: "ZoneValue", 5: "ZoneColor", 6: "ZoneName", 7: "MemFree",
                8: "MemAllocTwice", 9: "FrameEnd", 10: "FrameImageIndex",
                11: "FrameImageTwice", 12: "FiberLeave",
                13: "SourceLocationOverflow", 14: "LockThreadOverflow",
            }
            failure_code = w.get_failure_type()
            if failure_code:
                failure_name = failure_names.get(failure_code, str(failure_code))
                fd = w.get_failure_data()
                loc = f"{fd.get('file', '?')}:{fd.get('line', '?')} ({fd.get('zone_name', fd.get('function', '?'))})" if fd else "?"
                failure_hint = (
                    f" The server disconnected the client itself: Failure.{failure_name} "
                    f"on thread '{fd.get('thread_name', '?')}' at {loc} -- this is a "
                    f"CORRUPT event stream from the target's own instrumentation "
                    f"(e.g. an unbalanced Zone macro), not a network stall. Fix the "
                    f"target's instrumentation at that location."
                )
                print(
                    f"[live_connect] failure: Failure.{failure_name} thread={fd.get('thread_name')} "
                    f"loc={loc}",
                    file=sys.stderr,
                    flush=True,
                )
        except Exception:
            pass

        return (
            f"Reached {address}:{port} but the Tracy handshake did not complete "
            f"within {deadline_s:.1f}s.{status_hint}{failure_hint}{hint} Common causes: (1) the "
            f"Tracy client version embedded in the target program differs from "
            f"these server bindings; (2) the target was built with "
            f"TRACY_ON_DEMAND and is awaiting a profiler request; (3) another "
            f"client is already attached."
        )

    print(
        f"[live_connect] handshake completed after {elapsed:.1f}s",
        file=sys.stderr,
        flush=True,
    )
    name = alias or f"live_{address}_{port}"
    if name in instances:
        _shutdown_worker(instances[name].worker)
    evicted = _evict_for_capacity(exclude=name)
    instances[name] = TracyInstance(name, w)
    note = f" (evicted idle instance '{evicted}' to stay under the {_MAX_INSTANCES}-instance cap)" if evicted else ""
    limit_note = f"{limit_mb}MB memory limit" if limit_mb > 0 else "no memory limit"
    return (
        f"Connected to live instance as '{name}'{note} ({limit_note}). "
        f"Before your first eval, read resources tracy://prompt "
        f"(analysis guidance) and tracy://eval-guide (ctx object model, "
        f"ns time units, srcloc IDs)."
    )


@mcp_server.tool()
async def load_capture(path: str, alias: str | None = None) -> str:
    """
    Load a .tracy capture file by absolute path.

    Parameters:
      path  — absolute path to a .tracy file. On Windows use backslashes
              (e.g. 'E:\\\\traces\\\\foo.tracy').
      alias — optional instance name; overwrites existing on collision.
              If omitted, an ID is derived from filename and mtime.

    If you don't already have a path, call `list_captures` first — it lists
    .tracy files in the TRACY_CAPTURES_DIR environment directory.

    Loading returns before zone/symbol statistics finish building — a
    background thread populates them after the file itself is read. Check
    `background_done` in `list_instances` (or poll it) before relying on
    `get_all_zone_stats` and similar stats-based `eval` queries; they can
    return empty or partial results until it's true.
    """
    if not tracy_server:
        return "Error: Tracy Server bindings not found."
    try:
        mtime = os.path.getmtime(path)
        if alias:
            name = alias
        else:
            # unique name including mtime to avoid version collision
            name = f"{os.path.basename(path)}@{int(mtime):x}"

        if name in instances:
            inst = instances[name]
            if inst.path == path and inst.mtime == mtime:
                inst.touch()
                return f"Instance '{name}' is already loaded and up to date."
            _shutdown_worker(inst.worker)

        evicted = _evict_for_capacity(exclude=name)

        f = tracy_server.open_file(path)
        w = tracy_server.create_worker_from_file(f)
        inst = TracyInstance(name, w)
        inst.path = path
        inst.mtime = mtime
        instances[name] = inst
        note = f" (evicted idle instance '{evicted}' to stay under the {_MAX_INSTANCES}-instance cap)" if evicted else ""
        return (
            f"Loaded as '{name}'{note}. "
            f"Before your first eval, read resources tracy://prompt "
            f"(analysis guidance) and tracy://eval-guide (ctx object model, "
            f"ns time units, srcloc IDs)."
        )
    except Exception as e:
        return f"Failed to load: {str(e)}"


@mcp_server.tool()
async def save_trace(
    instance_id: str,
    path: str,
    level: int = 3,
    streams: int = 4,
    fi_dict: bool = False,
    overwrite: bool = False,
    async_mode: bool = True,
) -> object:
    """
    Snapshot a Tracy instance (live or loaded) to a .tracy file.

    Wraps `Worker::Write` under the main-thread data lock — safe for live
    instances; the receive thread yields cooperatively for the duration of
    the write. Concurrent `eval` calls against the same live instance may
    stall until the save completes.

    Parameters:
      instance_id — name returned by live_connect / load_capture.
      path        — output file. Absolute paths are used as-is; a bare
                    filename is resolved under TRACY_CAPTURES_DIR if set.
                    On Windows use backslashes (e.g. 'E:\\\\traces\\\\a.tracy').
      level       — Zstd compression level (default 3, matches capture tool).
      streams     — number of compression streams (default 4).
      fi_dict     — rebuild frame-image dedup dictionary on save. Only
                    meaningful for traces containing screenshots; default
                    False matches the capture tool and GUI default.
      overwrite   — refuse to clobber an existing file unless True.
      async_mode  — default True; large traces take seconds-to-minutes.
                    Returns {task_id, status: "running"}; poll with `task`.

    On success returns a dict with path, uncompressed_bytes, compressed_bytes,
    ratio, and elapsed_ms — the same numbers the capture tool prints.
    """
    if instance_id not in instances:
        return f"Error: Instance '{instance_id}' not found. Use list_instances to find valid IDs."
    instance = instances[instance_id]
    if not instance.worker:
        return f"Error: Instance '{instance_id}' has no worker."
    instance.touch()

    if not os.path.isabs(path):
        if captures_dir and os.path.basename(path) == path:
            path = os.path.join(captures_dir, path)
        else:
            return (
                f"Error: '{path}' is not absolute. Pass a full path, or a bare "
                f"filename with TRACY_CAPTURES_DIR set (currently "
                f"{captures_dir!r})."
            )
    if not path.endswith(".tracy"):
        path += ".tracy"
    if os.path.exists(path) and not overwrite:
        return (
            f"Error: '{path}' already exists. Pass overwrite=True to clobber, "
            f"or choose a different path."
        )

    if not async_mode:
        return await _execute_save(instance.worker, path, level, streams, fi_dict)

    task_id = str(uuid.uuid4())
    t = Task(task_id, f"save_trace({instance_id} -> {path})")
    tasks[task_id] = t
    asyncio.get_running_loop().run_in_executor(
        executor, _run_save_task_sync, t, instance.worker, path, level, streams, fi_dict
    )
    return {"task_id": task_id, "status": "running"}


def _save_worker_sync(worker: object, path: str, level: int, streams: int, fi_dict: bool) -> dict:
    t0 = time.time()
    uncompressed, compressed = tracy_server.save_worker(
        worker, path, level, streams, fi_dict
    )
    elapsed_ms = int((time.time() - t0) * 1000)
    ratio = (compressed / uncompressed) if uncompressed else 0.0
    return {
        "path": path,
        "uncompressed_bytes": uncompressed,
        "compressed_bytes": compressed,
        "ratio": ratio,
        "elapsed_ms": elapsed_ms,
    }


def _run_save_task_sync(t: Task, worker: object, path: str, level: int, streams: int, fi_dict: bool) -> None:
    t.status = "running"
    try:
        t.result = _save_worker_sync(worker, path, level, streams, fi_dict)
        t.status = "completed"
    except Exception as e:
        t.error = str(e)
        t.status = "failed"
    finally:
        t.end_time = time.time()


async def _execute_save(worker: object, path: str, level: int, streams: int, fi_dict: bool) -> dict:
    return await asyncio.get_running_loop().run_in_executor(
        executor, _save_worker_sync, worker, path, level, streams, fi_dict
    )


@mcp_server.tool()
async def unload_capture(instance_id: str) -> str:
    """Unload a Tracy instance and release its memory.

    Prefer calling this as soon as you're done analyzing a capture — every
    loaded instance holds the *entire* trace (zones, messages, callstacks,
    memory events) in memory for as long as this server process runs, and
    the server is a long-lived singleton shared across all MCP clients.
    Automatic eviction exists as a backstop (see `list_instances`), not a
    substitute for unloading what you no longer need.
    """
    if instance_id in instances:
        inst = instances.pop(instance_id)
        _shutdown_worker(inst.worker)
        return f"Instance '{instance_id}' unloaded."
    return f"Instance '{instance_id}' not found."


@mcp_server.tool(name="eval")
async def tracy_eval(code: str, instance_id: str, async_mode: bool = False) -> object:
    """
    Execute Python code against a specific Tracy Worker bound as `ctx`.

    On first use, read the `tracy://prompt` (analysis guidance) and
    `tracy://eval-guide` (ctx object model, units, source-location ID joins)
    resources. Time values returned by Worker methods are nanoseconds.

    If async_mode=True, returns a task_id immediately; poll via the `task` tool.
    """
    if instance_id not in instances:
        return f"Error: Instance '{instance_id}' not found. Use list_instances to find valid IDs."

    instance = instances[instance_id]
    if not instance.worker:
        return f"Error: Instance '{instance_id}' has no worker."
    instance.touch()

    if not async_mode:
        return await _execute_eval(code, instance.worker)

    # Async mode: spawn task and return immediately
    task_id = str(uuid.uuid4())
    task = Task(task_id, code)
    tasks[task_id] = task
    asyncio.get_running_loop().run_in_executor(
        executor, _run_task_sync, task, instance.worker
    )
    return {"task_id": task_id, "status": "running"}


def _run_task_sync(task: Task, worker: object) -> None:
    """Run a background eval task in the thread pool."""
    task.status = "running"
    try:
        task.result = _execute_eval_sync(task.code, worker)
        task.status = "completed"
    except Exception as e:
        task.error = str(e)
        task.status = "failed"
    finally:
        task.end_time = time.time()


def _execute_eval_sync(code: str, ctx: object) -> str:
    """Execute *code* with `ctx` bound to the Tracy worker. Captures stdout."""
    global_vars = {
        "__builtins__": builtins,
        "ctx": ctx,
        "tracy": tracy_server,
        "instances": {name: inst.worker for name, inst in instances.items()},
    }
    buf = io.StringIO()
    with redirect_stdout(buf):
        try:
            result = eval(compile(code, "<eval>", "eval"), global_vars)
        except SyntaxError:
            exec(compile(code, "<exec>", "exec"), global_vars)
            result = None
    output = buf.getvalue()
    if result is None:
        return output or ""
    return str(result)


async def _execute_eval(code: str, ctx: object) -> str:
    """Async wrapper: runs `_execute_eval_sync` in the thread-pool executor."""
    return await asyncio.get_running_loop().run_in_executor(
        executor, _execute_eval_sync, code, ctx
    )


@mcp_server.tool()
async def task(action: str, task_id: str | None = None) -> object:
    """
    Manage background analysis tasks.

    Actions: poll, cancel, list
    """
    if action == "list":
        return [
            {"id": t.id, "status": t.status, "elapsed": time.time() - t.start_time}
            for t in tasks.values()
        ]

    if not task_id or task_id not in tasks:
        return "Error: Task ID not found."

    t = tasks[task_id]
    if action == "poll":
        res: dict = {"id": t.id, "status": t.status}
        if t.status == "completed":
            res["result"] = t.result
        elif t.status == "failed":
            res["error"] = t.error
        return res

    if action == "cancel":
        # Cancellation of thread-pool work is not possible post-submission;
        # mark the task so callers know it was abandoned.
        if t.status == "running":
            t.status = "cancelled"
            return f"Task {task_id} marked as cancelled."
        return f"Task {task_id} is not running."

    return "Error: Unknown action."


@mcp_server.tool()
async def shutdown_server() -> str:
    """
    Shut down the Tracy MCP server.

    Because the server runs as a singleton (SSE transport, one process shared
    across all VS Code windows), this releases the TracyServerBindings.pyd lock
    for all clients at once. Restart tracy_mcp.py after rebuilding.
    """
    import threading
    def _exit() -> None:
        time.sleep(0.2)
        os._exit(0)
    threading.Thread(target=_exit, daemon=True).start()
    return "Server shutting down. Restart tracy_mcp.py to reconnect."


if __name__ == "__main__":
    # Enabled before anything touches the native bindings, so a genuine
    # fatal crash (segfault etc.) writes a thread-state traceback to disk
    # before the process dies — works on Windows via
    # SetUnhandledExceptionFilter, same mechanism as POSIX signals.
    _crash_log_fh = open(_CRASH_LOG_FILE, "a", encoding="utf-8")
    faulthandler.enable(file=_crash_log_fh, all_threads=True)

    atexit.register(_cleanup_pid_files)

    # streamable-http/SSE keep-alive connections never close on their own, so
    # uvicorn's graceful shutdown hangs waiting for them unless a second
    # interrupt arrives. Arm a watchdog on the first signal instead.
    import threading
    import uvicorn as _uvicorn

    _orig_handle_exit = _uvicorn.Server.handle_exit

    def _handle_exit_with_watchdog(self, sig, frame):
        _orig_handle_exit(self, sig, frame)

        def _force_exit_after_grace() -> None:
            time.sleep(3)
            os._exit(0)

        threading.Thread(target=_force_exit_after_grace, daemon=True).start()

    _uvicorn.Server.handle_exit = _handle_exit_with_watchdog

    if _TRANSPORT not in ("sse", "streamable-http"):
        print(
            "TRACY_MCP_TRANSPORT must be 'sse' or 'streamable-http'.",
            file=sys.stderr,
        )
        sys.exit(1)

    running, existing_port = _is_our_server_running()
    if running:
        print(
            f"Tracy MCP already running on port {existing_port}. "
            "All VS Code windows share that instance.",
            file=sys.stderr,
        )
        sys.exit(0)

    port = _find_free_port()
    _write_pid_and_port(port)

    path = _SSE_PATH if _TRANSPORT == "sse" else _STREAMABLE_HTTP_PATH
    proto = "unknown" if _OUR_PROTOCOL_VERSION is None else _OUR_PROTOCOL_VERSION
    print(
        f"Tracy MCP listening on http://127.0.0.1:{port}{path} (protocol v{proto})",
        file=sys.stderr,
    )

    # v2's MCPServer takes transport options on run(), not the constructor
    # or a mutable settings object (SDK v1 -> v2 migration).
    run_kwargs = {"host": "127.0.0.1", "port": port}
    if _TRANSPORT == "sse":
        run_kwargs["sse_path"] = _SSE_PATH
    else:
        run_kwargs["streamable_http_path"] = _STREAMABLE_HTTP_PATH
    try:
        mcp_server.run(transport=_TRANSPORT, **run_kwargs)
    except KeyboardInterrupt:
        print("\nTracy MCP server stopped.", file=sys.stderr)
        sys.exit(0)