diffusion-review 0.1.0__py3-none-any.whl

diffusion/__init__.py

@@ -0,0 +1 @@
+"""Local engine that turns a GitHub PR into a guided, linked review walkthrough."""

diffusion/agent.py

@@ -0,0 +1,134 @@
+"""Backend-neutral agent interface: requests, events, custom tools, and the runner seam.
+
+A *runner* takes an ``AgentRequest`` and yields ``Event``s. Backends (``backends/``) implement it on top of a
+model API (Pydantic AI) or an agent CLI (Claude Code, Codex). Everything else in the engine only sees this.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable, Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Protocol
+
+from .config import Config
+
+
+def inline_refs(schema: dict) -> dict:
+    """Inline ``$defs`` so tool schemas are flat (some tool-callers handle ``$ref`` poorly)."""
+    defs = schema.get("$defs", {})
+
+    def walk(node: Any, is_props: bool = False) -> Any:
+        if isinstance(node, list):
+            return [walk(v) for v in node]
+        if not isinstance(node, dict):
+            return node
+        if is_props:  # {property name: schema}: keys are names (e.g. "title"), not keywords
+            return {k: walk(v) for k, v in node.items()}
+        if "$ref" in node:
+            return walk(defs[node["$ref"].rsplit("/", 1)[-1]])
+        return {k: walk(v, k == "properties") for k, v in node.items() if k not in ("$defs", "title")}
+
+    return walk(schema)
+
+
+# ---------- events surfaced to callers (CLI progress, SSE) ----------
+
+
+@dataclass
+class ToolEvent:
+    name: str
+    summary: str
+
+
+@dataclass
+class TextEvent:
+    text: str
+
+
+@dataclass
+class DoneEvent:
+    turns: int
+    cost_usd: float | None
+    is_error: bool
+    session_id: str | None
+    result: str | None = None
+
+
+Event = ToolEvent | TextEvent | DoneEvent
+
+
+def summarize_tool(name: str, args: dict) -> str:
+    short = name.rsplit("__", 1)[-1]
+    if short.lower() in ("grep", "glob"):
+        return f"{short.capitalize()} {args.get('pattern', '')}".strip()
+    if short.lower() in ("read", "read_file"):
+        return f"Read {Path(args.get('file_path') or args.get('path') or '').name}"
+    if short == "get_hunk":
+        ids = args.get("ids") or []
+        return f"Hunk {', '.join(ids) if isinstance(ids, list) else ids}"
+    if short == "show_code":
+        return f"Show {Path(str(args.get('file', ''))).name}:{args.get('start', '')}"
+    return short
+
+
+@dataclass
+class ToolSpec:
+    """A custom tool the engine exposes to the agent (e.g. ``get_hunk``, ``submit_guide``).
+
+    ``terminal`` tools end the task: backends that support structured final output (Codex) may use the tool's
+    schema as the output schema and call the handler with the final result instead of exposing it as a tool.
+    """
+
+    name: str
+    description: str
+    schema: dict
+    handler: Callable[[dict], Awaitable[str]]
+    terminal: bool = False
+
+
+@dataclass
+class AgentRequest:
+    system_prompt: str
+    prompt: str
+    model: str
+    cwd: Path
+    tools: list[ToolSpec] = field(default_factory=list)
+    max_turns: int = 40
+    resume: str | None = None
+    include_partial: bool = False
+
+
+class Runner(Protocol):
+    def __call__(self, req: AgentRequest) -> AsyncIterator[Event]: ...
+
+
+class SandboxError(ValueError):
+    pass
+
+
+def within(root: Path, path: str | Path) -> Path:
+    """Resolve ``path`` (relative to ``root`` unless absolute) and require it to stay inside ``root``."""
+    root = root.resolve()
+    p = Path(path)
+    p = (p if p.is_absolute() else root / p).resolve()
+    if p != root and root not in p.parents:
+        raise SandboxError(f"{path} is outside the repository")
+    return p
+
+
+def make_runner(cfg: Config) -> Runner:
+    """The runner for the configured backend."""
+    if cfg.backend == "api":
+        from .backends.api import api_runner
+
+        return api_runner(cfg)
+    if cfg.backend == "claude-code":
+        from .backends.claude_code import claude_code_runner
+
+        return claude_code_runner(cfg)
+    if cfg.backend == "codex":
+        from .backends.codex import codex_runner
+
+        return codex_runner(cfg)
+    raise ValueError(f"unknown backend {cfg.backend!r}")

diffusion/app.py

@@ -0,0 +1,333 @@
+"""HTTP API (localhost only). Guide generation streams progress over SSE."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import time
+import uuid
+from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from .agent import DoneEvent, Runner, ToolEvent, make_runner
+from .chat_agent import ChatDelta, ChatFocus, Question, run_turn
+from .config import Config, load_config
+from .github import GitHubError
+from .guide_agent import GuideDone, GuideError, generate_guide
+from .jobs import Emit, Jobs
+from .pipeline import Prepared, prepare, target_ref
+from .repo import RepoError
+from .store import Store
+from .symbols import NAME, SymbolIndex
+
+PrepareFn = Callable[[Config, str], Prepared]
+MAX_FILE_CHARS = 2_000_000
+
+
+@dataclass
+class Deps:
+    cfg: Config
+    runner: Runner
+    prepare: PrepareFn
+    store: Store
+    jobs: Jobs
+
+
+def sse(event: str, data: dict) -> str:
+    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
+
+
+def sse_response(events: AsyncIterator[tuple[str, dict]]) -> StreamingResponse:
+    async def body():
+        async for ev, data in events:
+            yield sse(ev, data)
+
+    return StreamingResponse(
+        body(),
+        media_type="text/event-stream",
+        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+    )
+
+
+class GuideReq(BaseModel):
+    url: str
+    force: bool = False
+
+
+class ChatCreate(BaseModel):
+    url: str
+
+
+class TurnReq(BaseModel):
+    url: str
+    questions: list[Question]
+
+
+def chat_summary(c: dict) -> dict:
+    qs = [q for t in c.get("turns", []) for q in t["questions"]]
+    return {
+        "id": c["id"],
+        "title": qs[0]["text"][:80] if qs else "",
+        "questions": len(qs),
+        "anchored": sum(1 for q in qs if q.get("anchor")),
+        "created_at": c["created_at"],
+        "updated_at": c["updated_at"],
+        "head_sha": c.get("head_sha"),
+    }
+
+
+def create_app(deps: Deps | None = None) -> FastAPI:
+    app = FastAPI(title="diffusion")
+    _deps: list[Deps] = [deps] if deps else []
+
+    def d() -> Deps:
+        if not _deps:
+            cfg = load_config()
+            _deps.append(Deps(cfg, make_runner(cfg), prepare, Store(cfg.cache_dir), Jobs()))
+        return _deps[0]
+
+    _prepared: dict[str, Prepared] = {}
+    chat_locks: dict[str, asyncio.Lock] = {}
+    symbols: dict[Path, SymbolIndex] = {}
+    symbol_locks: dict[Path, asyncio.Lock] = {}
+
+    async def get_prepared(url: str, refresh: bool = False) -> Prepared:
+        try:
+            key = target_ref(url).key
+        except GitHubError as e:
+            raise HTTPException(400, str(e)) from e
+        if refresh or key not in _prepared:
+            try:
+                _prepared[key] = await asyncio.to_thread(d().prepare, d().cfg, url)
+            except (GitHubError, RepoError) as e:
+                raise HTTPException(502, str(e)) from e
+            except FileNotFoundError as e:
+                raise HTTPException(400, str(e)) from e
+        return _prepared[key]
+
+    @app.get("/api/health")
+    async def health() -> dict:
+        cfg = d().cfg
+        return {
+            "ok": True,
+            "profile": cfg.profile,
+            "backend": cfg.backend,
+            "guide_model": cfg.guide_model,
+            "chat_model": cfg.chat_model,
+        }
+
+    @app.get("/api/pr")
+    async def pr(url: str) -> dict:
+        prep = await get_prepared(url, refresh=True)
+        cached = d().store.load_guide(prep.meta.ref.key, prep.meta.head_sha)
+        return {**prep.to_dict(), "guide": cached}
+
+    @app.post("/api/guide")
+    async def guide(req: GuideReq) -> StreamingResponse:
+        deps = d()
+        prep = await get_prepared(req.url)
+        key, sha = prep.meta.ref.key, prep.meta.head_sha
+        job_key = f"guide:{key}:{sha}"
+
+        if not req.force and (cached := deps.store.load_guide(key, sha)):
+
+            async def replay():
+                yield "guide", cached
+
+            return sse_response(replay())
+
+        if req.force:
+            deps.jobs.drop(job_key)
+        job = deps.jobs.start(job_key, partial(run_guide, deps, prep))
+        return sse_response(job.follow())
+
+    @app.get("/api/file")
+    async def file(url: str, path: str) -> dict:
+        """A changed file at the PR head, so the diff can show lines around and between hunks."""
+        prep = await get_prepared(url)
+        if path not in {f.path for f in prep.index.files}:
+            raise HTTPException(404, f"{path} isn't in this diff")
+        try:
+            text = await asyncio.to_thread(prep.checkout.file_at_head, path)
+        except UnicodeDecodeError as e:
+            raise HTTPException(415, f"{path} isn't text") from e
+        if text is None:
+            raise HTTPException(404, f"{path} doesn't exist at the head commit")
+        if len(text) > MAX_FILE_CHARS:
+            raise HTTPException(413, f"{path} is too large to show in full")
+        return {"path": path, "text": text}
+
+    @app.get("/api/definition")
+    async def definition(url: str, symbol: str, path: str) -> dict:
+        """Where `symbol` is defined at the PR head, nearest to `path` first, with each definition's source."""
+        if not NAME.match(symbol):
+            raise HTTPException(400, "Not a name")
+        prep = await get_prepared(url)
+        root = prep.checkout.worktree
+        async with symbol_locks.setdefault(root, asyncio.Lock()):
+            if root not in symbols:
+                symbols[root] = await asyncio.to_thread(SymbolIndex, root)
+        index = symbols[root]
+        found = index.lookup(symbol, path)
+        return {
+            "symbol": symbol,
+            "matches": [
+                {"path": m.path, "line": m.line, "end": m.end, "text": index.snippet(m)} for m in found
+            ],
+        }
+
+    # ---------- chats ----------
+
+    def pr_key(url: str) -> str:
+        try:
+            return target_ref(url).key
+        except GitHubError as e:
+            raise HTTPException(400, str(e)) from e
+
+    @app.get("/api/chats")
+    async def list_chats(url: str) -> list[dict]:
+        return [chat_summary(c) for c in d().store.list_chats(pr_key(url)) if c.get("turns")]
+
+    @app.post("/api/chats")
+    async def create_chat(req: ChatCreate) -> dict:
+        now = time.time()
+        chat = {
+            "id": uuid.uuid4().hex[:12],
+            "created_at": now,
+            "updated_at": now,
+            "session_id": None,
+            "turns": [],
+        }
+        d().store.save_chat(pr_key(req.url), chat)
+        return chat
+
+    @app.get("/api/chats/{chat_id}")
+    async def get_chat(chat_id: str, url: str) -> dict:
+        chat = d().store.load_chat(pr_key(url), chat_id)
+        if not chat:
+            raise HTTPException(404, "chat not found")
+        return chat
+
+    @app.post("/api/chats/{chat_id}/turn")
+    async def turn(chat_id: str, req: TurnReq) -> StreamingResponse:
+        deps = d()
+        key = pr_key(req.url)
+        chat = deps.store.load_chat(key, chat_id)
+        if not chat:
+            raise HTTPException(404, "chat not found")
+        if not req.questions:
+            raise HTTPException(400, "no questions")
+        if len({q.n for q in req.questions}) != len(req.questions):
+            raise HTTPException(400, "duplicate question numbers")
+        prep = await get_prepared(req.url)
+        guide = deps.store.load_guide(key, prep.meta.head_sha)
+        lock = chat_locks.setdefault(chat_id, asyncio.Lock())
+
+        async def events():
+            # One turn per chat at a time: a turn resumes the session the previous one produced.
+            async with lock:
+                chat = deps.store.load_chat(key, chat_id)  # fresh: an earlier turn may have just saved
+                if chat is None:  # deleted while this turn waited for the lock
+                    yield "error", {"message": "chat not found"}
+                    return
+                record = {
+                    "at": time.time(),
+                    "head_sha": prep.meta.head_sha,
+                    "questions": [q.model_dump() for q in req.questions],
+                    "answers": {str(q.n): "" for q in req.questions},
+                    "tools": [],
+                    "shows": [],
+                    "status": "streaming",
+                }
+                chat["turns"].append(record)
+                chat["head_sha"] = prep.meta.head_sha
+                chat["updated_at"] = record["at"]
+                deps.store.save_chat(key, chat)
+                try:
+                    async for ev in run_turn(
+                        prep, guide, req.questions, deps.runner, deps.cfg.chat_model, chat.get("session_id")
+                    ):
+                        if isinstance(ev, ChatDelta):
+                            record["answers"][str(ev.n)] += ev.text
+                            yield "delta", {"n": ev.n, "text": ev.text}
+                        elif isinstance(ev, ChatFocus):
+                            record["shows"].append(ev.to_dict())
+                            yield "focus", ev.to_dict()
+                        elif isinstance(ev, ToolEvent):
+                            record["tools"].append(ev.summary)
+                            yield "tool", {"name": ev.name, "summary": ev.summary}
+                        elif isinstance(ev, DoneEvent):
+                            if ev.is_error:
+                                raise RuntimeError(ev.result or "agent error")
+                            chat["session_id"] = ev.session_id or chat.get("session_id")
+                    record["status"] = "done"
+                    yield "done", {"ns": [q.n for q in req.questions]}
+                except Exception as e:  # surface to the client; keep partial answers
+                    record["status"] = "error"
+                    record["error"] = str(e) or type(e).__name__
+                    yield "error", {"message": record["error"]}
+                finally:
+                    chat["updated_at"] = time.time()
+                    deps.store.save_chat(key, chat)
+
+        return sse_response(events())
+
+    dist = _web_dist()
+    if dist:
+        app.mount("/assets", StaticFiles(directory=dist / "assets"), name="assets")
+
+        @app.get("/{path:path}", include_in_schema=False)
+        async def spa(path: str):
+            f = dist / path
+            if path and f.is_file():
+                return FileResponse(f)
+            return FileResponse(dist / "index.html")
+
+    return app
+
+
+def _web_dist() -> Path | None:
+    """Built UI: $DIFFUSION_WEB_DIST, the copy bundled in the package, or web/dist in a source checkout."""
+    env = os.environ.get("DIFFUSION_WEB_DIST")
+    here = Path(__file__).resolve()
+    candidates = [Path(env)] if env else []
+    candidates += [here.parent / "web", here.parents[3] / "web" / "dist"]
+    return next((d for d in candidates if (d / "index.html").is_file()), None)
+
+
+async def run_guide(deps: Deps, prep: Prepared, emit: Emit) -> None:
+    t0 = time.monotonic()
+    done: GuideDone | None = None
+    async for ev in generate_guide(
+        prep.meta, prep.checkout, prep.index, deps.runner, deps.cfg.guide_model, deps.cfg.max_turns
+    ):
+        if isinstance(ev, ToolEvent):
+            emit("tool", {"name": ev.name, "summary": ev.summary})
+        elif isinstance(ev, DoneEvent) and ev.is_error:
+            raise GuideError(ev.result or "agent error")
+        elif isinstance(ev, GuideDone):
+            done = ev
+    assert done is not None
+    record = {
+        **done.guide.model_dump(by_alias=True),
+        "meta": {
+            "head_sha": prep.meta.head_sha,
+            "backend": deps.cfg.backend,
+            "model": deps.cfg.guide_model,
+            "turns": done.turns,
+            "cost_usd": done.cost_usd,
+            "attempts": done.attempts,
+            "seconds": round(time.monotonic() - t0, 1),
+            "created_at": time.time(),
+        },
+    }
+    deps.store.save_guide(prep.meta.ref.key, prep.meta.head_sha, record)
+    emit("guide", record)

diffusion/backends/__init__.py

@@ -0,0 +1 @@
+"""Agent backends. Each exposes a ``*_runner(cfg) -> Runner`` factory; see ``agent.make_runner``."""

diffusion/backends/api.py

@@ -0,0 +1,206 @@
+"""Backend: a Pydantic AI agent over a model API (Anthropic, OpenAI, OpenRouter or any OpenAI-compatible server).
+
+Pydantic AI runs the loop and speaks each provider's wire format. We add the read-only repository tools, the
+engine's custom tools, Anthropic prompt caching, and the stop rule: the run ends when the model answers without
+calling tools, a terminal tool (e.g. ``submit_guide``) is accepted, or ``max_turns`` requests have been made.
+Conversations are stored as JSON so a later request can ``resume`` them (chat follow-ups, guide retries).
+"""
+
+from __future__ import annotations
+
+import uuid
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import pydantic_ai
+from pydantic_ai import (
+    Agent,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    ModelMessage,
+    ModelMessagesTypeAdapter,
+    ModelResponse,
+    PartDeltaEvent,
+    PartStartEvent,
+    TextPart,
+    TextPartDelta,
+    Tool,
+    ToolFailed,
+    ToolReturnPart,
+    UsageLimitExceeded,
+    UsageLimits,
+)
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+from ..agent import AgentRequest, DoneEvent, Event, Runner, TextEvent, ToolEvent, ToolSpec, summarize_tool
+from ..config import Config
+from .tools import repo_tools
+
+pydantic_ai.BANNER_ENABLED = False  # our CLI owns its output
+
+ACCEPTED = ("received",)  # terminal tool replies that end the task
+MAX_TOKENS = 16000
+
+
+class SessionStore:
+    def __init__(self, root: Path):
+        self.root = root
+
+    def load(self, sid: str) -> list[ModelMessage]:
+        try:
+            return ModelMessagesTypeAdapter.validate_json((self.root / f"{sid}.json").read_bytes())
+        except (OSError, ValueError):
+            return []
+
+    def save(self, sid: str, messages: list[ModelMessage]) -> None:
+        self.root.mkdir(parents=True, exist_ok=True)
+        tmp = self.root / f"{sid}.json.tmp"
+        tmp.write_bytes(ModelMessagesTypeAdapter.dump_json(messages))
+        tmp.replace(self.root / f"{sid}.json")
+
+
+def make_model(cfg: Config, name: str, http_client: Any = None) -> tuple[Model, ModelSettings]:
+    """The Pydantic AI model for a profile, and the settings every request uses."""
+    if cfg.protocol == "anthropic":
+        from anthropic import AsyncAnthropic
+        from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+        from pydantic_ai.providers.anthropic import AnthropicProvider
+
+        kw: dict[str, Any] = {"base_url": cfg.base_url, "default_headers": cfg.headers or None}
+        kw["auth_token" if cfg.auth == "bearer" else "api_key"] = cfg.api_key
+        client = AsyncAnthropic(**kw, max_retries=3, http_client=http_client)
+        # Cache the system prompt, the tools, and the conversation so far. Explicit breakpoints rather than
+        # automatic caching, which some Anthropic-compatible gateways don't support.
+        settings = AnthropicModelSettings(
+            max_tokens=MAX_TOKENS,
+            anthropic_cache_instructions=True,
+            anthropic_cache_tool_definitions=True,
+            anthropic_cache_messages=True,
+        )
+        return AnthropicModel(name, provider=AnthropicProvider(anthropic_client=client)), settings
+
+    from openai import AsyncOpenAI
+
+    client = AsyncOpenAI(
+        base_url=cfg.base_url,
+        api_key=cfg.api_key or "none",
+        default_headers=cfg.headers or None,
+        max_retries=3,
+        http_client=http_client,
+    )
+    settings = ModelSettings(max_tokens=MAX_TOKENS)
+    if urlparse(cfg.base_url or "").hostname == "openrouter.ai":
+        from pydantic_ai.models.openrouter import OpenRouterModel
+        from pydantic_ai.providers.openrouter import OpenRouterProvider
+
+        return OpenRouterModel(name, provider=OpenRouterProvider(openai_client=client)), settings
+
+    from pydantic_ai.models.openai import OpenAIChatModel
+    from pydantic_ai.providers.openai import OpenAIProvider
+
+    return OpenAIChatModel(name, provider=OpenAIProvider(openai_client=client)), settings
+
+
+def _tool(spec: ToolSpec) -> Tool:
+    async def call(**args: Any) -> str:
+        try:
+            return await spec.handler(args)
+        except Exception as e:  # report to the model and let it recover
+            raise ToolFailed(f"error: {e}") from e
+
+    return Tool.from_schema(call, spec.name, spec.description, spec.schema)
+
+
+def _cost(responses: list[ModelResponse]) -> float | None:
+    """OpenRouter reports the exact cost; otherwise price the usage with genai-prices when the model is known."""
+    total = None
+    for r in responses:
+        cost = (r.provider_details or {}).get("cost")
+        if not isinstance(cost, int | float):
+            try:
+                cost = float(r.cost().total_price)
+            except Exception:  # unknown model or gateway
+                continue
+        total = (total or 0.0) + cost
+    return total
+
+
+def api_runner(
+    cfg: Config, model: Model | None = None, sessions: SessionStore | None = None, http_client: Any = None
+) -> Runner:
+    store = sessions or SessionStore(cfg.cache_dir / "sessions")
+
+    async def run(req: AgentRequest) -> AsyncIterator[Event]:
+        specs: list[ToolSpec] = repo_tools(req.cwd) + list(req.tools)
+        terminal = {t.name for t in specs if t.terminal}
+        llm, settings = (model, None) if model else make_model(cfg, req.model, http_client)
+        agent = Agent(llm, instructions=req.system_prompt, tools=[_tool(t) for t in specs], retries=3)
+        sid = req.resume or uuid.uuid4().hex[:16]
+        history = store.load(req.resume) if req.resume else []
+        messages = list(history)
+        turns = 0
+        error: str | None = None
+        final_text = ""
+
+        try:
+            async with agent.iter(
+                req.prompt,
+                message_history=history,
+                model_settings=settings,
+                usage_limits=UsageLimits(request_limit=req.max_turns),
+            ) as agent_run:
+                finished = False
+                async for node in agent_run:
+                    if finished:  # a terminal tool was accepted: keep its results so the session can resume
+                        if Agent.is_model_request_node(node):
+                            messages = [*messages, node.request]
+                        break
+                    if Agent.is_model_request_node(node):
+                        async with node.stream(agent_run.ctx) as stream:
+                            async for ev in stream:
+                                if not req.include_partial:
+                                    continue
+                                if isinstance(ev, PartStartEvent) and isinstance(ev.part, TextPart):
+                                    text = ev.part.content
+                                elif isinstance(ev, PartDeltaEvent) and isinstance(ev.delta, TextPartDelta):
+                                    text = ev.delta.content_delta
+                                else:
+                                    continue
+                                if text:
+                                    yield TextEvent(text)
+                        turns += 1
+                        response = stream.response
+                        final_text = "".join(p.content for p in response.parts if isinstance(p, TextPart))
+                        if final_text and not req.include_partial:
+                            yield TextEvent(final_text)
+                    elif Agent.is_call_tools_node(node):
+                        async with node.stream(agent_run.ctx) as stream:
+                            async for ev in stream:
+                                if isinstance(ev, FunctionToolCallEvent):
+                                    yield ToolEvent(
+                                        ev.part.tool_name,
+                                        summarize_tool(ev.part.tool_name, ev.part.args_as_dict()),
+                                    )
+                                elif isinstance(ev, FunctionToolResultEvent):
+                                    r = ev.part
+                                    finished |= (
+                                        isinstance(r, ToolReturnPart)
+                                        and r.outcome == "success"
+                                        and r.tool_name in terminal
+                                        and str(r.content).startswith(ACCEPTED)
+                                    )
+                    messages = agent_run.all_messages()
+        except UsageLimitExceeded:
+            error = f"stopped after {req.max_turns} turns"
+        except Exception as e:
+            error = f"{type(e).__name__}: {e}"
+        finally:
+            store.save(sid, messages)
+
+        new = [m for m in messages[len(history) :] if isinstance(m, ModelResponse)]
+        yield DoneEvent(turns, _cost(new), error is not None, sid, error or final_text)
+
+    return run