Build an MCP Client and Host: Handshake, Tool Discovery, Error Handling and Approvals

"""An MCP host for Pebble & Pine's support desk: a client for each MCP server, and a model that uses their tools.
You write the parts marked TODO, one step at a time."""
import json
import queue
import re
import subprocess
import threading
import time

import harness as H

PROTOCOL = "2025-06-18"
SUPPORTED = {"2025-06-18", "2025-03-26", "2024-11-05"}
SERVERS = {"orders": ["python3", "orders_server.py"], "tickets": ["python3", "tickets_server.py"]}


class McpError(Exception):
    """The server answered with a JSON-RPC error object."""

    def __init__(self, code, message):
        super().__init__(f"{code}: {message}")
        self.code, self.message = code, message


class ServerGone(Exception):
    """The server process exited or closed its output."""


class StdioClient:
    """One MCP server, started as a subprocess. Messages are JSON-RPC 2.0, one JSON object per line: we write to its
    stdin and read its stdout. Its stderr is its log, and goes to <name>.stderr.log."""

    def __init__(self, name, command):
        self.name = name
        self.proc = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                     stderr=open(f"{name}.stderr.log", "w"), text=True, bufsize=1)
        self.inbox = queue.Queue()   # parsed messages from the server; None when it has gone
        self.notifications = []      # messages with a method and no id
        self.stale = 0               # responses to requests we had stopped waiting for
        self.next_id = 1
        self.server_info, self.capabilities = None, None
        threading.Thread(target=self._reader, daemon=True).start()

    def _reader(self):
        for line in self.proc.stdout:
            line = line.strip()
            if line:
                try:
                    self.inbox.put(json.loads(line))
                except ValueError:
                    pass  # not a JSON-RPC message
        self.inbox.put(None)

    def send(self, message):
        try:
            self.proc.stdin.write(json.dumps(message) + "\n")
            self.proc.stdin.flush()
        except (BrokenPipeError, ValueError, OSError):
            raise ServerGone(f"{self.name} is not running")

    def notify(self, method, params=None):
        """A notification: a message with no id, which gets no response."""
        msg = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            msg["params"] = params
        self.send(msg)

    # ---------- Step 1: request and handshake ----------

    def request(self, method, params=None, timeout=10.0):
        """Send a request with the next id and wait for the response with that id: return its "result", or raise
        McpError for an "error". While waiting, keep notifications in self.notifications and count responses to
        other ids in self.stale, and answer requests from the server (a "ping" with an empty result, anything else
        with error -32601). Raise ServerGone if the server goes (None in the inbox). After `timeout` seconds, tell
        the server with a notifications/cancelled for this id, and raise TimeoutError."""
        # TODO (Step 1): send {jsonrpc, id, method, params}; read self.inbox until the response with that id;
        # keep notifications, count stale responses, answer server requests; ServerGone on None; cancel + TimeoutError.
        raise NotImplementedError("Step 1: write request()")

    def initialize(self):
        """The MCP handshake: an "initialize" request (protocolVersion PROTOCOL, capabilities {}, clientInfo with a
        name and version), check the server's protocolVersion is in SUPPORTED (else close() and raise McpError
        with code -32602), keep serverInfo and capabilities, then send the notifications/initialized notification."""
        # TODO (Step 1): request initialize, check protocolVersion in SUPPORTED, keep serverInfo/capabilities,
        # then notify notifications/initialized.
        raise NotImplementedError("Step 1: write initialize()")

    def close(self):
        try:
            self.proc.stdin.close()
        except OSError:
            pass
        try:
            self.proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.proc.kill()

    # ---------- Step 2: discover tools ----------

    def list_tools(self):
        """Every tool the server offers: call tools/list, and while the result has a "nextCursor", call it again with
        {"cursor": that value}."""
        raise NotImplementedError("list_tools() arrives in Step 2")


NAME_OK = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")


def to_openai(server, tool):
    """The tool as an OpenAI function tool, named "<server>__<tool name>" so that two servers' tools cannot clash;
    the description starts with "[<server>] "; parameters is the inputSchema. Raise ValueError if the name does not
    match NAME_OK."""
    raise NotImplementedError("to_openai() arrives in Step 2")


def split_name(name):
    """("<server>", "<tool>") from "<server>__<tool>" (split at the first "__")."""
    raise NotImplementedError("split_name() arrives in Step 2")


# ---------- Step 3: call tools, and every way that can fail ----------

def content_text(content):
    """The text of a tool result's content list: text items joined with newlines; any other item as
    "[<type> omitted]"."""
    raise NotImplementedError("content_text() arrives in Step 3")


def call_tool(client, tool, arguments, timeout=5.0):
    """{"ok": bool, "text": str} for a tools/call, never raising: a result is ok unless "isError" is true; a McpError
    gives "MCP error <code>: <message>", a TimeoutError "timed out after <timeout>s", ServerGone "server <name> has
    stopped"."""
    raise NotImplementedError("call_tool() arrives in Step 3")


# ---------- Step 4: the host ----------

SYSTEM = ("You are Pebble & Pine's support assistant. Use the tools to look things up and act; never guess order, "
          "parcel or ticket details. If a tool fails or an action is declined, say so plainly. Answer in at most "
          "three sentences.")

# ---------- Step 5: who approves what ----------

TRUSTED = {"orders"}                        # servers whose annotations we trust (we run them)
REVIEWED = {"tickets__search", "tickets__get_ticket"}  # tools on other servers we have read and allow freely


def needs_approval(name, tool):
    """True unless the tool is known safe. On a TRUSTED server, trust its annotations: readOnlyHint true means no
    approval; otherwise destructiveHint decides, and it counts as true when missing. On any other server, only the
    tools in REVIEWED run without approval."""
    raise NotImplementedError("needs_approval() arrives in Step 5")


class Host:
    """Starts and initializes a client per server and exposes all their tools to the model."""

    def __init__(self, servers=SERVERS):
        self.clients = {name: StdioClient(name, cmd) for name, cmd in servers.items()}
        for c in self.clients.values():
            c.initialize()
        self.tools, self.specs = self.discover()
        self.log = []

    def discover(self):
        """(tools, specs): tools maps "<server>__<tool>" to the MCP tool dict; specs is the list of to_openai() tools."""
        raise NotImplementedError("discover() arrives in Step 4")

    def execute(self, call, approve):
        """Run one tool call from the model ({"function": {"name", "arguments"}}): {"ok", "text"}. Unknown names and
        arguments that are not a JSON object fail without reaching a server. Tools that need_approval() only run if
        approve(name, arguments) returns True; otherwise the result is not ok and says the user declined."""
        raise NotImplementedError("execute() arrives in Step 4")
        raise NotImplementedError("execute() arrives in Step 4")

    def run(self, question, approve=lambda name, args: True, max_turns=8):
        """The tool loop: send SYSTEM and the question with self.specs to H.chat(); while the reply has tool_calls,
        append it, execute() each call and append {"role": "tool", "tool_call_id", "content": result text}, then ask
        again. Return the final content, or "Stopped after <max_turns> turns." Every executed call is appended to
        self.log as (name, arguments string, result)."""
        raise NotImplementedError("run() arrives in Step 4")

    def close(self):
        for c in self.clients.values():
            c.close()