Todo tracking

import asyncio
import itertools
import json
import re

from claude_agent_sdk import ClaudeAgentOptions, Transport

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):
        # 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:]
        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.made, self.created = 0, None
        self.ready = False

    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"])
            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 work(self):
        while (prompt := await self.prompts.get()) is not None:
            await self.turn(prompt)

    async def turn(self, prompt):
        self.said.append(prompt)
        await self.send({"type": "system", "subtype": "init", "model": "shop-desk",
                         "session_id": self.session_id, "tools": []})

        # the loop: a step either asks for tools or has the text to answer with
        results, turns = [], 0
        while "calls" in (step := self.model.reply(self.said, results)):
            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 use(self, name, args):
        call = f"toolu_{next(self.ids)}"
        await self.send(self.message("assistant", [{"type": "tool_use", "id": call,
                                                    "name": name, "input": args}]))

        # a tool nobody allowed never runs, and the refusal goes back as its result
        allowed = name in self.options.allowed_tools
        text = self.task(name, args) if allowed else f"{name} was not allowed."
        await self.send(self.message("user", [{"type": "tool_result", "tool_use_id": call,
                                               "content": text, "is_error": not allowed}]))
        return text

    def task(self, name, args):
        # Claude Code is the side that assigns a task id, so the stand-in does too
        if name == "TaskCreate":
            self.made += 1
            self.created = {"task": {"id": f"t{self.made}", "subject": args["subject"]}}
        return "ok"

    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)}")
        line = {"type": kind, "message": body, "parent_tool_use_id": parent,
                "session_id": self.session_id}

        # Claude Code hangs a tool's structured output on the user line beside its result
        if kind == "user" and self.created is not None:
            line["tool_use_result"], self.created = self.created, None
        return line