External MCP servers

import asyncio
import itertools
import json
import os
import re

from claude_agent_sdk import (AssistantMessage, ClaudeAgentOptions, ResultMessage,
                              SystemMessage, TextBlock, ToolResultBlock, ToolUseBlock,
                              Transport, UserMessage, query)

ORDER = re.compile(r"\b[A-Z]\d+\b")
SESSIONS = itertools.count(1)


def read_file(args):
    """What the desk serves for a Read call: the file's lines, numbered."""
    lines = open(args["file_path"]).read().splitlines()
    return "\n".join(f"{number:>5}  {line}" for number, line in enumerate(lines, 1))


BUILT_IN = {"Read": read_file}


class ShopModel:
    """Stands in for Claude: reads the conversation, picks the next step."""

    def __init__(self, script=None):
        self.script = list(script or [])

    def reply(self, said, results, tools):
        # a scripted reply lets a lesson force a step the stand-in would not pick
        if self.script:
            return self.script.pop(0)
        if results:
            return {"text": " ".join(results)}

        # the newest message first, then the last id named earlier in the chat
        orders = ORDER.findall(said[-1]) or ORDER.findall(" ".join(said))[-1:]
        wanted = "refund_order" if "refund" in said[-1].lower() else "lookup_order"
        tool = next((name for name in tools if name.endswith(wanted)), None)
        if orders and tool:
            return {"calls": [(tool, {"order_id": order}) for order in orders]}
        return {"text": f"I cannot help with order {orders[0]} yet." if orders
                else "Which order is this about?"}


class ShopDesk(Transport):
    """Stands in for the Claude Code binary: it runs the loop over JSON lines."""

    def __init__(self, options=None, model=None):
        self.options = options or ClaudeAgentOptions()
        self.model = model or ShopModel()
        self.session_id = f"session-{next(SESSIONS)}"
        self.inbox, self.outbox = asyncio.Queue(), asyncio.Queue()
        self.prompts, self.said = asyncio.Queue(), []
        self.ids, self.setup = itertools.count(1), {}
        self.ready, self.tools = False, {}
        self.waiting = {}

        # one running process per stdio server, and a status line per configured server
        self.external, self.servers = {}, []

    async def connect(self):
        self.ready = True

    def is_ready(self):
        return self.ready

    async def write(self, data):
        # one JSON line at a time, which is what the SDK writes to a real process
        await self.inbox.put(json.loads(data))

    async def end_input(self):
        # for a real process this would close its standard input
        await self.inbox.put(None)

    async def close(self):
        self.ready = False

    async def send(self, message):
        await self.outbox.put(message)

    async def read_messages(self):
        # the listener and the worker run alongside the SDK's own read loop
        asyncio.create_task(self.listen())
        self.worker = asyncio.create_task(self.work())

        while (message := await self.outbox.get()) is not None:
            if isinstance(message, Exception):
                raise message
            yield message

    async def listen(self):
        while (message := await self.inbox.get()) is not None:
            # a control request is the SDK asking a question; initialize is first
            if message["type"] == "control_request":
                if message["request"]["subtype"] == "initialize":
                    self.setup = message["request"]
                await self.respond(message["request_id"])

            # a control response answers a question the desk itself asked
            elif message["type"] == "control_response":
                answer = message["response"]
                self.waiting.pop(answer["request_id"]).set_result(answer)
            elif message["type"] == "user":
                await self.prompts.put(message["message"]["content"])

        # the input has ended: finish the work in hand, then close the stream
        await self.prompts.put(None)
        await self.worker
        await self.send(None)

    async def respond(self, request_id):
        # the stand-in says yes to every request the SDK makes of it
        await self.send({"type": "control_response", "response": {
            "subtype": "success", "request_id": request_id, "response": {}}})

    async def ask(self, request):
        # a future per request, so the answer can be awaited where it was asked
        request_id = f"desk_{next(self.ids)}"
        self.waiting[request_id] = asyncio.get_running_loop().create_future()
        await self.send({"type": "control_request", "request_id": request_id,
                         "request": request})
        return await self.waiting[request_id]

    async def work(self):
        while (prompt := await self.prompts.get()) is not None:
            await self.turn(prompt)

    async def start(self, prompt):
        self.said.append(prompt)
        if not self.tools:
            self.tools = await self.find_tools()

        # the init message reports what each configured server is doing
        await self.send({"type": "system", "subtype": "init", "model": "shop-desk",
                         "session_id": self.session_id, "tools": list(self.tools),
                         "mcp_servers": self.servers})

    async def turn(self, prompt):
        await self.start(prompt)

        # keep asking the model until it answers instead of calling a tool
        results, turns = [], 0
        while "calls" in (step := self.model.reply(self.said, results, self.tools)):
            turns += 1
            results = [await self.use(name, args) for name, args in step["calls"]]

        await self.say(step["text"])
        await self.finish("success", turns + 1, step["text"])

    def offered(self):
        # the tools option is the availability layer: None means every built-in the
        # desk has, a list names the ones to keep, and [] leaves Claude none
        asked = self.options.tools
        if asked is None or isinstance(asked, dict):
            return list(BUILT_IN)
        return [name for name in asked if name in BUILT_IN]

    async def mcp(self, server, method, params):
        message = {"jsonrpc": "2.0", "id": next(self.ids), "method": method,
                   "params": params}
        answer = await self.ask({"subtype": "mcp_message", "server_name": server,
                                 "message": message})
        return answer["response"]["mcp_response"]["result"]

    async def connect_stdio(self, name, config):
        # Claude Code starts the command and talks MCP over its standard input
        process = await asyncio.create_subprocess_exec(
            config["command"], *config.get("args", []),
            stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,
            env={**os.environ, **config.get("env", {})})
        self.external[name] = process

        hello = await self.stdio(name, "initialize", {"protocolVersion": "2025-06-18"})
        if hello is None:
            return {"name": name, "status": "failed",
                    "error": "the server gave no answer to initialize"}
        return {"name": name, "status": "connected", "serverInfo": hello["serverInfo"]}

    async def stdio(self, name, method, params):
        process = self.external[name]
        request = {"jsonrpc": "2.0", "id": next(self.ids), "method": method,
                   "params": params}
        process.stdin.write((json.dumps(request) + "\n").encode())
        await process.stdin.drain()

        # a server that died writes nothing and its pipe reaches end of file
        line = await process.stdout.readline()
        return json.loads(line)["result"] if line else None

    async def list_from(self, name, config):
        # an in-process server is handed over by the SDK; a stdio one is a command to run
        if not isinstance(config, dict) or config.get("type") == "sdk":
            await self.mcp(name, "initialize", {"protocolVersion": "2025-06-18",
                           "capabilities": {}, "clientInfo": {"name": "desk", "version": "1"}})
            listed = await self.mcp(name, "tools/list", {})
            return {"name": name, "status": "connected", "tools": listed["tools"]}

        status = await self.connect_stdio(name, config)
        if status["status"] == "connected":
            status["tools"] = (await self.stdio(name, "tools/list", {}))["tools"]
        return status

    async def find_tools(self):
        tools = {name: ("built-in", name) for name in self.offered()}
        self.servers = []

        # one status per configured server, whichever kind of server it is
        for name, config in self.options.mcp_servers.items():
            status = await self.list_from(name, config)
            self.servers.append(status)
            for tool in status.pop("tools", []):
                tools[f"mcp__{name}__{tool['name']}"] = (name, tool["name"])

        return {name: where for name, where in tools.items()
                if name not in self.options.disallowed_tools}

    async def use(self, name, args, parent=None):
        # one id ties the request and its result together, as the message pair does
        call = f"toolu_{next(self.ids)}"
        block = {"type": "tool_use", "id": call, "name": name, "input": args}
        await self.send(self.message("assistant", [block], parent))

        text, failed = await self.run(name, args, call)
        result = {"type": "tool_result", "tool_use_id": call, "content": text,
                  "is_error": failed}
        await self.send(self.message("user", [result], parent))
        return text

    async def run(self, name, args, call):
        where, tool = self.tools[name]
        if where == "built-in":
            try:
                return BUILT_IN[tool](args), False
            except Exception as error:
                return f"{type(error).__name__}: {error}", True

        # the same tools/call either way, over the channel that server came in on
        call_args = {"name": tool, "arguments": args}
        result = (await self.stdio(where, "tools/call", call_args) if where in self.external
                  else await self.mcp(where, "tools/call", call_args))
        return result["content"][0]["text"], result.get("isError", False)

    async def say(self, text):
        # with include_partial_messages on, the text goes out a word at a time
        if self.options.include_partial_messages:
            for word in text.split(" "):
                delta = {"type": "text_delta", "text": word + " "}
                await self.send({"type": "stream_event", "uuid": f"ev_{next(self.ids)}",
                                 "session_id": self.session_id, "event": {
                                     "type": "content_block_delta", "index": 0,
                                     "delta": delta}})

        # the whole message still follows, so code written earlier keeps working
        await self.send(self.message("assistant", [{"type": "text", "text": text}]))

    async def finish(self, subtype, turns, text):
        ok = subtype == "success"
        result = {"type": "result", "subtype": subtype, "is_error": not ok,
                  "duration_ms": 0, "duration_api_ms": 0, "num_turns": turns,
                  "session_id": self.session_id, "total_cost_usd": 0.0,
                  "usage": {"input_tokens": len(" ".join(self.said).split()),
                            "output_tokens": len(text.split())}}

        # cost is zero and usage counts words, because no model was called
        result.update({"result": text} if ok else {"errors": [text]})
        await self.send(result)

    def message(self, kind, content, parent=None):
        body = {"role": kind, "content": content}
        if kind == "assistant":
            body.update(model="shop-desk", id=f"msg_{next(self.ids)}")
        return {"type": kind, "message": body, "parent_tool_use_id": parent,
                "session_id": self.session_id}


def describe(message):
    # the two messages with no content blocks are read from their own fields
    if isinstance(message, SystemMessage):
        return f"system     {message.subtype}"
    if isinstance(message, ResultMessage):
        return f"result     {message.subtype}: {message.result or message.errors}"

    # a type this course does not use prints its name rather than failing
    if not isinstance(message, (AssistantMessage, UserMessage)):
        return type(message).__name__
    return describe_block(message.content[0])


def describe_block(block):
    if isinstance(block, TextBlock):
        return f"assistant  {block.text}"
    if isinstance(block, ToolUseBlock):
        return f"assistant  calls {block.name} {block.input}"

    # a tool result sits on the user's side, and its error flag is worth seeing
    if isinstance(block, ToolResultBlock):
        flag = " (error)" if block.is_error else ""
        return f"user       result{flag}: {block.content}"
    return type(block).__name__


async def watch(prompt, options, model=None):
    """Run one question through a fresh desk and print every message that comes back."""
    async for message in query(prompt=prompt, options=options,
                               transport=ShopDesk(options, model)):
        print(describe(message))

from claude_agent_sdk import create_sdk_mcp_server, tool

ORDERS = {"A17": "Order A17 shipped on 3 March."}


# the description is what Claude reads to decide this is the right tool
@tool("lookup_order", "Look up an order's shipping status by its id, such as A17.",
      {"order_id": str})
async def lookup_order(args):
    order = args["order_id"]

    # an unknown id is answered in words, not raised, so Claude can pass it on
    status = ORDERS.get(order, f"No order {order} was found.")
    return {"content": [{"type": "text", "text": status}]}


shop = create_sdk_mcp_server("shop", tools=[lookup_order])