app/progress_workflow.py

"""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)