A tool that raises

import asyncio
import itertools
import json
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)


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()

        # a counted id, not a UUID, so the printed output is the same every run
        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 = {}

    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)

        # the tool list is fetched once per desk, not once per question
        if not self.tools:
            self.tools = await self.find_tools()
        await self.send({"type": "system", "subtype": "init", "model": "shop-desk",
                         "session_id": self.session_id, "tools": list(self.tools)})

    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"])

    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 find_tools(self):
        tools = {}
        for server in self.options.mcp_servers:
            # MCP wants a handshake before a server will list anything
            await self.mcp(server, "initialize", {"protocolVersion": "2025-06-18",
                           "capabilities": {}, "clientInfo": {"name": "desk", "version": "1"}})

            for tool in (await self.mcp(server, "tools/list", {}))["tools"]:
                tools[f"mcp__{server}__{tool['name']}"] = (server, tool["name"])
        return 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):
        # the full name carries the server it came from, so the call goes back there
        server, tool = self.tools[name]
        result = await self.mcp(server, "tools/call", {"name": tool, "arguments": 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])