"""ProgressTracker — yields ChatKit `WorkflowItem` events for one respond() turn. The bridge does NOT use `chatkit.agents.AgentContext` (it proxies /v1/chat/completions, not the Agents SDK). To get the same on-the-wire event shape that AgentContext produces, we replicate `start_workflow` / `add_workflow_task` / `update_workflow_task` / `end_workflow` by hand using the same documented primitives: ThreadItemAddedEvent(WorkflowItem) ThreadItemUpdatedEvent(update=WorkflowTaskAdded(...)) ThreadItemUpdatedEvent(update=WorkflowTaskUpdated(...)) ThreadItemDoneEvent(WorkflowItem) Design notes: - Lazy open: a `WorkflowItem` is only emitted once `add_task` is called for the first time. Turns that produce zero progress signals emit no workflow item (AC-13). - Pure event generator: no I/O, no asyncio.create_task, no SSE connection. Caller drains the async generators directly onto the SSE stream. Keeps the cancellation pattern in `respond()` (commit 1dde7f57) intact. - Disabled path: when `enabled=False`, all generators yield nothing. The bridge then falls through to its existing `ProgressUpdateEvent` emit sites unchanged. - In-place mutation: heartbeats call `update_task(idx, ...)` and mutate the same task — never `add_task` per tick. This is the duplication fix (the 2026-05-13 regression that motivated this spec). - Sub-1s coalescing: rapid duplicate `update_task` calls on the same task are dropped (last-wins) when neither status nor title changed materially, but terminal status changes (`complete`) are always emitted. """ from __future__ import annotations import time import uuid from collections.abc import AsyncIterator from datetime import datetime from typing import Literal from chatkit.types import ( CustomSummary, CustomTask, DurationSummary, ThreadItemAddedEvent, ThreadItemDoneEvent, ThreadItemUpdatedEvent, ThreadStreamEvent, Workflow, WorkflowItem, WorkflowTaskAdded, WorkflowTaskUpdated, ) _COALESCE_WINDOW_S = 1.0 class ProgressTracker: """Build a single WorkflowItem for one respond() turn. Usage:: tracker = ProgressTracker(thread_id=thread.id, enabled=True) async for ev in tracker.add_task(title="Thinking", icon="sparkle"): yield ev thinking_idx = tracker.last_index # ... later, in a heartbeat: async for ev in tracker.update_task(thinking_idx, title="Still working"): yield ev # ... at end of turn: async for ev in tracker.end_workflow(): yield ev """ def __init__(self, thread_id: str, enabled: bool) -> None: self._thread_id = thread_id self._enabled = enabled self._item: WorkflowItem | None = None self._opened = False self._closed = False self.last_index: int = -1 self._last_emit_at: dict[int, float] = {} @property def enabled(self) -> bool: return self._enabled @property def opened(self) -> bool: return self._opened @property def item(self) -> WorkflowItem | None: return self._item def _ensure_item(self) -> WorkflowItem: if self._item is None: self._item = WorkflowItem( id=f"workflow_{uuid.uuid4().hex}", thread_id=self._thread_id, created_at=datetime.now(), workflow=Workflow( type="custom", tasks=[], summary=CustomSummary(title="Working…", icon="sparkle"), expanded=False, ), ) return self._item async def add_task( self, *, title: str, icon: str | None = None, ) -> AsyncIterator[ThreadStreamEvent]: """Append a CustomTask and yield the matching add/update event. The first `add_task` of the turn yields a `ThreadItemAddedEvent` carrying the freshly-built `WorkflowItem` (with the new task already inside). Subsequent calls yield `ThreadItemUpdatedEvent` with a `WorkflowTaskAdded` update. After iteration completes, `self.last_index` is set to the new task's index. """ if not self._enabled or self._closed: return item = self._ensure_item() task = CustomTask( title=title, icon=icon, status_indicator="loading", ) item.workflow.tasks.append(task) idx = len(item.workflow.tasks) - 1 self.last_index = idx # Note: do NOT seed _last_emit_at here. The first update_task after # an add_task carries genuine new info (e.g. heartbeat elapsed) and # must always emit — coalescing applies only between successive # update_task calls. if not self._opened: self._opened = True yield ThreadItemAddedEvent(item=item) else: yield ThreadItemUpdatedEvent( item_id=item.id, update=WorkflowTaskAdded(task=task, task_index=idx), ) async def update_task( self, task_index: int, *, title: str | None = None, status: Literal["loading", "complete"] | None = None, icon: str | None = None, ) -> AsyncIterator[ThreadStreamEvent]: """Mutate an existing task in place and yield a WorkflowTaskUpdated. Sub-1s repeats with no terminal status change are dropped (last-wins). When `status="complete"`, the event is always emitted regardless of cadence. """ if not self._enabled or self._closed or self._item is None: return if task_index < 0 or task_index >= len(self._item.workflow.tasks): return task = self._item.workflow.tasks[task_index] if title is not None: task.title = title if icon is not None: task.icon = icon status_changed_to_complete = False if status is not None: if status == "complete" and task.status_indicator != "complete": status_changed_to_complete = True task.status_indicator = status now = time.monotonic() last = self._last_emit_at.get(task_index, 0.0) if not status_changed_to_complete and (now - last) < _COALESCE_WINDOW_S: # Coalesce: state was mutated above (last-wins) but no event emit. return self._last_emit_at[task_index] = now yield ThreadItemUpdatedEvent( item_id=self._item.id, update=WorkflowTaskUpdated(task=task, task_index=task_index), ) async def end_workflow( self, ) -> AsyncIterator[ThreadStreamEvent]: """Finalize the workflow with a ThreadItemDoneEvent. No-ops if disabled, never opened, or already closed. Flips any leftover loading tasks to complete (best effort) and attaches a `DurationSummary` when no summary is set. """ if not self._enabled or not self._opened or self._closed or self._item is None: return self._closed = True for t in self._item.workflow.tasks: if t.status_indicator == "loading": t.status_indicator = "complete" if self._item.workflow.summary is None: elapsed = (datetime.now() - self._item.created_at).total_seconds() self._item.workflow.summary = DurationSummary(duration=max(int(elapsed), 0)) self._item.workflow.expanded = False yield ThreadItemDoneEvent(item=self._item)