MCP primitives lab.

"""
server.py — the MCP primitives lab: every server-side feature of the protocol in one small server.

    Tools        echo, negotiation, count, summarize, roots, add_note, delete_note
    Resources    lab://notes (list), lab://notes/{note_id} (template), lab://clock (changes every call)
    Prompts      explain
    Sampling     summarize asks the CLIENT's model to write the summary (sampling/createMessage)
    Roots        roots asks the client which directories it is allowed to touch (roots/list)
    Elicitation  delete_note asks the person to confirm before it deletes (elicitation/create)
    Progress     count streams notifications/progress and log messages while it works
    Transports   --transport stdio | streamable-http  (plus --stateless and --json-response for HTTP)

The `negotiation` tool returns the flags the server sees for the current connection: the protocol
version that was negotiated, the client's declared capabilities, and, over HTTP, the session id.

Run:   python server.py                                   stdio (a host launches this as a subprocess)
       python server.py --transport streamable-http       http://127.0.0.1:8765/mcp
Needs: pip install "mcp>=2.2"
"""
from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import logging
import os
from typing import Annotated, Any

from pydantic import BaseModel, Field

from mcp.server.mcpserver import Context, Elicit, ListRoots, MCPServer, Resolve, Sample
from mcp.server.mcpserver.exceptions import ToolError
from mcp.types import (
    CreateMessageResult,
    ListRootsResult,
    ModelHint,
    ModelPreferences,
    SamplingMessage,
    TextContent,
    ToolAnnotations,
)

log = logging.getLogger("primitives-lab")  # stderr only: on stdio, stdout is the protocol channel

mcp = MCPServer(
    "primitives-lab",
    version="1.0.0",
    instructions=(
        "A teaching server. Call `negotiation` first to see the flags for this connection, "
        "then try every tool once. `summarize` uses your own model through sampling, "
        "`roots` asks which directories you allow, `delete_note` asks you to confirm."
    ),
)

NOTES: dict[str, str] = {
    "welcome": "Resources are application-controlled context: the host decides what the model sees.",
    "flags": "Track: protocolVersion, capabilities on both sides, initialized, and the HTTP session id.",
}

READ_ONLY = ToolAnnotations(read_only_hint=True, idempotent_hint=True, open_world_hint=False)


# ------------------------------------------------------------------ tools
@mcp.tool(title="Echo", annotations=READ_ONLY)
def echo(text: str) -> str:
    """Return the text unchanged. The simplest possible tool: one argument, one result."""
    return text


class Negotiation(BaseModel):
    protocol_version: str = Field(description="The version both sides agreed on during initialize")
    client_capabilities: dict[str, Any] = Field(description="What the client declared it can do (sampling, roots, elicitation...)")
    transport: str = Field(description="How this server was started")
    http_session_id: str | None = Field(description="Mcp-Session-Id header, present on stateful Streamable HTTP only")
    request_id: str = Field(description="JSON-RPC id of this very tools/call request")


@mcp.tool(title="Show the negotiation", annotations=READ_ONLY)
def negotiation(ctx: Context) -> Negotiation:
    """The flags the server sees for this connection: negotiated protocol version, the client's
    declared capabilities, the transport, the HTTP session id if any, and this request's id."""
    caps = ctx.client_capabilities
    try:
        session_id = ctx.headers.get("mcp-session-id")
    except Exception:  # no HTTP request behind this call (stdio)
        session_id = None
    return Negotiation(
        protocol_version=str(ctx.protocol_version),
        client_capabilities=caps.model_dump(exclude_none=True) if caps is not None else {},
        transport=os.environ.get("LAB_TRANSPORT", "stdio"),
        http_session_id=session_id,
        request_id=str(ctx.request_id),
    )


@mcp.tool(title="Count slowly", annotations=READ_ONLY)
async def count(n: int, ctx: Context) -> str:
    """Count to n, sending a notifications/progress message for every step.
    (The separate logging channel, notifications/message, is deprecated as of 2026-07-28: put a
    human-readable `message` on the progress notification instead, and log to stderr for operators.)"""
    if n < 1 or n > 50:
        raise ToolError("n must be between 1 and 50")
    for i in range(1, n + 1):
        await ctx.report_progress(progress=i, total=n, message=f"step {i} of {n}")
        log.info("counted %d", i)  # stderr: never stdout on stdio
        await asyncio.sleep(0.05)
    return f"counted to {n}"


# --- Server-to-client asks: sampling, roots and elicitation ------------------------------------
# Up to spec 2025-11-25 the server sent these as JSON-RPC requests BACK over the connection, which
# needs a back-channel (stdio, or the SSE stream). Since 2026-07-28 they are "input required" rounds:
# the tool result says what it needs, the client supplies it, and the call is retried with the answer
# riding along in request_state. That is what lets them work over stateless HTTP. In this SDK you
# declare the need with a RESOLVER: a function that returns Sample(...), ListRoots() or Elicit(...);
# the framework runs the round trip and injects the answer as a tool argument the model never sees.


def ask_client_model(text: str) -> Sample:
    """Resolver for `summarize`: what to ask the client's model. Must render identically on every retry round."""
    return Sample(
        [SamplingMessage(role="user", content=TextContent(type="text", text=f"Summarise in one sentence:\n\n{text}"))],
        max_tokens=120,
        system_prompt="You summarise plainly. One sentence, no preamble.",
        model_preferences=ModelPreferences(hints=[ModelHint(name="claude")], intelligencePriority=0.3, speedPriority=0.8),
    )


@mcp.tool(title="Summarise with the client's model", annotations=READ_ONLY)
def summarize(text: str, answer: Annotated[CreateMessageResult, Resolve(ask_client_model)]) -> str:
    """Ask the CLIENT's language model to summarise the text in one sentence (sampling/createMessage).
    The server has no model and no API key: the client owns the model, and can review the request."""
    content = answer.content
    reply = content.text if isinstance(content, TextContent) else str(content)
    return f"[{answer.model}] {reply}"


def ask_roots() -> ListRoots:
    """Resolver for `roots`: request the client's roots/list."""
    return ListRoots()


@mcp.tool(title="Which roots may I touch?", annotations=READ_ONLY)
def roots(allowed: Annotated[ListRootsResult, Resolve(ask_roots)]) -> list[str]:
    """Ask the client for its roots (roots/list): the directories or URIs this server is allowed to work in."""
    return [f"{r.name or '(unnamed)'} -> {r.uri}" for r in allowed.roots]


@mcp.tool(title="Add a note", annotations=ToolAnnotations(read_only_hint=False, destructive_hint=False, idempotent_hint=True, open_world_hint=False))
async def add_note(note_id: str, text: str, ctx: Context) -> str:
    """Create or replace a note, then tell subscribed clients the resource list changed."""
    if not note_id.isidentifier():
        raise ToolError("note_id must be a simple identifier, e.g. meeting_notes")
    NOTES[note_id] = text
    await ctx.notify_resources_changed()
    return f"lab://notes/{note_id}"


class ConfirmDelete(BaseModel):
    sure: bool = Field(description="Really delete this note?")


def confirm_delete(note_id: str) -> ConfirmDelete | Elicit[ConfirmDelete]:
    """Resolver for `delete_note`: ask the PERSON (elicitation/create) unless there is nothing to delete."""
    if note_id not in NOTES:
        return ConfirmDelete(sure=False)  # nothing to ask; the tool will report the missing note
    return Elicit(f"Delete the note '{note_id}'? This cannot be undone.", ConfirmDelete)


@mcp.tool(title="Delete a note", annotations=ToolAnnotations(read_only_hint=False, destructive_hint=True, idempotent_hint=True, open_world_hint=False))
async def delete_note(note_id: str, confirm: Annotated[ConfirmDelete, Resolve(confirm_delete)], ctx: Context) -> str:
    """Delete a note, but only after the person confirms (elicitation/create). The model cannot answer
    this question: the client shows it to the human and sends back the human's answer."""
    if note_id not in NOTES:
        raise ToolError(f"no note called {note_id}")
    if not confirm.sure:
        return f"kept {note_id}: the person said no"
    del NOTES[note_id]
    await ctx.notify_resources_changed()
    return f"deleted {note_id}"


# ------------------------------------------------------------------ resources
@mcp.resource("lab://notes", title="All notes", mime_type="application/json")
def list_notes() -> str:
    """Every note id with its length. A plain, static-looking resource."""
    import json
    return json.dumps({k: len(v) for k, v in NOTES.items()})


@mcp.resource("lab://notes/{note_id}", title="One note", mime_type="text/plain")
def read_note(note_id: str) -> str:
    """A resource template: the client fills in {note_id}. Unknown ids are an error, not empty text."""
    if note_id not in NOTES:
        raise ToolError(f"no note called {note_id}")
    return NOTES[note_id]


@mcp.resource("lab://clock", title="Server clock", mime_type="text/plain")
def clock() -> str:
    """Changes on every read: shows that a resource can be dynamic."""
    return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds")


# ------------------------------------------------------------------ prompts
@mcp.prompt(title="Explain a concept")
def explain(concept: str, level: str = "beginner") -> str:
    """A reusable prompt template the USER picks (user-controlled), unlike tools the model picks."""
    return (
        f"Explain '{concept}' to a {level}. Start with a one-sentence definition, give one concrete "
        "example, then name the most common mistake people make with it."
    )


# ------------------------------------------------------------------ entry point
def main() -> None:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--transport", choices=["stdio", "streamable-http"], default=os.environ.get("MCP_TRANSPORT", "stdio"))
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=8765)
    p.add_argument("--stateless", action="store_true", help="HTTP: no session id, every request stands alone")
    p.add_argument("--json-response", action="store_true", help="HTTP: answer with application/json instead of an SSE stream")
    a = p.parse_args()
    logging.basicConfig(level=logging.INFO, stream=__import__("sys").stderr, format="%(name)s %(levelname)s %(message)s")
    os.environ["LAB_TRANSPORT"] = a.transport + (" stateless" if a.stateless else "") + (" json" if a.json_response else "")
    if a.transport == "streamable-http":
        log.info("Streamable HTTP on http://%s:%d/mcp  stateless=%s json=%s", a.host, a.port, a.stateless, a.json_response)
        mcp.run(transport="streamable-http", host=a.host, port=a.port, stateless_http=a.stateless, json_response=a.json_response)
    else:
        mcp.run(transport="stdio")


if __name__ == "__main__":
    main()

"""
wire_trace.py — the handshakes with NO SDK, so you can see exactly what is on the wire.

    python wire_trace.py stdio      spawn server.py and speak newline-delimited JSON-RPC over its stdin/stdout
    python wire_trace.py http       talk to a running `python server.py --transport streamable-http`
                                    with plain HTTP requests, and read the SSE stream by hand

Every line marked  ->  is bytes we sent, every line marked  <-  is bytes we received.

Two generations of handshake are shown, because you will meet both:

  2026-07-28  `server/discover`, then every request carries the version and the client's capabilities
              in `_meta`. No session state is required, so it works over stateless HTTP. A server that
              needs the client's model, roots or the person answers `resultType: "input_required"` and
              the client RETRIES the call with `inputResponses` and the opaque `requestState`.
  2025-11-25  `initialize` -> reply -> `notifications/initialized`. The version and capabilities are
              fixed once per connection; over HTTP the server issues an `Mcp-Session-Id`.

The flags to keep track of: the request id that must match; the protocol version proposed and
confirmed; the capabilities each side declares; `initialized` (legacy) or `_meta` on every request
(modern); over HTTP the `Mcp-Session-Id` and `MCP-Protocol-Version` headers; and, for an input-required
round, the `requestState` that must go back unchanged.
"""
from __future__ import annotations

import json
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
MODERN = "2026-07-28"
LEGACY = "2025-11-25"
CLIENT_CAPS = {"roots": {"listChanged": True}, "sampling": {}, "elicitation": {"form": {}}}
META = {  # modern: rides on EVERY request instead of being fixed by initialize
    "io.modelcontextprotocol/protocolVersion": MODERN,
    "io.modelcontextprotocol/clientInfo": {"name": "wire-trace", "version": "1.0.0"},
    "io.modelcontextprotocol/clientCapabilities": CLIENT_CAPS,
}


def show(direction: str, payload: dict | str, note: str = "", width: int = 420) -> None:
    text = payload if isinstance(payload, str) else json.dumps(payload)
    if len(text) > width:
        text = text[:width] + f" …({len(text)} chars)"
    print(f"{direction} {text}")
    if note:
        print(f"   {note}")


def stub_model_answer(request: dict) -> dict:
    """What a host would do with a sampling/createMessage request: call its own model. We fake it."""
    asked = request["params"]["messages"][-1]["content"]["text"]
    return {"role": "assistant", "model": "stub-model-1", "stopReason": "endTurn",
            "content": {"type": "text", "text": "In short: " + " ".join(asked.split("\n\n", 1)[-1].split()[:8]) + "…"}}


# ------------------------------------------------------------------ stdio
def trace_stdio() -> None:
    print("== stdio: one JSON-RPC message per line on the child's stdin/stdout; stderr is free for logs ==")

    def spawn():
        return subprocess.Popen([sys.executable, str(HERE / "server.py"), "--transport", "stdio"],
                                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                                text=True, encoding="utf-8")

    def rpc(proc, msg: dict, note_out: str = "", note_in: str = "") -> dict | None:
        show("->", msg, note_out)
        proc.stdin.write(json.dumps(msg) + "\n")        # newline-delimited; a message may not contain a raw newline
        proc.stdin.flush()
        if "id" not in msg:
            return None                                 # a notification gets no reply
        reply = json.loads(proc.stdout.readline())
        show("<-", reply, note_in)
        assert reply["id"] == msg["id"], "flag: the reply must carry the request's id"
        return reply

    print("\n-- A. modern handshake (2026-07-28): server/discover, then _meta on every request --")
    p = spawn()
    disc = rpc(p, {"jsonrpc": "2.0", "id": 1, "method": "server/discover", "params": {"_meta": META}},
               "flag: we say which version we speak and what we can do, inside _meta",
               "flag: supportedVersions, capabilities, instructions, cacheScope/ttlMs (the client may cache this), serverInfo in _meta")
    r = disc["result"]
    print(f"   supportedVersions={r['supportedVersions']} capabilities={json.dumps(r['capabilities'])} cache={r['cacheScope']}/{r['ttlMs']}ms")
    rpc(p, {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"_meta": META, "name": "echo", "arguments": {"text": "raw bytes"}}},
        "flag: no notifications/initialized; _meta carries the version and capabilities again",
        "resultType: \"complete\" — content[] for any client, structuredContent because the SDK derived an output schema")

    call = {"jsonrpc": "2.0", "id": 3, "method": "tools/call",
            "params": {"_meta": META, "name": "summarize", "arguments": {"text": "MCP separates tools, resources and prompts."}}}
    first = rpc(p, call, "the server needs OUR model for this one",
                "flag: resultType \"input_required\" + inputRequests (keyed) + requestState (opaque, encrypted, must go back unchanged)")
    res = first["result"]
    assert res["resultType"] == "input_required"
    key, req = next(iter(res["inputRequests"].items()))
    print(f"   the server asked for {req['method']} under key {key!r}")
    retry = {"jsonrpc": "2.0", "id": 4, "method": "tools/call",
             "params": {**call["params"], "inputResponses": {key: stub_model_answer(req)}, "requestState": res["requestState"]}}
    rpc(p, retry, "the RETRY: same call, plus inputResponses for that key and the requestState we were given",
        "resultType \"complete\": the tool ran with the model's answer injected as an argument")
    p.stdin.close(); p.wait(timeout=10)
    print(f"   server exited with code {p.returncode} once stdin closed")

    print("\n-- B. legacy handshake (2025-11-25): initialize, reply, notifications/initialized --")
    p = spawn()
    init = rpc(p, {"jsonrpc": "2.0", "id": 1, "method": "initialize",
                   "params": {"protocolVersion": MODERN, "capabilities": CLIENT_CAPS, "clientInfo": {"name": "wire-trace", "version": "1.0.0"}}},
               "flag: we PROPOSE a protocolVersion",
               "flag: result.protocolVersion is what the server CONFIRMED (an initialize request lands you on the legacy version)")
    print(f"   negotiated protocolVersion = {init['result']['protocolVersion']}")
    rpc(p, {"jsonrpc": "2.0", "method": "notifications/initialized"},
        "flag: the handshake is complete only after this notification; before it, only ping and logging are allowed")
    rpc(p, {"jsonrpc": "2.0", "id": 2, "method": "ping"}, "ping exists on the legacy version", "an empty result is the pong")
    tools = rpc(p, {"jsonrpc": "2.0", "id": 3, "method": "tools/list", "params": {}})
    print(f"   tools = {[t['name'] for t in tools['result']['tools']]}")
    p.stdin.close(); p.wait(timeout=10)
    print(f"   server exited with code {p.returncode} once stdin closed")


# ------------------------------------------------------------------ Streamable HTTP
def sse_events(body: str) -> list[dict]:
    """Parse a text/event-stream body: blank-line separated events, `data:` lines carry the JSON."""
    events = []
    for block in body.split("\n\n"):
        data = "\n".join(line[5:].lstrip() for line in block.splitlines() if line.startswith("data:"))
        if data:
            events.append(json.loads(data))
    return events


def trace_http(url: str = "http://127.0.0.1:8765/mcp") -> None:
    import httpx  # pip install httpx

    print(f"== Streamable HTTP: every client message is a POST to {url}; replies come as JSON or as an SSE stream ==")
    accept = {"Accept": "application/json, text/event-stream", "Content-Type": "application/json"}

    def messages(r: httpx.Response) -> list[dict]:
        ct = r.headers.get("content-type", "")
        return sse_events(r.text) if ct.startswith("text/event-stream") else ([r.json()] if r.content else [])

    with httpx.Client(timeout=10) as http:
        print("\n-- A. modern (2026-07-28): server/discover, _meta on every POST, no session required --")
        disc = {"jsonrpc": "2.0", "id": 1, "method": "server/discover", "params": {"_meta": META}}
        show("->", disc, "first, WITHOUT the MCP-Protocol-Version header, to see what the flag does")
        r = http.post(url, json=disc, headers=accept)
        print(f"<- HTTP {r.status_code}: {r.text[:100]!r}")
        print("   flag: without the header the server applies the session rules of the older protocol and refuses")
        def modern_headers(method: str, tool: str | None = None) -> dict[str, str]:
            """The three modern flags ride as HTTP headers so a proxy or cache can route without parsing the body:
            MCP-Protocol-Version (which rules apply), Mcp-Method (must equal the JSON-RPC method), Mcp-Name (the tool)."""
            h = {**accept, "MCP-Protocol-Version": MODERN, "Mcp-Method": method}
            if tool:
                h["Mcp-Name"] = tool
            return h

        show("->", disc, "again, with headers MCP-Protocol-Version: 2026-07-28 and Mcp-Method: server/discover (Accept lists BOTH media types)")
        r = http.post(url, json=disc, headers=modern_headers("server/discover"))
        print(f"<- HTTP {r.status_code} content-type={r.headers.get('content-type')} Mcp-Session-Id={r.headers.get('mcp-session-id')}")
        show("<-", messages(r)[-1], "flag: no session id issued; the answer is plain JSON because nothing needs streaming")
        r = http.post(url, json=disc, headers={**accept, "MCP-Protocol-Version": MODERN, "Mcp-Method": "tools/list"})
        print(f"-> the same body with a WRONG Mcp-Method header\n<- HTTP {r.status_code}: {r.text[:110]!r}")
        print("   flag: the header and the body must agree; the transport checks it before the method runs")
        call = {"jsonrpc": "2.0", "id": 2, "method": "tools/call", "params": {"_meta": META, "name": "negotiation", "arguments": {}}}
        show("->", call, "a tool call straight away: headers + _meta, no session, no initialized notification (Mcp-Name: negotiation)")
        r = http.post(url, json=call, headers=modern_headers("tools/call", "negotiation"))
        print(f"<- HTTP {r.status_code}")
        show("<-", messages(r)[-1]["result"]["structuredContent"], "what the server sees: the version and capabilities we put in _meta")

        print("\n-- B. legacy (2025-11-25): initialize creates a session; every later request must carry its id --")
        init = {"jsonrpc": "2.0", "id": 1, "method": "initialize",
                "params": {"protocolVersion": MODERN, "capabilities": CLIENT_CAPS, "clientInfo": {"name": "wire-trace", "version": "1.0.0"}}}
        show("->", init, "no session header yet: this request creates the session")
        r = http.post(url, json=init, headers=accept)
        session = r.headers.get("mcp-session-id")
        print(f"<- HTTP {r.status_code} content-type={r.headers.get('content-type')} Mcp-Session-Id={session}")
        version = messages(r)[-1]["result"]["protocolVersion"]
        print(f"   negotiated protocolVersion = {version}")
        head = {**accept, "MCP-Protocol-Version": version, **({"Mcp-Session-Id": session} if session else {})}

        note = {"jsonrpc": "2.0", "method": "notifications/initialized"}
        show("->", note, "flag: a notification -> 202 Accepted, empty body")
        r = http.post(url, json=note, headers=head)
        print(f"<- HTTP {r.status_code} (body {len(r.content)} bytes)")

        lst = {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}
        show("->", lst, "same POST, now with Mcp-Session-Id and MCP-Protocol-Version headers")
        r = http.post(url, json=lst, headers=head)
        print(f"<- HTTP {r.status_code} content-type={r.headers.get('content-type')} tools={[t['name'] for t in messages(r)[-1]['result']['tools']]}")

        if session:
            r = http.post(url, json=lst, headers=accept)
            print(f"-> the same request WITHOUT the session header\n<- HTTP {r.status_code}: {r.text[:110]!r}")
            print("   flag: a stateful server rejects requests that do not carry the session id it issued")
            print("-> GET with Accept: text/event-stream opens the standalone stream for server-initiated messages")
            with http.stream("GET", url, headers={"Accept": "text/event-stream", "Mcp-Session-Id": session, "MCP-Protocol-Version": version}) as s:
                print(f"<- HTTP {s.status_code} content-type={s.headers.get('content-type')} (kept open; nothing to deliver, so we close it)")
            r = http.delete(url, headers={"Mcp-Session-Id": session, "MCP-Protocol-Version": version})
            print(f"-> DELETE with the session id\n<- HTTP {r.status_code}: the session is over; the id is now invalid")
        else:
            print("   (started with --stateless: no session id was issued, every request stands alone)")


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else "stdio"
    if mode == "http":
        trace_http(sys.argv[2] if len(sys.argv) > 2 else "http://127.0.0.1:8765/mcp")
    else:
        trace_stdio()

"""
client.py — an MCP client that exercises every primitive of server.py and prints the flags as it goes.

It is the HOST side of the protocol: it owns the model (sampling), the filesystem policy (roots),
the person (elicitation) and the log. Every server-to-client request lands in one of the callbacks below.

Run:   python client.py                                      launches server.py over stdio
       python client.py --transport http                     connects to http://127.0.0.1:8765/mcp
       python client.py --transport http --url http://127.0.0.1:8765/mcp
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path

from mcp import types
from mcp.client import Client
from mcp.client.stdio import StdioServerParameters
from mcp.shared.exceptions import MCPError

HERE = Path(__file__).resolve().parent

# The flags a host keeps track of for one connection. Printed at the end.
FLAGS: dict[str, object] = {
    "initialized": False,            # true once initialize + notifications/initialized are done
    "protocol_version": None,        # the version both sides agreed on
    "server_info": None,             # name/version the server announced
    "server_capabilities": None,     # tools, resources (subscribe? listChanged?), prompts, logging, completions
    "client_capabilities_sent": ["sampling", "roots", "elicitation"],   # implied by the callbacks we register
    "sampling_requests": 0,          # times the server asked our model for a completion
    "roots_requests": 0,             # times the server asked which directories it may touch
    "elicitations": 0,               # times the server asked the person a question
    "progress_notifications": 0,
    "log_messages": 0,
    "other_notifications": [],
}


def step(n: int, title: str) -> None:
    print(f"\n[{n:02d}] {title}")


# ------------------------------------------------------------------ server -> client callbacks
async def on_sampling(context, params: types.CreateMessageRequestParams):
    """sampling/createMessage: the server wants a completion from OUR model.
    A real host would show or filter this, then call its LLM. This stub answers deterministically
    so the lab runs offline; swap the body for an Anthropic or OpenAI call to make it real."""
    FLAGS["sampling_requests"] = int(FLAGS["sampling_requests"]) + 1
    last = params.messages[-1].content
    asked = last.text if isinstance(last, types.TextContent) else str(last)
    prefs = params.model_preferences
    hint = prefs.hints[0].name if prefs and prefs.hints else "any"
    print(f"     ↳ server asked my model (hint: {hint}, max_tokens={params.max_tokens}): {asked[:70]!r}…")
    body = asked.split("\n\n", 1)[-1]
    summary = " ".join(body.split()[:12]) + "…"
    return types.CreateMessageResult(
        role="assistant",
        content=types.TextContent(type="text", text=f"In short: {summary}"),
        model="stub-model-1",
        stop_reason="endTurn",
    )


async def on_list_roots(context):
    """roots/list: the server asks where it may work. We allow exactly this folder."""
    FLAGS["roots_requests"] = int(FLAGS["roots_requests"]) + 1
    print("     ↳ server asked for my roots")
    return types.ListRootsResult(roots=[types.Root(uri=HERE.as_uri(), name="primitives-lab folder")])


async def on_elicit(context, params):
    """elicitation/create: the server needs the PERSON. We auto-accept so the lab runs unattended;
    a real host renders params.message and the requested schema as a form."""
    FLAGS["elicitations"] = int(FLAGS["elicitations"]) + 1
    print(f"     ↳ server asked the person: {getattr(params, 'message', params)!r}")
    return types.ElicitResult(action="accept", content={"sure": True})


async def on_log(params: types.LoggingMessageNotificationParams) -> None:
    FLAGS["log_messages"] = int(FLAGS["log_messages"]) + 1
    print(f"     log/{params.level}: {params.data}")


async def on_progress(progress: float, total: float | None, message: str | None) -> None:
    FLAGS["progress_notifications"] = int(FLAGS["progress_notifications"]) + 1
    print(f"     progress {progress:g}/{total:g}  {message or ''}")


async def on_message(message) -> None:
    """Everything else the server sends that is not a request we handle above (e.g. list-changed)."""
    method = getattr(getattr(message, "root", message), "method", None)
    if method and method not in ("notifications/progress", "notifications/message"):
        FLAGS["other_notifications"].append(method)  # type: ignore[union-attr]
        print(f"     notification: {method}")


# ------------------------------------------------------------------ the walk-through
async def run(client: Client) -> None:
    step(1, "initialize: the handshake fixed the flags for this connection")
    FLAGS["initialized"] = True
    FLAGS["protocol_version"] = str(client.protocol_version)
    FLAGS["server_info"] = f"{client.server_info.name} {client.server_info.version}"
    caps = client.server_capabilities
    FLAGS["server_capabilities"] = caps.model_dump(exclude_none=True) if caps else {}
    print(f"     protocol {FLAGS['protocol_version']} | server {FLAGS['server_info']}")
    print(f"     server capabilities: {json.dumps(FLAGS['server_capabilities'])}")
    print(f"     instructions: {(client.instructions or '')[:90]}…")

    step(2, "tools/list: names, schemas and annotations are the contract")
    tools = await client.list_tools()
    for t in tools.tools:
        a = t.annotations
        flags = "read-only" if a and a.read_only_hint else ("destructive" if a and a.destructive_hint else "writes")
        print(f"     {t.name:12} {flags:11} args={list((t.input_schema or {}).get('properties', {}))}")

    step(3, "tools/call echo: the plain case")
    r = await client.call_tool("echo", {"text": "hello, protocol"})
    print("     ->", r.content[0].text)

    step(4, "tools/call negotiation: the flags as the SERVER sees them")
    r = await client.call_tool("negotiation", {})
    print("     ->", json.dumps(r.structured_content, indent=2).replace("\n", "\n        "))

    step(5, "tools/call count: progress and log notifications while a tool runs")
    r = await client.call_tool("count", {"n": 3}, progress_callback=on_progress)
    print("     ->", r.content[0].text)

    step(6, "tools/call summarize: the server asks OUR model (sampling)")
    r = await client.call_tool("summarize", {"text": NOTE_TEXT})
    print("     ->", r.content[0].text)

    step(7, "tools/call roots: the server asks where it may work (roots)")
    r = await client.call_tool("roots", {})
    print("     ->", r.structured_content)

    step(8, "resources: list, read a template, read a dynamic one, then change the list")
    res = await client.list_resources()
    print("     resources:", [str(x.uri) for x in res.resources])
    tpl = await client.list_resource_templates()
    print("     templates:", [x.uri_template for x in tpl.resource_templates])
    note = await client.read_resource("lab://notes/welcome")
    print("     lab://notes/welcome ->", note.contents[0].text)
    clock = await client.read_resource("lab://clock")
    print("     lab://clock ->", clock.contents[0].text)
    r = await client.call_tool("add_note", {"note_id": "todo", "text": "Read the spec's lifecycle section."})
    print("     add_note ->", r.content[0].text)
    print("     lab://notes ->", (await client.read_resource("lab://notes")).contents[0].text)

    step(9, "prompts: list and get a user-controlled template")
    prompts = await client.list_prompts()
    print("     prompts:", [(p.name, [a.name for a in (p.arguments or [])]) for p in prompts.prompts])
    got = await client.get_prompt("explain", {"concept": "sampling", "level": "senior Java engineer"})
    print("     explain ->", got.messages[0].content.text[:96], "…")

    step(10, "tools/call delete_note: the server asks the PERSON first (elicitation)")
    r = await client.call_tool("delete_note", {"note_id": "todo"})
    print("     ->", r.content[0].text)

    step(11, "ping: a liveness check that 2026-07-28 no longer answers")
    try:
        await client.send_ping()
        print("     -> pong (the server negotiated an older protocol, where ping exists)")
    except MCPError as e:
        print(f"     -> server said {e.message!r}: on 2026-07-28 liveness belongs to the transport, not the protocol")

    step(12, "tools/call with a bad argument: an error result, not a crashed connection")
    r = await client.call_tool("count", {"n": 999})
    print(f"     -> is_error={r.is_error}: {r.content[0].text}")

    print("\nFLAGS after the session:")
    print(json.dumps(FLAGS, indent=2, default=str))


NOTE_TEXT = (
    "The Model Context Protocol separates the three things a model needs: tools it can call, "
    "resources the host can show it, and prompts a person can pick. Sampling, roots and elicitation "
    "run the other way: the server asks the client for a completion, for permitted directories, or for a human answer."
)


async def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--transport", choices=["stdio", "http"], default="stdio")
    p.add_argument("--url", default="http://127.0.0.1:8765/mcp")
    a = p.parse_args()
    if a.transport == "stdio":
        target = StdioServerParameters(command=sys.executable, args=[str(HERE / "server.py"), "--transport", "stdio"])
        print(f"connecting over stdio: {sys.executable} server.py")
    else:
        target = a.url
        print(f"connecting over Streamable HTTP: {a.url}")
    async with Client(
        target,
        sampling_callback=on_sampling,
        list_roots_callback=on_list_roots,
        elicitation_callback=on_elicit,
        logging_callback=on_log,
        message_handler=on_message,
        client_info=types.Implementation(name="primitives-lab-client", version="1.0.0"),
    ) as client:
        await run(client)


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

"""
test_lab.py — every primitive, on three transports.

    pytest -q

In-memory: the Client talks to the MCPServer object directly (no process, no socket): fastest, for logic.
stdio:     the Client launches server.py as a subprocess, as Claude Desktop or Claude Code would.
HTTP:      server.py --transport streamable-http runs in a subprocess; the Client connects by URL.
"""
from __future__ import annotations

import socket
import subprocess
import sys
import time
from pathlib import Path

import pytest

from mcp import types
from mcp.client import Client
from mcp.client.stdio import StdioServerParameters

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import server as lab  # noqa: E402  (the module under test; importing it does not start a transport)

pytestmark = pytest.mark.anyio


@pytest.fixture
def anyio_backend():
    return "asyncio"


# ---------------------------------------------------------------- the host side, recorded
class Host:
    """The callbacks a real host implements, with counters so tests can assert they were used."""

    def __init__(self) -> None:
        self.sampling = 0
        self.roots = 0
        self.elicit = 0
        self.progress: list[tuple[float, float | None, str | None]] = []
        self.say_yes = True

    async def on_sampling(self, context, params: types.CreateMessageRequestParams):
        self.sampling += 1
        return types.CreateMessageResult(role="assistant", model="stub", stop_reason="endTurn",
                                         content=types.TextContent(type="text", text="a one-sentence summary"))

    async def on_roots(self, context):
        self.roots += 1
        return types.ListRootsResult(roots=[types.Root(uri=HERE.as_uri(), name="lab")])

    async def on_elicit(self, context, params):
        self.elicit += 1
        return types.ElicitResult(action="accept", content={"sure": self.say_yes})

    async def on_progress(self, progress, total, message):
        self.progress.append((progress, total, message))

    def client(self, target) -> Client:
        return Client(target, sampling_callback=self.on_sampling, list_roots_callback=self.on_roots,
                      elicitation_callback=self.on_elicit,
                      client_info=types.Implementation(name="test-host", version="0"))


async def exercise_everything(client: Client, host: Host) -> None:
    """The same assertions for every transport."""
    # negotiation flags
    assert str(client.protocol_version) == "2026-07-28"
    assert client.server_info.name == "primitives-lab"
    caps = client.server_capabilities
    assert caps.tools is not None and caps.resources is not None and caps.prompts is not None

    # tools/list: names and annotations
    tools = {t.name: t for t in (await client.list_tools()).tools}
    assert set(tools) == {"echo", "negotiation", "count", "summarize", "roots", "add_note", "delete_note"}
    assert tools["echo"].annotations.read_only_hint is True
    assert tools["delete_note"].annotations.destructive_hint is True

    # plain call, structured result
    assert (await client.call_tool("echo", {"text": "hi"})).content[0].text == "hi"
    neg = (await client.call_tool("negotiation", {})).structured_content
    assert neg["protocol_version"] == "2026-07-28"
    assert "sampling" in neg["client_capabilities"] and "roots" in neg["client_capabilities"]

    # progress notifications
    r = await client.call_tool("count", {"n": 3}, progress_callback=host.on_progress)
    assert r.content[0].text == "counted to 3"
    assert [p[0] for p in host.progress] == [1, 2, 3] and host.progress[-1][2] == "step 3 of 3"

    # sampling: the server used OUR model
    r = await client.call_tool("summarize", {"text": "MCP separates tools, resources and prompts."})
    assert host.sampling == 1 and r.content[0].text == "[stub] a one-sentence summary"

    # roots: the server asked where it may work
    r = await client.call_tool("roots", {})
    assert host.roots == 1 and r.structured_content["result"] == [f"lab -> {HERE.as_uri()}"]

    # resources: list, template, dynamic, and a change
    uris = {str(x.uri) for x in (await client.list_resources()).resources}
    assert {"lab://notes", "lab://clock"} <= uris
    assert [t.uri_template for t in (await client.list_resource_templates()).resource_templates] == ["lab://notes/{note_id}"]
    assert (await client.read_resource("lab://notes/welcome")).contents[0].text.startswith("Resources are")
    assert "T" in (await client.read_resource("lab://clock")).contents[0].text
    assert (await client.call_tool("add_note", {"note_id": "t1", "text": "x"})).content[0].text == "lab://notes/t1"
    assert (await client.read_resource("lab://notes/t1")).contents[0].text == "x"

    # prompts: user-controlled templates
    assert [p.name for p in (await client.list_prompts()).prompts] == ["explain"]
    got = await client.get_prompt("explain", {"concept": "roots", "level": "beginner"})
    assert "Explain 'roots' to a beginner" in got.messages[0].content.text

    # elicitation: the person is asked, and can say no
    host.say_yes = False
    assert (await client.call_tool("delete_note", {"note_id": "t1"})).content[0].text.startswith("kept t1")
    host.say_yes = True
    assert (await client.call_tool("delete_note", {"note_id": "t1"})).content[0].text == "deleted t1"
    assert host.elicit == 2

    # errors are results, not crashes
    r = await client.call_tool("count", {"n": 999})
    assert r.is_error and "between 1 and 50" in r.content[0].text
    r = await client.call_tool("delete_note", {"note_id": "nope"})
    assert r.is_error and "no note called nope" in r.content[0].text


# ---------------------------------------------------------------- three transports
async def test_in_memory():
    host = Host()
    async with host.client(lab.mcp) as client:
        await exercise_everything(client, host)


async def test_stdio_subprocess():
    host = Host()
    params = StdioServerParameters(command=sys.executable, args=[str(HERE / "server.py"), "--transport", "stdio"])
    async with host.client(params) as client:
        await exercise_everything(client, host)


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


@pytest.fixture
def http_server():
    port = _free_port()
    proc = subprocess.Popen([sys.executable, str(HERE / "server.py"), "--transport", "streamable-http", "--port", str(port)],
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        for _ in range(50):
            with socket.socket() as s:
                s.settimeout(0.2)
                if s.connect_ex(("127.0.0.1", port)) == 0:
                    break
            time.sleep(0.2)
        else:
            raise RuntimeError("HTTP server did not start")
        yield f"http://127.0.0.1:{port}/mcp"
    finally:
        proc.terminate()
        proc.wait(timeout=10)


async def test_streamable_http(http_server):
    host = Host()
    async with host.client(http_server) as client:
        await exercise_everything(client, host)


# ---------------------------------------------------------------- the raw wire, no SDK
def test_wire_trace_stdio_runs():
    out = subprocess.run([sys.executable, str(HERE / "wire_trace.py"), "stdio"], capture_output=True, text=True,
                         encoding="utf-8", timeout=60)
    assert out.returncode == 0, out.stderr
    assert '"method": "server/discover"' in out.stdout and 'input_required' in out.stdout
    assert "negotiated protocolVersion = 2025-11-25" in out.stdout   # the legacy path, for comparison


def test_wire_trace_http_runs(http_server):
    out = subprocess.run([sys.executable, str(HERE / "wire_trace.py"), "http", http_server], capture_output=True,
                         text=True, encoding="utf-8", timeout=60)
    assert out.returncode == 0, out.stderr
    assert "Mcp-Session-Id=" in out.stdout and "HTTP 202" in out.stdout and "Missing session ID" in out.stdout