feat(progress): ProgressBus with sessions, ticks, and turn-complete events

server/progress.py

@@ -0,0 +1,174 @@
+"""Progress event bus for in-flight generations.
+
+Endpoints (`/api/generate`, `/api/generate/dialog`) wrap their work in
+`bus.session(...)` which emits `start` and `done`/`error` events plus a
+0.5s `tick` heartbeat. Dialog mode also emits `turn_complete` between
+adapter calls. Subscribers receive events via `subscribe()` (used by
+the SSE endpoint).
+"""
+from __future__ import annotations
+
+import asyncio
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import AsyncIterator, Literal
+
+
+EventType = Literal["start", "tick", "turn_complete", "done", "error"]
+
+
+@dataclass
+class ProgressEvent:
+    type: EventType
+    elapsed_s: float
+    payload: dict = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        return {"type": self.type, "elapsed_s": round(self.elapsed_s, 2), **self.payload}
+
+
+class ProgressBus:
+    def __init__(self) -> None:
+        self._subscribers: list[asyncio.Queue[ProgressEvent]] = []
+        self._lock = asyncio.Lock()
+        self._current_session: "_Session | None" = None
+
+    async def publish(self, event: ProgressEvent) -> None:
+        async with self._lock:
+            subs = list(self._subscribers)
+        for q in subs:
+            await q.put(event)
+
+    @asynccontextmanager
+    async def subscribe(self) -> AsyncIterator[asyncio.Queue[ProgressEvent]]:
+        q: asyncio.Queue[ProgressEvent] = asyncio.Queue()
+        async with self._lock:
+            self._subscribers.append(q)
+            if self._current_session is not None:
+                snapshot = self._current_session.snapshot_event()
+                if snapshot is not None:
+                    await q.put(snapshot)
+        try:
+            yield q
+        finally:
+            async with self._lock:
+                if q in self._subscribers:
+                    self._subscribers.remove(q)
+
+    @asynccontextmanager
+    async def session(
+        self, kind: Literal["single", "dialog"], total_turns: int = 1,
+    ) -> AsyncIterator["_Session"]:
+        session = _Session(bus=self, kind=kind, total_turns=total_turns)
+        async with self._lock:
+            self._current_session = session
+        await self.publish(
+            ProgressEvent(
+                type="start",
+                elapsed_s=0.0,
+                payload={"kind": kind, "total_turns": total_turns, "turn": 0},
+            ),
+        )
+        ticker = asyncio.create_task(session._tick_loop())
+        try:
+            yield session
+            await self.publish(
+                ProgressEvent(
+                    type="done",
+                    elapsed_s=session.elapsed(),
+                    payload={
+                        "kind": kind,
+                        "seed_used": session.seed_used,
+                        "turn": session.turn,
+                        "total_turns": total_turns,
+                    },
+                ),
+            )
+        except Exception as exc:
+            await self.publish(
+                ProgressEvent(
+                    type="error",
+                    elapsed_s=session.elapsed(),
+                    payload={"message": str(exc)},
+                ),
+            )
+            raise
+        finally:
+            ticker.cancel()
+            try:
+                await ticker
+            except asyncio.CancelledError:
+                pass
+            async with self._lock:
+                if self._current_session is session:
+                    self._current_session = None
+
+
+@dataclass
+class _Session:
+    bus: ProgressBus
+    kind: Literal["single", "dialog"]
+    total_turns: int
+    started_at: float = field(default_factory=time.monotonic)
+    turn: int = 0
+    seed_used: int | None = None
+
+    def elapsed(self) -> float:
+        return time.monotonic() - self.started_at
+
+    def set_seed(self, seed: int) -> None:
+        self.seed_used = seed
+
+    async def turn_complete(self, turn_index: int) -> None:
+        self.turn = turn_index
+        await self.bus.publish(
+            ProgressEvent(
+                type="turn_complete",
+                elapsed_s=self.elapsed(),
+                payload={
+                    "turn": turn_index,
+                    "total_turns": self.total_turns,
+                    "kind": self.kind,
+                },
+            ),
+        )
+
+    async def _tick_loop(self) -> None:
+        try:
+            while True:
+                await asyncio.sleep(0.5)
+                await self.bus.publish(
+                    ProgressEvent(
+                        type="tick",
+                        elapsed_s=self.elapsed(),
+                        payload={
+                            "kind": self.kind,
+                            "turn": self.turn,
+                            "total_turns": self.total_turns,
+                        },
+                    ),
+                )
+        except asyncio.CancelledError:
+            pass
+
+    def snapshot_event(self) -> ProgressEvent | None:
+        return ProgressEvent(
+            type="tick",
+            elapsed_s=self.elapsed(),
+            payload={
+                "kind": self.kind,
+                "turn": self.turn,
+                "total_turns": self.total_turns,
+            },
+        )
+
+
+_BUS: ProgressBus | None = None
+
+
+def get_bus() -> ProgressBus:
+    global _BUS
+    if _BUS is None:
+        _BUS = ProgressBus()
+    return _BUS

tests/conftest.py

@@ -59,3 +59,14 @@ async def lifespan_ctx(app):
     """Run an ASGI app's lifespan startup/shutdown around an `httpx.AsyncClient`."""
     async with app.router.lifespan_context(app):
         yield
+
+
+@pytest.fixture(autouse=False)
+def reset_progress_bus():
+    """Reset server.progress._BUS so each test gets a fresh bus."""
+    import server.progress as p
+    p._BUS = None
+    try:
+        yield
+    finally:
+        p._BUS = None

tests/test_progress.py

@@ -0,0 +1,121 @@
+import asyncio
+
+import pytest
+
+from server.progress import ProgressEvent, get_bus
+
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_subscribe_receives_published_events(reset_progress_bus):
+    bus = get_bus()
+    async with bus.subscribe() as q:
+        await bus.publish(ProgressEvent(type="tick", elapsed_s=0.1, payload={"foo": 1}))
+        evt = await asyncio.wait_for(q.get(), 0.5)
+    assert evt.type == "tick"
+    assert evt.payload == {"foo": 1}
+
+
+async def test_two_subscribers_both_receive_events(reset_progress_bus):
+    bus = get_bus()
+    async with bus.subscribe() as q1, bus.subscribe() as q2:
+        await bus.publish(ProgressEvent(type="tick", elapsed_s=0.0))
+        a = await asyncio.wait_for(q1.get(), 0.5)
+        b = await asyncio.wait_for(q2.get(), 0.5)
+    assert a.type == "tick"
+    assert b.type == "tick"
+
+
+async def test_session_emits_start_and_done(reset_progress_bus):
+    bus = get_bus()
+    received: list[ProgressEvent] = []
+
+    async def collect():
+        async with bus.subscribe() as q:
+            while True:
+                received.append(await q.get())
+                if received[-1].type == "done":
+                    return
+
+    consumer = asyncio.create_task(collect())
+    await asyncio.sleep(0)  # let subscriber register
+
+    async with bus.session("single", total_turns=1) as sess:
+        sess.set_seed(42)
+
+    await asyncio.wait_for(consumer, 1.0)
+    types = [e.type for e in received]
+    assert types[0] == "start"
+    assert types[-1] == "done"
+    done_payload = received[-1].payload
+    assert done_payload["seed_used"] == 42
+
+
+async def test_session_emits_error_on_exception_and_reraises(reset_progress_bus):
+    bus = get_bus()
+    received: list[ProgressEvent] = []
+
+    async def collect():
+        async with bus.subscribe() as q:
+            while True:
+                received.append(await q.get())
+                if received[-1].type in ("done", "error"):
+                    return
+
+    consumer = asyncio.create_task(collect())
+    await asyncio.sleep(0)
+
+    with pytest.raises(RuntimeError):
+        async with bus.session("single", total_turns=1):
+            raise RuntimeError("boom")
+
+    await asyncio.wait_for(consumer, 1.0)
+    types = [e.type for e in received]
+    assert "error" in types
+    assert any(e.payload.get("message") == "boom" for e in received)
+
+
+async def test_turn_complete_event_carries_turn_payload(reset_progress_bus):
+    bus = get_bus()
+    received: list[ProgressEvent] = []
+
+    async def collect():
+        async with bus.subscribe() as q:
+            while True:
+                received.append(await q.get())
+                if received[-1].type == "done":
+                    return
+
+    consumer = asyncio.create_task(collect())
+    await asyncio.sleep(0)
+
+    async with bus.session("dialog", total_turns=3) as sess:
+        await sess.turn_complete(1)
+        await sess.turn_complete(2)
+        await sess.turn_complete(3)
+
+    await asyncio.wait_for(consumer, 1.0)
+    turn_events = [e for e in received if e.type == "turn_complete"]
+    assert [e.payload["turn"] for e in turn_events] == [1, 2, 3]
+    assert all(e.payload["total_turns"] == 3 for e in turn_events)
+
+
+async def test_late_subscriber_gets_snapshot(reset_progress_bus):
+    bus = get_bus()
+    received: list[ProgressEvent] = []
+
+    async def collect():
+        async with bus.subscribe() as q:
+            received.append(await asyncio.wait_for(q.get(), 1.0))
+
+    async with bus.session("dialog", total_turns=4) as sess:
+        await sess.turn_complete(2)
+        # join AFTER the session started
+        consumer = asyncio.create_task(collect())
+        await asyncio.wait_for(consumer, 1.0)
+
+    assert received[0].type == "tick"
+    assert received[0].payload["kind"] == "dialog"
+    assert received[0].payload["turn"] == 2
+    assert received[0].payload["total_turns"] == 4