Topic 5: Module 2 lab

"""Module 2 Lab: one script that exercises the notes server across both transports and both eras.

It starts its own HTTP server on port 8024, runs every check, prints a report, and stops the server.
"""
from __future__ import annotations

import contextlib
import io
import json
import os
import statistics
import subprocess
import sys
import time

import anyio
import httpx2

from mcp import Client, StdioServerParameters
from mcp.types import ResourceTemplateReference

from m02_wire import StdioWire, meta

PORT = 8024
URL = f"http://127.0.0.1:{PORT}/mcp"
STDIO = StdioServerParameters(command=sys.executable, args=["examples/m02_server.py"], env={"PYTHONPATH": "."})
report: list[tuple[str, bool, str]] = []


def check(name: str, ok: bool, detail: str) -> None:
    report.append((name, ok, detail))


def start_http_server() -> subprocess.Popen[bytes]:
    env = {**os.environ, "PYTHONPATH": ".", "NOTES_TRANSPORT": "streamable-http", "NOTES_PORT": str(PORT)}
    proc = subprocess.Popen([sys.executable, "examples/m02_server.py"], env=env, stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
    for _ in range(100):  # wait until the port answers
        try:
            httpx2.post(URL, content=b"{}", headers={"Content-Type": "application/json"})
            return proc
        except httpx2.TransportError:
            time.sleep(0.1)
    raise RuntimeError("HTTP server did not start")


def raw_wire_checks() -> None:
    """Hand-written JSON-RPC over stdio: discovery, the envelope rule, and stdio line sizes."""
    with contextlib.redirect_stdout(io.StringIO()):  # keep the lab report short
        wire = StdioWire(sys.executable, "examples/m02_server.py")
        discovered = wire.request("server/discover")
        wire.send({"jsonrpc": "2.0", "id": 50, "method": "tools/list",
                   "params": {"_meta": {"io.modelcontextprotocol/protocolVersion": "2026-07-28"}}})
        missing = wire.receive()
        request = {"jsonrpc": "2.0", "id": 51, "method": "tools/call",
                   "params": {"name": "search_notes", "arguments": {"query": "sleep memory", "limit": 3}, "_meta": meta()}}
        wire.send(request)
        response_line = wire.lines.get(timeout=5)  # the raw line, to count its bytes
        wire.close()
    check("server/discover over stdio", discovered["result"]["supportedVersions"] == ["2026-07-28"],
          f"supportedVersions={discovered['result']['supportedVersions']}")
    check("missing clientCapabilities rejected", missing["error"]["code"] == -32602, f"code {missing['error']['code']}")
    sent = len(json.dumps(request, separators=(",", ":")).encode()) + 1  # +1 for the newline
    received = len(response_line.encode()) + 1
    check("stdio bytes for one search_notes call", True, f"{sent} sent + {received} received = {sent + received}")


async def sdk_checks() -> None:
    async with Client(URL) as client:
        check("modern HTTP negotiation", client.protocol_version == "2026-07-28", client.protocol_version)
        created = await client.call_tool("start_reading_list", {"name": "Lab"})
        list_id = created.structured_content["list_id"]
        await client.call_tool("add_to_reading_list", {"list_id": list_id, "note_id": "sleep-and-memory"})
    async with Client(URL, mode="legacy") as legacy:  # a different connection, a different era
        shown = await legacy.call_tool("show_reading_list", {"list_id": list_id})
        links = [str(b.uri) for b in shown.content if b.type == "resource_link"]
        check("handle survives a new legacy connection", links == ["notes://sleep-and-memory"],
              f"{legacy.protocol_version} sees {links}")
        progress: list[float] = []

        async def on_progress(value: float, total: float | None, message: str | None) -> None:
            progress.append(value)

        await legacy.call_tool("reindex_notes", {"delay_seconds": 0.01}, progress_callback=on_progress)
        await anyio.sleep(0.2)  # progress notifications are delivered beside the response
        check("progress over legacy HTTP", len(progress) == 8, f"{len(progress)} notifications")
    async with Client(URL) as client:
        done = await client.complete(ref=ResourceTemplateReference(type="ref/resource", uri="notes://{note_id}"),
                                     argument={"name": "note_id", "value": "sp"})
        check("completion for notes://{note_id}", done.completion.values == ["spaced-repetition"], str(done.completion.values))


def security_checks() -> None:
    body = b'{"jsonrpc":"2.0","id":1,"method":"tools/list"}'
    evil = httpx2.post(URL, content=body, headers={"Content-Type": "application/json", "Origin": "http://evil.example"})
    check("foreign Origin rejected", evil.status_code == 403, f"HTTP {evil.status_code}")
    rebound = httpx2.post(URL, content=body, headers={"Content-Type": "application/json", "Host": f"attacker.example:{PORT}"})
    check("foreign Host rejected", rebound.status_code == 421, f"HTTP {rebound.status_code}")


async def latency(label: str, target: object, calls: int = 30) -> None:
    async with Client(target) as client:
        for _ in range(3):
            await client.call_tool("search_notes", {"query": "sleep memory", "limit": 3})
        samples = []
        for _ in range(calls):
            t0 = time.perf_counter()
            await client.call_tool("search_notes", {"query": "sleep memory", "limit": 3})
            samples.append((time.perf_counter() - t0) * 1000)
    check(f"latency {label}", True, f"median {statistics.median(samples):.1f} ms over {calls} calls")


async def main() -> None:
    server = start_http_server()
    try:
        raw_wire_checks()
        await sdk_checks()
        security_checks()
        await latency("stdio", STDIO)
        await latency("http", URL)
    finally:
        server.terminate()
        server.wait(timeout=10)
    width = max(len(name) for name, _, _ in report)
    for name, ok, detail in report:
        print(f"{'PASS' if ok else 'FAIL'}  {name:{width}}  {detail}")
    print(f"{sum(ok for _, ok, _ in report)}/{len(report)} checks passed")


if __name__ == "__main__":
    anyio.run(main)