tests/test_realtime_progress_feed.py

"""Integration tests for RealtimeProgressFeed.

Spec: `specs/realtime-progress-via-gateway-ws.md` §8.4.

Wires a real GatewayWsClient against the FakeGateway (reused from
test_gateway_ws_client.py), a real translator, and a real ProgressTracker.
Asserts the resulting ChatKit event stream matches the expected workflow
timeline for a multi-tool Hub→subagent run.
"""
from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any

import pytest

from chatkit.types import (
    CustomTask,
    ThreadItemAddedEvent,
    ThreadItemUpdatedEvent,
    WorkflowItem,
    WorkflowTaskAdded,
    WorkflowTaskUpdated,
)

from app.gateway_ws_client import GatewayWsClient
from app.progress_workflow import ProgressTracker
from app.realtime_progress_feed import RealtimeProgressFeed

from tests.test_gateway_ws_client import FakeGateway


pytestmark = pytest.mark.asyncio


@asynccontextmanager
async def _feed(
    gw: FakeGateway, parent_key: str = "agent:main:abc"
) -> AsyncIterator[tuple[GatewayWsClient, ProgressTracker, RealtimeProgressFeed]]:
    # These tests exercise the WS path explicitly — the fake gateway is a
    # WebSocket server. Pass mode="ws" so the feed uses the supplied
    # ws_client instead of the new default `poll` (HTTP sessions_history).
    GatewayWsClient.reset_process_instance()
    c = GatewayWsClient(ws_url=gw.ws_url, token="test-token")
    tracker = ProgressTracker(thread_id="thread-1", enabled=True)
    feed = RealtimeProgressFeed(
        parent_session_key=parent_key,
        tracker=tracker,
        mode="ws",
        ws_client=c,
    )
    try:
        yield c, tracker, feed
    finally:
        await feed.stop()
        await c.aclose()


async def _drain(feed: RealtimeProgressFeed) -> list[Any]:
    out = []
    async for ev in feed.drain():
        out.append(ev)
    return out


async def _wait_for_queue(
    feed: RealtimeProgressFeed, n: int, timeout: float = 2.0
) -> None:
    """Spin until the feed's internal queue has at least `n` events ready,
    or timeout. Spec: feed runs fully async, so we need a settle period
    after broadcasting fake frames.
    """
    deadline = asyncio.get_running_loop().time() + timeout
    while asyncio.get_running_loop().time() < deadline:
        if feed._event_queue.qsize() >= n:
            return
        await asyncio.sleep(0.01)


# --------------------------------------------------------------------------
# Tests
# --------------------------------------------------------------------------


async def test_tool_use_final_emits_workflow_added_then_task_added() -> None:
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="parent:1") as (_c, _t, feed):
            await feed.start()
            await gw.broadcast_chat(
                "parent:1",
                state="final",
                content=[
                    {
                        "type": "tool_use",
                        "id": "tu-a",
                        "name": "Read",
                        "input": {"file_path": "/x/y.py"},
                    }
                ],
            )
            await _wait_for_queue(feed, 1)
            events = await _drain(feed)

    # Lazy open: first AddTask should emit the WorkflowItem ADDED.
    assert events, "no events drained"
    assert isinstance(events[0], ThreadItemAddedEvent)
    assert isinstance(events[0].item, WorkflowItem)
    workflow_id = events[0].item.id
    # Subsequent task adds (none here — only one task) would be Updated.
    assert len(events[0].item.workflow.tasks) == 1
    task = events[0].item.workflow.tasks[0]
    assert isinstance(task, CustomTask)
    assert task.title == "Reading y.py"
    assert task.icon == "bolt"
    assert task.status_indicator == "loading"
    # workflow_id is stable for subsequent updates
    assert workflow_id.startswith("workflow_")


async def test_tool_use_then_tool_result_marks_task_complete() -> None:
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="parent:1") as (_c, _t, feed):
            await feed.start()
            await gw.broadcast_chat(
                "parent:1",
                state="final",
                run_id="r",
                seq=1,
                content=[
                    {
                        "type": "tool_use",
                        "id": "tu-1",
                        "name": "Bash",
                        "input": {"command": "ls -la"},
                    }
                ],
            )
            await _wait_for_queue(feed, 1)
            await gw.broadcast_chat(
                "parent:1",
                state="final",
                run_id="r",
                seq=2,
                content=[
                    {"type": "tool_result", "tool_use_id": "tu-1", "content": "ok"}
                ],
            )
            await _wait_for_queue(feed, 2)
            events = await _drain(feed)

    # Sequence: ThreadItemAddedEvent (workflow with one loading task)
    # then ThreadItemUpdatedEvent with WorkflowTaskUpdated marking complete.
    assert len(events) >= 2
    assert isinstance(events[0], ThreadItemAddedEvent)
    update_events = [e for e in events if isinstance(e, ThreadItemUpdatedEvent)]
    assert update_events, "expected an update event"
    # Look for the complete update.
    completes = [
        e
        for e in update_events
        if isinstance(e.update, WorkflowTaskUpdated)
        and e.update.task.status_indicator == "complete"
    ]
    assert completes, "no complete-status update emitted"


async def test_multi_tool_run_produces_ordered_task_timeline() -> None:
    """A realistic Hub→tool sequence: Read → Bash → Grep, all completed."""
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="hub") as (_c, _t, feed):
            await feed.start()
            # Three sequential tool_use finals.
            for seq, (tid, name, inp) in enumerate(
                [
                    ("a", "Read", {"file_path": "x.py"}),
                    ("b", "Bash", {"command": "ls"}),
                    ("c", "Grep", {"pattern": "foo"}),
                ],
                start=1,
            ):
                await gw.broadcast_chat(
                    "hub",
                    state="final",
                    run_id="r",
                    seq=seq,
                    content=[
                        {"type": "tool_use", "id": tid, "name": name, "input": inp}
                    ],
                )
            await _wait_for_queue(feed, 3)
            # Then their tool_results in order.
            for seq, tid in enumerate([("a"), ("b"), ("c")], start=4):
                await gw.broadcast_chat(
                    "hub",
                    state="final",
                    run_id="r",
                    seq=seq,
                    content=[{"type": "tool_result", "tool_use_id": tid}],
                )
            await _wait_for_queue(feed, 6)
            events = await _drain(feed)

    # First event is ThreadItemAdded (workflow + first task).
    assert isinstance(events[0], ThreadItemAddedEvent)
    wf = events[0].item
    # After all events applied, the workflow should have 3 tasks all complete.
    tasks = wf.workflow.tasks
    assert len(tasks) == 3
    for t in tasks:
        assert isinstance(t, CustomTask)
        assert t.status_indicator == "complete"
    # Titles in order
    assert tasks[0].title == "Reading x.py"
    assert tasks[1].title == "Running: ls"
    assert tasks[2].title.startswith('Searching: "foo"')


async def test_no_workflow_emitted_when_only_delta_arrives() -> None:
    """Delta-state envelopes must NOT trigger workflow creation."""
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="parent:1") as (_c, tracker, feed):
            await feed.start()
            for seq in range(1, 5):
                await gw.broadcast_chat(
                    "parent:1",
                    state="delta",
                    run_id="r",
                    seq=seq,
                    content=[{"type": "text", "text": "hi" * seq}],
                )
            # Give the reader loop a chance.
            await asyncio.sleep(0.15)
            events = await _drain(feed)
    assert events == []
    assert tracker.opened is False  # never lazy-opened


async def test_error_state_marks_all_open_complete_and_appends_terminal() -> None:
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="p") as (_c, tracker, feed):
            await feed.start()
            # Open two loading tasks.
            await gw.broadcast_chat(
                "p",
                state="final",
                run_id="r",
                seq=1,
                content=[
                    {
                        "type": "tool_use",
                        "id": "x",
                        "name": "Read",
                        "input": {"file_path": "a"},
                    }
                ],
            )
            await gw.broadcast_chat(
                "p",
                state="final",
                run_id="r",
                seq=2,
                content=[
                    {
                        "type": "tool_use",
                        "id": "y",
                        "name": "Bash",
                        "input": {"command": "ls"},
                    }
                ],
            )
            await _wait_for_queue(feed, 2)
            # Now an error. broadcast_chat doesn't carry errorMessage, so
            # emit a tailored frame directly via the fake gateway socket.
            import json as _json

            for ws in list(gw.subscriptions.get("p", [])):
                await ws.send(
                    _json.dumps(
                        {
                            "type": "event",
                            "event": "chat",
                            "payload": {
                                "runId": "r",
                                "sessionKey": "p",
                                "seq": 4,
                                "state": "error",
                                "errorMessage": "model timeout",
                            },
                            "seq": 4,
                        }
                    )
                )
            await _wait_for_queue(feed, 3)
            events = await _drain(feed)

    # All open tasks should be marked complete + a "Run errored: ..." terminal
    # task should be appended (also complete).
    tasks = events[0].item.workflow.tasks
    assert any(
        isinstance(t, CustomTask) and t.title.startswith("Run errored:")
        for t in tasks
    )
    for t in tasks:
        assert t.status_indicator == "complete", (t.title, t.status_indicator)
    assert feed.terminal is True


async def test_child_subscription_opens_on_spawnedby() -> None:
    """When a frame for an unknown sessionKey arrives carrying
    spawnedBy=<parent>, the feed should open a child subscription."""
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="hub") as (_c, _t, feed):
            await feed.start()
            # Simulate a parent-side `sessions_spawn` tool_use (so we know
            # the child exists conceptually).
            await gw.broadcast_chat(
                "hub",
                state="final",
                run_id="hub-r",
                seq=1,
                content=[
                    {
                        "type": "tool_use",
                        "id": "sp",
                        "name": "sessions_spawn",
                        "input": {"agentId": "andy", "task": "do work"},
                    }
                ],
            )
            # Now the child broadcasts a frame BEFORE we explicitly subscribe.
            # Fake gateway only routes to explicit subscribers, so we send via
            # all connections (the child key is not subscribed yet by client).
            # We use a side-channel: broadcast through the parent's connection
            # with the child's sessionKey + spawnedBy field, so the WS client's
            # filter sees a chat event on a known subscription (the parent) —
            # then the feed detects the new sessionKey via spawnedBy.
            for ws in list(gw.subscriptions.get("hub", [])):
                import json

                await ws.send(
                    json.dumps(
                        {
                            "type": "event",
                            "event": "chat",
                            "payload": {
                                "runId": "child-r",
                                "sessionKey": "andy-child",
                                "spawnedBy": "hub",
                                "seq": 1,
                                "state": "final",
                                "message": {
                                    "role": "assistant",
                                    "content": [
                                        {
                                            "type": "tool_use",
                                            "id": "ct",
                                            "name": "Read",
                                            "input": {"file_path": "z.py"},
                                        }
                                    ],
                                },
                            },
                            "seq": 99,
                        }
                    )
                )
            # Wait for feed to open the child subscription.
            await asyncio.sleep(0.2)
    # The fake gateway should have recorded a subscribe call for the child key.
    keys_subscribed = [k for (k, _) in gw.subscribe_calls]
    assert "andy-child" in keys_subscribed


async def test_replay_idempotent_via_translator() -> None:
    """A duplicate final envelope (post-reconnect replay) must not double-add."""
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="p") as (_c, _t, feed):
            await feed.start()
            for _ in range(2):  # second call simulates replay before client dedup
                await gw.broadcast_chat(
                    "p",
                    state="final",
                    run_id="r",
                    seq=1,
                    content=[
                        {
                            "type": "tool_use",
                            "id": "z",
                            "name": "Read",
                            "input": {"file_path": "x"},
                        }
                    ],
                )
            await asyncio.sleep(0.15)
            events = await _drain(feed)
    # WS client itself drops the duplicate by (sessionKey, runId, seq) before
    # the feed sees it. We should have exactly one workflow-added event with
    # one task.
    added = [e for e in events if isinstance(e, ThreadItemAddedEvent)]
    assert len(added) == 1
    assert len(added[0].item.workflow.tasks) == 1


async def test_stop_is_idempotent_and_closes_subscriptions() -> None:
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="p") as (_c, _t, feed):
            await feed.start()
            await feed.stop()
            await feed.stop()  # idempotent
            assert feed.active is False


async def test_drain_returns_empty_when_queue_empty() -> None:
    async with FakeGateway() as gw:
        async with _feed(gw, parent_key="p") as (_c, _t, feed):
            await feed.start()
            out = await _drain(feed)
    assert out == []


# ============================================================================
# Option 1 — poll mode integration (sessions_history HTTP polling).
# Spec §7. Mocks the gateway HTTP via a fake httpx.AsyncClient transport.
# ============================================================================


class _FakeAsyncHttp:
    """In-process httpx.AsyncClient shim. Returns a queue of canned
    sessions_history responses for sequential POST calls."""

    def __init__(self, responses: list[list[dict]]) -> None:
        self._responses = list(responses)
        self.call_count = 0

    async def post(self, url: str, json: dict, headers: dict):  # noqa: A002
        self.call_count += 1
        if self._responses:
            msgs = self._responses.pop(0)
        else:
            msgs = []
        body = {
            "ok": True,
            "result": {
                "content": [
                    {
                        "type": "text",
                        "text": __import__("json").dumps(
                            {"sessionKey": "x", "messages": msgs}
                        ),
                    }
                ]
            },
        }
        class _Resp:
            def __init__(self, data):
                self._data = data
            def raise_for_status(self):
                return None
            def json(self):
                return self._data
        return _Resp(body)

    async def aclose(self):
        return None


async def test_poll_mode_polls_sessions_history_and_emits_tasks() -> None:
    """End-to-end: poll mode tails sessions_history, dedupes by
    __openclaw.seq, runs new toolcall + tool_result messages through the
    translator, and emits WorkflowItem + per-tool task events."""
    # First poll: empty.
    # Second poll: assistant message carrying toolcall+tool_result both in
    # one content array, seq=5.
    # Third poll: same message again (dedup must skip).
    msg_with_tool = {
        "role": "assistant",
        "content": [
            {
                "type": "toolcall",
                "id": "toolu_42",
                "name": "Bash",
                "arguments": {"command": "ls /tmp", "description": ""},
            },
            {
                "type": "tool_result",
                "tool_use_id": "toolu_42",
                "content": "a\nb",
                "name": "Bash",
                "is_error": False,
            },
        ],
        "__openclaw": {"seq": 5, "id": "msg-1"},
    }
    fake_http = _FakeAsyncHttp(
        responses=[
            [],
            [msg_with_tool],
            [msg_with_tool],  # replay — must be deduped
        ]
    )

    tracker = ProgressTracker(thread_id="t-poll-1", enabled=True)
    feed = RealtimeProgressFeed(
        parent_session_key="thread_poll_session",
        tracker=tracker,
        mode="poll",
        http_client=fake_http,  # type: ignore[arg-type]
        gateway_url="http://localhost:18789",
        gateway_token="ignored-in-fake",
        poll_interval_s=0.02,
    )
    try:
        await feed.start()
        # Spin until we see at least 2 events (workflow ADDED + complete) AND
        # at least 3 polls have happened (so the dedupe replay was exercised).
        deadline = asyncio.get_running_loop().time() + 2.0
        while asyncio.get_running_loop().time() < deadline:
            if feed._event_queue.qsize() >= 2 and fake_http.call_count >= 3:
                break
            await asyncio.sleep(0.02)
        events = await _drain(feed)
    finally:
        await feed.stop()

    assert events, "no events from poll feed"
    assert isinstance(events[0], ThreadItemAddedEvent)
    wf = events[0].item
    assert isinstance(wf, WorkflowItem)
    # The workflow should now have exactly ONE task — the Bash one — and it
    # should be marked complete (since result was co-located on same envelope).
    # Replay-dedupe proof: even though the same message arrives twice from
    # the fake http, only ONE task exists.
    tasks = wf.workflow.tasks
    assert len(tasks) == 1, [t.title for t in tasks]
    t = tasks[0]
    assert isinstance(t, CustomTask)
    assert t.title.startswith("Running: ls /tmp")
    assert t.status_indicator == "complete"
    assert fake_http.call_count >= 3  # all three poll-responses consumed


async def test_poll_mode_skips_messages_without_seq() -> None:
    """Defensive: messages missing __openclaw.seq are skipped (not
    translated) — they can't be deduped reliably."""
    fake_http = _FakeAsyncHttp(
        responses=[
            [{"role": "assistant", "content": "no meta", "__openclaw": None}],
            [{"role": "assistant", "content": [{"type": "text", "text": "x"}]}],
        ]
    )
    tracker = ProgressTracker(thread_id="t-poll-2", enabled=True)
    feed = RealtimeProgressFeed(
        parent_session_key="thread_x",
        tracker=tracker,
        mode="poll",
        http_client=fake_http,  # type: ignore[arg-type]
        poll_interval_s=0.02,
    )
    try:
        await feed.start()
        await asyncio.sleep(0.15)
        events = await _drain(feed)
    finally:
        await feed.stop()
    assert events == []  # no envelope had a usable seq -> nothing emitted


async def test_poll_mode_stop_cancels_loop_and_closes_http() -> None:
    fake_http = _FakeAsyncHttp(responses=[[], []])
    tracker = ProgressTracker(thread_id="t-poll-3", enabled=True)
    feed = RealtimeProgressFeed(
        parent_session_key="thread_y",
        tracker=tracker,
        mode="poll",
        http_client=fake_http,  # type: ignore[arg-type]
        poll_interval_s=0.02,
    )
    await feed.start()
    await asyncio.sleep(0.1)
    await feed.stop()
    # Internal task cleared.
    assert feed._poll_task is None
    # http_client was passed in (owns_http=False) so feed should NOT close it.
    # (Caller is responsible.) The fake's aclose is idempotent regardless.


async def test_poll_mode_unknown_mode_raises_at_start() -> None:
    tracker = ProgressTracker(thread_id="t-bad", enabled=True)
    feed = RealtimeProgressFeed(
        parent_session_key="k",
        tracker=tracker,
        mode="bogus",  # type: ignore[arg-type]
    )
    with pytest.raises(ValueError):
        await feed.start()