Turn and budget limits

import asyncio
import itertools
import json
import os
import re
from fnmatch import fnmatch

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

ORDER = re.compile(r"\b[A-Z]\d+\b")
SESSIONS = itertools.count(1)
WANTS = {"refund": "refund_order", "stock": "stock_level"}
RATE = 0.002  # the desk's own price for a word, since it has no model and no bill


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}


def matches(rule, name, args):
    """Does one allow or deny rule match this call?"""
    head, _, scope = rule.partition("(")
    if not fnmatch(name, head):
        return False

    # a bare rule covers the whole tool; a scoped one only matching calls
    target = args.get("file_path") or args.get("command") or ""
    return not scope or fnmatch(target, scope.rstrip(")"))


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

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

    def reply(self, said, results, tools):
        if self.script:
            return self.script.pop(0)

        # a search result is not an answer, so a searching turn does not stop here
        if results and not self.searched:
            return {"text": " ".join(results)}
        orders = ORDER.findall(said[-1]) or ORDER.findall(" ".join(said))[-1:]
        wanted = next((name for word, name in WANTS.items() if word in said[-1].lower()),
                      "lookup_order")
        tool = next((name for name in tools if name.endswith(wanted)), None)

        # the tool this message needs has no definition in context: search for it first
        if orders and tool is None and not self.searched and "ToolSearch" in tools:
            self.searched = True
            return {"calls": [("ToolSearch", {"query": wanted.replace("_", " ")})]}

        # the definition is in context now, so the call itself can be made
        if orders and tool:
            self.searched = False
            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 = {}
        self.external, self.servers = {}, []
        self.deferring = self.options.env.get("ENABLE_TOOL_SEARCH", "false").startswith(
            ("true", "auto"))

        # spent is the desk's own running estimate, and failed is the exit a cap causes
        self.loaded, self.spent, self.failed = set(), 0.0, None

    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:
            # an exception in the outbox is raised at the SDK, as a dead process would
            if isinstance(message, Exception):
                raise message
            yield message

    async def listen(self):
        while (message := await self.inbox.get()) is not None:
            if message["type"] == "control_request":
                if message["request"]["subtype"] == "initialize":
                    self.setup = message["request"]
                await self.respond(message["request_id"])
            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 pass on any failure
        await self.prompts.put(None)
        await self.worker
        await self.send(self.failed)

    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)
        results, turns = [], 0

        # the caps are read inside the loop, so a run is stopped between turns
        while "calls" in (step := self.model.reply(self.said, results, self.in_context())):
            turns += 1
            if stop := self.capped(turns):
                return await self.finish(*stop)
            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):
        asked = self.options.tools
        if asked is None or isinstance(asked, dict):
            return list(BUILT_IN) + ["AskUserQuestion"]

        # a tools list that leaves AskUserQuestion out takes away Claude's questions
        return [name for name in asked if name in BUILT_IN or name == "AskUserQuestion"]

    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 = []

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

        # a deny rule with no scope removes the tool before any call is made
        bare = [rule for rule in self.options.disallowed_tools if "(" not in rule]
        return {name: where for name, where in tools.items()
                if not any(fnmatch(name, rule) for rule in bare)}

    def in_context(self):
        # what the model has definitions for: everything, or the searcher plus what it found
        if not self.deferring:
            return dict(self.tools)
        found = {name: self.tools[name] for name in self.loaded}
        return {"ToolSearch": ("built-in", "ToolSearch"), **found}

    def search(self, query):
        words = set(query.lower().split())

        # the search matches the query against tool names, as the docs describe
        hits = [name for name in self.tools
                if words & set(name.lower().replace("mcp__", "").split("_"))]

        # five at most, and each one's definition stays in context from now on
        self.loaded.update(hits[:5])
        return "Loaded: " + ", ".join(hits[:5]) if hits else "No tool matched."

    async def use(self, name, args, parent=None):
        call = f"toolu_{next(self.ids)}"
        block = {"type": "tool_use", "id": call, "name": name, "input": args}
        await self.send(self.message("assistant", [block], parent))

        # a refused call reuses args as the text to send back, so Claude is told why
        allowed, args = await self.permit(name, args, call)
        text, failed = (await self.run(name, args, call)) if allowed else (args, True)

        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 answers you collected are what the question tool returns to the model
        if name == "AskUserQuestion":
            picked = args.get("answers") or {}
            return "; ".join(f"{q}: {a}" for q, a in picked.items()), False
        if name == "ToolSearch":
            return self.search(args["query"]), False

        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
        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": self.spent,
                  "usage": {"input_tokens": len(" ".join(self.said).split()),
                            "output_tokens": len(text.split())}}

        # total_cost_usd is now the desk's own arithmetic, not a bill from anyone
        result.update({"result": text} if ok else {"errors": [text]})
        await self.send(result)

        # a real Claude Code process exits non-zero after a failure, so note one here
        self.failed = ProcessError("Command failed", exit_code=1) if not ok else None

    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}

    async def permit(self, name, args, call):
        o = self.options
        if any(matches(rule, name, args) for rule in o.disallowed_tools):
            return False, f"Permission to use {name} was denied by a rule."

        # a question is always put to you, whatever the rules and the mode say
        asks = name == "AskUserQuestion"
        if not asks and (o.permission_mode == "bypassPermissions"
                         or any(matches(rule, name, args) for rule in o.allowed_tools)):
            return True, args

        # dontAsk never prompts, so everything still undecided is denied
        if o.can_use_tool is None or o.permission_mode == "dontAsk":
            return False, f"Permission to use {name} was not granted."
        return await self.ask_permission(name, args, call)

    async def ask_permission(self, name, args, call):
        # a suggestion is the rule your code can echo back to stop the asking
        suggest = [{"type": "addRules", "behavior": "allow", "destination": "session",
                    "rules": [{"toolName": name}]}]

        decision = (await self.ask({"subtype": "can_use_tool", "tool_name": name,
                                    "input": args, "tool_use_id": call,
                                    "permission_suggestions": suggest,
                                    "blocked_path": None}))["response"]
        if decision["behavior"] != "allow":
            return False, decision["message"]

        # a rule you sent back is added to the session, so the next call matches it
        for update in decision.get("updatedPermissions") or []:
            for rule in update["rules"]:
                self.options.allowed_tools.append(rule["toolName"])
        return True, decision["updatedInput"]

    def capped(self, turns):
        """Which cap this turn has passed, as the arguments finish takes."""
        limit = self.options.max_turns
        if limit and turns > limit:
            return "error_max_turns", turns, f"Reached maximum number of turns ({limit})"

        # a turn re-sends the conversation, so the desk adds its words again at RATE
        self.spent = round(self.spent + len(" ".join(self.said).split()) * RATE, 4)
        if self.options.max_budget_usd and self.spent >= self.options.max_budget_usd:
            return ("error_max_budget_usd", turns,
                    f"Exceeded max budget: ${self.options.max_budget_usd}")
        return None


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."}
REFUNDED = []


# 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}]}


# the money moves, so this is the call that must not run unasked
@tool("refund_order", "Refund an order by its id. The money goes back to the customer.",
      {"order_id": str})
async def refund_order(args):
    # recorded so a lesson can show whether the refund happened at all
    REFUNDED.append(args["order_id"])
    text = f"Order {args['order_id']} has been refunded."
    return {"content": [{"type": "text", "text": text}]}


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