Session storage with a SessionStore

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

from shop_desk import ShopDesk


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):
    async for message in query(prompt=prompt, options=options,
                               transport=ShopDesk(options, model)):
        print(describe(message))

import asyncio
import itertools
import json
import os
import re
from pathlib import Path

from claude_agent_sdk import (ClaudeAgentOptions, ProcessError, Transport,
                              project_key_for_directory)

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

# Claude Code keeps one .jsonl file per session; the desk keeps one dictionary
SAVED = {}


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, self.marks = asyncio.Queue(), [], []
        self.ids, self.setup, self.waiting = itertools.count(1), {}, {}
        self.tools, self.failed, self.ready = {}, None, False

    async def connect(self):
        # a refused resume has to fail where the SDK starts the session
        self.restore()
        self.ready = True

    def restore(self):
        o = self.options
        source = o.resume or (self.newest() if o.continue_conversation else None)
        if source is None or source not in SAVED:
            return
        marks, said = SAVED[source]
        if o.resume_session_at:
            marks, said = self.truncate(marks, said)

        # the loaded turns are the context the next question is answered against
        self.marks, self.said = list(marks), list(said)

        # a fork keeps the copied history under its own new id, so the original
        # session is never appended to again
        self.session_id = self.session_id if o.fork_session else source

    def newest(self):
        # continue_conversation takes the most recent session and needs no id
        return list(SAVED)[-1] if SAVED else None

    def truncate(self, marks, said):
        # resume_session_at loads up to and including that message, and no further
        end = marks.index(self.options.resume_session_at) + 1
        drops = self.options.resume_drops_turn

        # resume_drops_turn names the one turn the cut is allowed to throw away
        if drops is not None and marks[end:end + 1] != [drops]:
            raise ProcessError("Resume rejected by --resume-drops-turn: "
                               "the dropped range is not that turn", exit_code=1)
        return marks[:end], said[:end]

    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 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)
        self.marks.append(f"uuid_{len(self.marks) + 1}")

        # replay-user-messages asks for the prompt back with the uuid it was
        # stored under, which is the only way to learn that id
        if "replay-user-messages" in self.options.extra_args:
            await self.send({"type": "user", "uuid": self.marks[-1],
                             "parent_tool_use_id": None, "session_id": self.session_id,
                             "message": {"role": "user", "content": prompt}})
        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):
        # UserPromptSubmit fires before the model is given the prompt
        await self.hooks("UserPromptSubmit", {"prompt": prompt})
        await self.start(prompt)
        results, turns, limit = [], 0, self.options.max_turns
        while "calls" in (step := self.model.reply(self.said, results, self.tools)):
            turns += 1

            # max_turns is the desk's own stop, not something the model respects
            if limit and turns > limit:
                message = f"Reached maximum number of turns ({limit})"
                return await self.finish("error_max_turns", turns, message)
            results = [await self.use(name, args) for name, args in step["calls"]]

        await self.say(step["text"])

        # Stop fires once the model has answered and wants no more tools
        await self.hooks("Stop", {"stop_hook_active": False})
        await self.finish("success", turns + 1, step["text"])

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

        # the turn is on the shelf now, which is what a resume will read back
        SAVED[self.session_id] = (list(self.marks), list(self.said))

        # the frame goes before the result, because the result is what makes the
        # SDK flush the batch to your store
        await self.mirror()
        await self.send(result)
        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 mirror(self):
        # one frame per turn, holding the lines Claude Code wrote to the file
        if self.options.session_store is None:
            return

        entries = [{"type": "user", "uuid": mark, "sessionId": self.session_id,
                    "message": {"role": "user", "content": text}}
                   for mark, text in zip(self.marks, self.said)]
        await self.send({"type": "transcript_mirror", "entries": entries,
                         "filePath": self.transcript()})

    def transcript(self):
        # the SDK reads the store key off this path, so it has to look like the
        # file Claude Code would have written
        config = (self.options.env.get("CLAUDE_CONFIG_DIR")
                  or os.environ.get("CLAUDE_CONFIG_DIR") or f"{Path.home()}/.claude")
        project = project_key_for_directory(self.options.cwd)
        return f"{config}/projects/{project}/{self.session_id}.jsonl"

    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:
            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"]:
                name = f"mcp__{server}__{tool['name']}"

                # a tool left out here never reaches the model at all
                if name not in self.options.disallowed_tools:
                    tools[name] = (server, tool["name"])
        return tools

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

        allowed, args = await self.permit(name, args, call)
        text, failed = (await self.run(name, args, call)) if allowed else (args, True)

        # a denied or raising call gets its own event, so a failure is never silent
        if failed:
            await self.hooks("PostToolUseFailure", {"tool_name": name, "error": text,
                                                    "tool_input": args, "tool_use_id": call})

        result = {"type": "tool_result", "tool_use_id": call, "content": text,
                  "is_error": failed}
        await self.send(self.message("user", [result]))
        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})

        # the hooks see the answer before the model does, with the same call id
        text = result["content"][0]["text"]
        for output in await self.hooks("PostToolUse", {"tool_name": name, "tool_input": args,
                                                       "tool_response": text, "tool_use_id": call}):
            # updatedToolOutput replaces what the model is about to read
            decision = output.get("hookSpecificOutput", {})
            text = decision.get("updatedToolOutput", text)
        return text, result.get("isError", False)

    def hook_input(self, event, data):
        # the fields a hook input carries whatever the event is, plus the
        # event's own fields on top
        return {"hook_event_name": event, "session_id": self.session_id,
                "cwd": ".", "transcript_path": "", **data}

    async def hooks(self, event, data):
        outputs = []
        for group in (self.setup.get("hooks") or {}).get(event, []):
            # a matcher of None means the group runs for every tool
            # an event with no tool name tests the matcher against an empty string
            if group["matcher"] and not re.fullmatch(group["matcher"],
                                                     data.get("tool_name", "")):
                continue

            # the SDK sent ids, not functions, so the call goes back by id
            for callback in group["hookCallbackIds"]:
                await self.lifecycle(event, "hook_started")
                answer = await self.ask({"subtype": "hook_callback",
                                         "callback_id": callback,
                                         "tool_use_id": data.get("tool_use_id"),
                                         "input": self.hook_input(event, data)})
                outputs.append(answer["response"])
                await self.lifecycle(event, "hook_response")
        return outputs

    async def lifecycle(self, event, phase):
        # include_hook_events asks for these two frames around every callback
        if self.options.include_hook_events:
            await self.send({"type": "system", "subtype": phase,
                             "hook_event": event, "session_id": self.session_id,
                             "uuid": f"ev_{next(self.ids)}"})

    async def pre_hooks(self, name, args, call):
        verdict = None
        for output in await self.hooks("PreToolUse", {"tool_name": name,
                                                      "tool_input": args, "tool_use_id": call}):
            decision = output.get("hookSpecificOutput", {})

            # one deny ends it, and the other hooks' opinions no longer matter
            if decision.get("permissionDecision") == "deny":
                return "deny", decision.get("permissionDecisionReason", "Blocked by a hook")
            verdict = decision.get("permissionDecision", verdict)
            args = decision.get("updatedInput", args)
        return verdict, args

    async def permit(self, name, args, call):
        # hooks are the first step, so nothing below here can skip them
        verdict, args = await self.pre_hooks(name, args, call)
        if verdict in ("allow", "deny"):
            return verdict == "allow", args

        o = self.options
        if o.permission_mode == "bypassPermissions" or name in o.allowed_tools:
            return True, args

        # nothing approved the call and there is nobody left to ask
        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):
        # the SDK runs your callback and writes its answer back on this channel
        decision = (await self.ask({"subtype": "can_use_tool", "tool_name": name,
                                    "input": args, "tool_use_id": call,
                                    "permission_suggestions": None,
                                    "blocked_path": None}))["response"]

        # an allow may hand back different arguments to run with
        if decision["behavior"] == "allow":
            return True, decision["updatedInput"]
        return False, decision["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}]}

REFUNDED = []


@tool("refund_order", "Refund an order by its id. The money goes back to the customer.",
      {"order_id": str})
async def refund_order(args):
    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])