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