deerflow-code/offline-backend-20260512/backend/tests/test_report_collaboration_sse.py

"""RC-BE-011: durable SSE, Live Hub, replay, heartbeat, coalescing."""

from __future__ import annotations

import asyncio
import json
from typing import Any

import httpx
import pytest
from fastapi import FastAPI

from app.gateway.routers import report_collaboration
from app.report_collaboration.agentscope_runtime.event_adapter import ConversationReplica, MemberStreamContext, SeqAllocator, make_emitter
from app.report_collaboration.execution.live_hub import (
    CoalescingEventSink,
    PublishingReportCollaborationStore,
    ReportCollaborationLiveHub,
    SseStreamSettings,
    coalesce_consecutive_deltas,
    iter_run_sse,
    replay_cursor,
    sanitize_event_payload,
)
from deerflow.config.model_config import ModelConfig
from deerflow.config.report_collaboration_config import ReportCollaborationConfig
from deerflow.persistence.report_collaboration import MemoryReportCollaborationStore


class _Cfg:
    report_collaboration = ReportCollaborationConfig(enabled=True)
    models = [ModelConfig(name="demo-chat", use="langchain_openai:ChatOpenAI", model="demo")]

    def get_model_config(self, name: str):
        return next((item for item in self.models if item.name == name), None)


def _sample_plan() -> dict:
    return {
        "id": "plan_sse",
        "proposal_group_id": "grp_sse",
        "title": "SSE",
        "strategy": "focused_fast",
        "summary": "sse",
        "rationale": "sse",
        "recommended": True,
        "estimated_duration_seconds": 60,
        "estimated_cost_level": "low",
        "requirement_revision": 0,
        "roles": [{"key": "researcher", "display_name": "研究员", "responsibility": "检索"}],
        "nodes": [
            {
                "id": "research",
                "label": "检索",
                "role_key": "researcher",
                "role_display_name": "研究员",
                "mission": "收集证据",
                "output_artifact_type": "EvidenceBundle",
                "allowed_tools": ["web_search"],
                "acceptance_criteria": ["每条主张有来源"],
                "depends_on": [],
            }
        ],
        "edges": [],
        "quality_gates": [],
        "validation": {"ok": True, "errors": []},
        "revision": 1,
        "status": "proposed",
    }


async def _seed_run(store, *, owner: str = "user_a") -> tuple[dict, dict]:
    session = await store.create_session(owner_id=owner, title="sse", idempotency_key="ik-session")
    plan = await store.insert_plan(session["id"], _sample_plan())
    await store.select_plan(session["id"], plan["id"], idempotency_key="ik-select", expected_revision=0)
    run = await store.create_run(session["id"], plan_id=plan["id"], idempotency_key="ik-run", expected_revision=0)
    return session, run


def _sse_events(body: str) -> list[dict[str, Any]]:
    return [json.loads(line[len("data: ") :]) for line in body.splitlines() if line.startswith("data: ")]


def _app(monkeypatch: pytest.MonkeyPatch, *, user: str | None = "user_a"):
    monkeypatch.setattr(report_collaboration, "get_app_config", lambda: _Cfg())

    async def _fake_user(_request=None):
        return user

    monkeypatch.setattr(report_collaboration, "get_current_user", _fake_user)
    hub = ReportCollaborationLiveHub()
    inner = MemoryReportCollaborationStore()
    store = PublishingReportCollaborationStore(inner, hub)
    app = FastAPI()
    app.state.report_collaboration_store = store
    app.state.report_collaboration_live_hub = hub
    app.include_router(report_collaboration.router)
    return app, store, hub


def test_replay_cursor_prefers_last_event_id() -> None:
    assert replay_cursor("12", 0) == 12
    assert replay_cursor("evt_not_a_seq", 4) == 4
    assert replay_cursor(None, 3) == 3


def test_sanitize_event_payload_redacts_secrets_and_stacks() -> None:
    cleaned = sanitize_event_payload(
        {
            "api_key": "sk-live-secret",
            "traceback": "Traceback (most recent call last):\n  File x",
            "delta": "份额上升",
            "nested": {"password": "hunter2", "excerpt": "ok"},
        }
    )
    assert cleaned["api_key"] == "[REDACTED]"
    assert "traceback" not in cleaned
    assert cleaned["delta"] == "份额上升"
    assert cleaned["nested"]["password"] == "[REDACTED]"
    assert cleaned["nested"]["excerpt"] == "ok"
    assert "sk-" not in str(cleaned)


def test_coalesce_consecutive_deltas_keeps_identity_and_concatenates() -> None:
    items = [
        ("message.created", {"message_id": "m1"}),
        ("message.delta", {"message_id": "m1", "delta": "份"}),
        ("message.delta", {"message_id": "m1", "delta": "额", "reasoning_delta": "想"}),
        ("message.delta", {"message_id": "m2", "delta": "其他"}),
        ("message.completed", {"message_id": "m1"}),
    ]
    merged = coalesce_consecutive_deltas(items)
    assert [item[0] for item in merged] == ["message.created", "message.delta", "message.delta", "message.completed"]
    assert merged[1][1]["delta"] == "份额"
    assert merged[1][1]["reasoning_delta"] == "想"
    assert merged[2][1]["delta"] == "其他"


@pytest.mark.asyncio
async def test_coalescing_sink_persists_one_delta_per_identity() -> None:
    store = MemoryReportCollaborationStore()
    session, run = await _seed_run(store)
    sink = CoalescingEventSink(store.append_event)
    await sink.emit(session_id=session["id"], run_id=run["id"], event_type="message.created", data={"message_id": "m1"})
    await sink.emit(session_id=session["id"], run_id=run["id"], event_type="message.delta", data={"message_id": "m1", "delta": "A"})
    await sink.emit(session_id=session["id"], run_id=run["id"], event_type="message.delta", data={"message_id": "m1", "delta": "B"})
    await sink.emit(session_id=session["id"], run_id=run["id"], event_type="message.completed", data={"message_id": "m1"})
    await sink.flush()
    types = [item["type"] for item in await store.list_events(run["id"], after_seq=0)]
    assert types == ["message.created", "message.delta", "message.completed"]
    delta = next(item for item in await store.list_events(run["id"], after_seq=0) if item["type"] == "message.delta")
    assert delta["data"]["delta"] == "AB"
    assert delta["seq"] == 2


@pytest.mark.asyncio
async def test_live_hub_fans_out_to_every_subscriber_and_drops_oldest_when_full() -> None:
    hub = ReportCollaborationLiveHub(queue_size=2)
    first = hub.subscribe("run_1")
    second = hub.subscribe("run_1")
    await hub.publish({"runId": "run_1", "seq": 1, "type": "heartbeat"})
    await hub.publish({"runId": "run_1", "seq": 2, "type": "heartbeat"})
    await hub.publish({"runId": "run_1", "seq": 3, "type": "heartbeat"})
    assert first.qsize() == 2
    assert second.qsize() == 2
    assert first.get_nowait()["seq"] == 2
    assert second.get_nowait()["seq"] == 2
    hub.unsubscribe("run_1", first)
    hub.unsubscribe("run_1", second)
    assert hub.subscriber_count("run_1") == 0


@pytest.mark.asyncio
async def test_publishing_store_sanitizes_then_publishes() -> None:
    hub = ReportCollaborationLiveHub()
    queue = hub.subscribe("pending")
    inner = MemoryReportCollaborationStore()
    store = PublishingReportCollaborationStore(inner, hub)
    session, run = await _seed_run(store)
    hub.unsubscribe("pending", queue)
    live = hub.subscribe(run["id"])
    envelope = await store.append_event(
        session_id=session["id"],
        run_id=run["id"],
        event_type="heartbeat",
        data={"api_key": "sk-should-hide", "traceback": "boom", "audit": "ok"},
    )
    assert envelope["data"]["api_key"] == "[REDACTED]"
    assert "traceback" not in envelope["data"]
    published = live.get_nowait()
    assert published["seq"] == envelope["seq"]
    assert published["data"]["audit"] == "ok"
    assert "sk-" not in str(published)


@pytest.mark.asyncio
async def test_iter_run_sse_replays_then_closes_on_terminal() -> None:
    store = MemoryReportCollaborationStore()
    session, run = await _seed_run(store)
    await store.append_event(session_id=session["id"], run_id=run["id"], event_type="run.status.changed", data={"status": "running"})
    await store.append_event(session_id=session["id"], run_id=run["id"], event_type="message.delta", data={"message_id": "m1", "delta": "hi"})
    await store.patch_run(run["id"], status="completed")
    chunks: list[str] = []

    async def never() -> bool:
        return False

    async for chunk in iter_run_sse(
        store=store,
        hub=None,
        run_id=run["id"],
        cursor=0,
        is_disconnected=never,
        settings=SseStreamSettings(flush_padding="", live_wait_seconds=0.01, db_poll_seconds=0.01, heartbeat_seconds=30, deadline_seconds=2),
    ):
        chunks.append(chunk)
    body = "".join(chunks)
    events = _sse_events(body)
    assert [item["seq"] for item in events] == [1, 2]
    assert events[0]["type"] == "run.status.changed"
    assert events[1]["data"]["delta"] == "hi"


@pytest.mark.asyncio
async def test_iter_run_sse_storage_error_keeps_the_connection() -> None:
    inner = MemoryReportCollaborationStore()
    session, run = await _seed_run(inner)
    await inner.append_event(session_id=session["id"], run_id=run["id"], event_type="heartbeat", data={"n": 1})
    await inner.patch_run(run["id"], status="completed")

    class _Flaky:
        def __init__(self) -> None:
            self.calls = 0

        def __getattr__(self, name: str):
            return getattr(inner, name)

        async def list_events(self, *args, **kwargs):
            self.calls += 1
            if self.calls == 1:
                raise RuntimeError("disk blip")
            return await inner.list_events(*args, **kwargs)

    chunks: list[str] = []

    async def never() -> bool:
        return False

    async for chunk in iter_run_sse(
        store=_Flaky(),
        hub=None,
        run_id=run["id"],
        cursor=0,
        is_disconnected=never,
        settings=SseStreamSettings(flush_padding="", live_wait_seconds=0.01, db_poll_seconds=0.05, heartbeat_seconds=30, deadline_seconds=2),
    ):
        chunks.append(chunk)
    body = "".join(chunks)
    assert ": storage-retry" in body
    events = _sse_events(body)
    assert events[0]["type"] == "heartbeat"


@pytest.mark.asyncio
async def test_http_stream_replays_after_seq_and_last_event_id(monkeypatch: pytest.MonkeyPatch) -> None:
    app, store, _hub = _app(monkeypatch)
    session, run = await _seed_run(store)
    await store.append_event(session_id=session["id"], run_id=run["id"], event_type="message.created", data={"message_id": "m1"})
    await store.append_event(session_id=session["id"], run_id=run["id"], event_type="message.delta", data={"message_id": "m1", "delta": "一"})
    await store.append_event(session_id=session["id"], run_id=run["id"], event_type="message.completed", data={"message_id": "m1"})
    await store.patch_run(run["id"], status="failed")
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        listed = await client.get(f"/api/report-collaboration/runs/{run['id']}/events?after_seq=1&limit=10")
        assert listed.status_code == 200
        assert [item["seq"] for item in listed.json()["events"]] == [2, 3]

        async with client.stream("GET", f"/api/report-collaboration/runs/{run['id']}/stream?after_seq=0") as response:
            assert response.status_code == 200
            body = "".join([chunk async for chunk in response.aiter_text()])
        events = _sse_events(body)
        assert [item["seq"] for item in events] == [1, 2, 3]

        async with client.stream(
            "GET",
            f"/api/report-collaboration/runs/{run['id']}/stream?after_seq=0",
            headers={"Last-Event-ID": "1"},
        ) as response:
            assert response.status_code == 200
            tail = "".join([chunk async for chunk in response.aiter_text()])
        assert [item["seq"] for item in _sse_events(tail)] == [2, 3]


@pytest.mark.asyncio
async def test_live_event_reaches_open_stream_and_does_not_reexecute_commands(monkeypatch: pytest.MonkeyPatch) -> None:
    app, store, hub = _app(monkeypatch)
    session, run = await _seed_run(store)
    transport = httpx.ASGITransport(app=app)

    async def produce() -> None:
        await asyncio.sleep(0.15)
        await store.append_event(
            session_id=session["id"],
            run_id=run["id"],
            event_type="message.delta",
            data={"message_id": "m-live", "delta": "并行成员A", "agent_run_id": "ar-a"},
        )
        await store.append_event(
            session_id=session["id"],
            run_id=run["id"],
            event_type="message.delta",
            data={"message_id": "m-live-b", "delta": "并行成员B", "agent_run_id": "ar-b"},
        )
        await store.patch_run(run["id"], status="completed")

    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        producer = asyncio.create_task(produce())
        async with client.stream("GET", f"/api/report-collaboration/runs/{run['id']}/stream?after_seq=0") as response:
            assert response.status_code == 200
            body = "".join([chunk async for chunk in response.aiter_text()])
        await producer
        events = _sse_events(body)
        texts = [item["data"].get("delta") for item in events if item["type"] == "message.delta"]
        assert "并行成员A" in texts
        assert "并行成员B" in texts
        snapshot = await store.get_snapshot(session["id"])
        assert all(item.get("operation") != "stream" for item in snapshot["commands"])
        assert [item["data"].get("agent_run_id") for item in events if item["type"] == "message.delta"] == ["ar-a", "ar-b"]
    assert hub.subscriber_count(run["id"]) == 0


@pytest.mark.asyncio
async def test_duplicate_streams_do_not_cancel_the_run(monkeypatch: pytest.MonkeyPatch) -> None:
    app, store, _hub = _app(monkeypatch)
    app.state.report_collaboration_sse_settings = SseStreamSettings(
        deadline_seconds=0.35,
        live_wait_seconds=0.05,
        db_poll_seconds=0.2,
        heartbeat_seconds=30,
        flush_padding="",
    )
    _session, run = await _seed_run(store)
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        first = await client.get(f"/api/report-collaboration/runs/{run['id']}/stream", timeout=2.0)
        second = await client.get(f"/api/report-collaboration/runs/{run['id']}/stream", timeout=2.0)
        assert first.status_code == 200
        assert second.status_code == 200
        assert first.headers["content-type"].startswith("text/event-stream")
        current = await client.get(f"/api/report-collaboration/runs/{run['id']}")
    assert current.json()["status"] == "queued"
    assert current.json()["cancel_requested"] is False


@pytest.mark.asyncio
async def test_other_user_cannot_stream(monkeypatch: pytest.MonkeyPatch) -> None:
    app, store, _hub = _app(monkeypatch, user="user_b")
    _session, run = await _seed_run(store, owner="user_a")
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        response = await client.get(f"/api/report-collaboration/runs/{run['id']}/stream")
    assert response.status_code == 404


@pytest.mark.asyncio
async def test_replay_includes_interrupted_tool_result(monkeypatch: pytest.MonkeyPatch) -> None:
    app, store, _hub = _app(monkeypatch)
    session, run = await _seed_run(store)
    replica = ConversationReplica(
        MemberStreamContext(
            session_id=session["id"],
            run_id=run["id"],
            agent_run_id="ar-1",
            node_run_id="research-attempt1",
            phase_id="phase_1",
            display_name="研究员",
        ),
        make_emitter(session["id"], run["id"], SeqAllocator()),
    )
    mirrored = []
    mirrored.extend(replica.ingest({"type": "REPLY_START", "reply_id": "msg_1", "name": "研究员"}))
    mirrored.extend(replica.ingest({"type": "TOOL_CALL_START", "reply_id": "msg_1", "tool_call_id": "call_1", "tool_call_name": "web_search"}))
    mirrored.extend(replica.ingest({"type": "REPLY_END", "reply_id": "msg_1", "finished_reason": "interrupted"}))
    sink = CoalescingEventSink(store.append_event)
    for event in mirrored:
        await sink.emit(session_id=session["id"], run_id=run["id"], event_type=event.type, data=event.data)
    await sink.flush()
    await store.patch_run(run["id"], status="cancelled")
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
        async with client.stream("GET", f"/api/report-collaboration/runs/{run['id']}/stream") as response:
            body = "".join([chunk async for chunk in response.aiter_text()])
    types = [item["type"] for item in _sse_events(body)]
    assert "tool_call.created" in types
    assert "tool_result.completed" in types
    completed = next(item for item in _sse_events(body) if item["type"] == "tool_result.completed")
    assert completed["data"]["tool_call_id"] == "call_1"