#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["tiktoken"]
# ///
"""Measure what the GitHub MCP server's tool definitions cost in context.

    uv run measure-github-mcp.py dump [--toolsets=all] > tools.json
    uv run measure-github-mcp.py tiktoken tools.json ...

dump talks to the server's Docker image over stdio and needs
GITHUB_PERSONAL_ACCESS_TOKEN set (`gh auth token` works). MCP_IMAGE overrides
the image; the default is the v1.14.0 tag the post measured.

tiktoken counts the compact JSON of name, description and schema, in the
shape the Claude API takes tool definitions. Its o200k_base encoding stands
in for Claude's tokenizer, so the counts are approximate.
"""
import itertools
import json
import os
import subprocess
import sys

import tiktoken

IMAGE = os.environ.get("MCP_IMAGE", "ghcr.io/github/github-mcp-server:v1.14.0")


def send(proc, message):
    proc.stdin.write(json.dumps(message) + "\n")
    proc.stdin.flush()


def call(proc, rid, method, params):
    """Send one JSON-RPC request and return the result of the matching response."""
    send(proc, {"jsonrpc": "2.0", "id": rid, "method": method, "params": params})
    for line in proc.stdout:
        if line.startswith("{"):
            message = json.loads(line)
            if message.get("id") == rid:
                return message["result"]
    sys.exit("server closed the pipe:\n" + proc.stderr.read()[-2000:])


def dump(server_args):
    """Print the server's complete tool list as JSON."""
    cmd = ["docker", "run", "-i", "--rm", "-e", "GITHUB_PERSONAL_ACCESS_TOKEN", IMAGE, "stdio", *server_args]
    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    info = call(proc, 1, "initialize", {
        "protocolVersion": "2025-06-18",
        "capabilities": {},
        "clientInfo": {"name": "measure", "version": "0"},
    })
    send(proc, {"jsonrpc": "2.0", "method": "notifications/initialized"})

    tools = []
    cursor = None
    for rid in itertools.count(2):
        page = call(proc, rid, "tools/list", {"cursor": cursor} if cursor else {})
        tools += page["tools"]
        cursor = page.get("nextCursor")
        if not cursor:
            break
    proc.terminate()
    json.dump({"serverInfo": info["serverInfo"], "args": server_args, "tools": tools}, sys.stdout)


def count(paths):
    """Print the token count of each dump."""
    enc = tiktoken.get_encoding("o200k_base")
    for path in paths:
        with open(path) as f:
            d = json.load(f)
        tools = [
            {"name": t["name"], "description": t.get("description", ""), "input_schema": t.get("inputSchema", {})}
            for t in d["tools"]
        ]
        tokens = len(enc.encode(json.dumps(tools, separators=(",", ":"))))
        print(f"{path}: {d['serverInfo']['version']}, {len(tools)} tools, {tokens} tokens")


def main(argv):
    if argv[:1] == ["dump"]:
        dump(argv[1:])
    elif argv[:1] == ["tiktoken"] and argv[1:]:
        count(argv[1:])
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main(sys.argv[1:])