openswarm/backend/tests/test_disconnect_resilience.py

"""Stress test for the WS resilience layer.

Goal: actively try to break the agent run with disconnects, drops,
reconnects, and concurrent broadcasts. The visible bug we're chasing
is "Network issue" toasts that flip a still-running task to a
terminal state. After these fixes the contract should be:

  1. The agent task NEVER dies because of a WS drop.
  2. Every event the server emits is replayable, in order, with no
     duplicates and no gaps, after any number of disconnects.
  3. Terminal events (completed/stopped/error) are always observable
     by a client that reconnects later — even if the only persistence
     of the event is the on-disk terminal log.
  4. Concurrent broadcasts (thinking deltas + tool calls + status
     changes from many tasks) preserve seq order == wire order.
  5. A client that's been gone too long for the ring buffer gets a
     `agent:gap_detected` instead of silent loss.

We don't run real Claude Code; we install a stub `agent_loop` that
emits the same WS event shapes a real run would (status, stream,
message, completed). That keeps the test fast (>200 iterations in
seconds) and self-contained.

Run:
    cd backend && .venv/bin/python -m pytest tests/test_disconnect_resilience.py -v
"""

from __future__ import annotations

import asyncio
import json
import os
import random
import tempfile
from unittest.mock import patch

import pytest
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.testclient import TestClient

# ---------------------------------------------------------------------------
# Boot env: route data to a tempdir BEFORE importing backend modules so
# the persistence dir for terminal events lives under our control.
# ---------------------------------------------------------------------------

_TMPROOT = tempfile.mkdtemp(prefix="openswarm-disconnect-test-")
os.environ.setdefault("OPENSWARM_DATA_DIR", _TMPROOT)

# Push the seq_log persist dir to a deterministic location too.
_SEQ_DIR = os.path.join(_TMPROOT, "seq_terminals")
os.makedirs(_SEQ_DIR, exist_ok=True)


@pytest.fixture(autouse=True)
def _patch_persist_dir():
    """Force the seq_log to use our tmp dir so we can assert on disk state."""
    from backend.apps.agents import seq_log as sl_mod

    # Rebuild the singleton with our test dir.
    new_store = sl_mod.SeqLogStore(persist_dir=_SEQ_DIR)
    monkey = patch.object(sl_mod, "seq_log", new_store)
    monkey.start()
    # Also patch the symbol re-exported into ws_manager's import scope.
    from backend.apps.agents import ws_manager as wm_mod
    wm_monkey = patch.object(wm_mod, "seq_log", new_store)
    wm_monkey.start()
    yield new_store
    monkey.stop()
    wm_monkey.stop()


# ---------------------------------------------------------------------------
# Minimal FastAPI app with the real WS endpoint logic. We import
# ws_manager directly and replicate the handler from backend/main.py
# without any of its auth middleware so the TestClient can connect
# without a token.
# ---------------------------------------------------------------------------


def _build_app(seq_log):
    """Replicates main.py's WS handler + adds a /test/emit endpoint
    so the test thread can drive event emission through the same
    event loop as the WS handler — avoiding the cross-loop hazards
    of `asyncio.run()` mid-test."""
    from backend.apps.agents.ws_manager import ws_manager

    app = FastAPI()

    @app.websocket("/ws/agents/{session_id}")
    async def ws_session(websocket: WebSocket, session_id: str):
        await ws_manager.connect_session(session_id, websocket)
        try:
            while True:
                data = await websocket.receive_text()
                msg = json.loads(data)
                event = msg.get("event")
                payload = msg.get("data", {})
                if event == "client:hello":
                    last_seq = int(payload.get("last_seq") or 0)
                    ack = await ws_manager.replay_to(session_id, websocket, last_seq)
                    await websocket.send_text(json.dumps({
                        "event": "server:hello",
                        "session_id": session_id,
                        "data": {
                            "connection_uuid": payload.get("connection_uuid", ""),
                            "current_seq": seq_log.current_seq(session_id),
                            "ack": ack,
                        },
                    }))
                elif event == "client:ping":
                    await websocket.send_text(json.dumps({
                        "event": "server:pong",
                        "session_id": session_id,
                        "data": {"nonce": payload.get("nonce")},
                    }))
        except WebSocketDisconnect:
            ws_manager.disconnect_session(session_id, websocket)

    @app.post("/test/emit/{session_id}")
    async def emit_events(session_id: str, body: dict):
        n = int(body.get("n", 0))
        terminate = body.get("terminate")  # str or None
        concurrent = int(body.get("concurrent", 1))
        await _emit_run(session_id, n, terminate=terminate, concurrent_tasks=concurrent)
        return {"ok": True, "current_seq": seq_log.current_seq(session_id)}

    return app


def _emit(client, session_id: str, n: int, terminate: str | None = None, concurrent: int = 1):
    """Drive event emission via the test-only HTTP endpoint."""
    r = client.post(f"/test/emit/{session_id}", json={
        "n": n, "terminate": terminate, "concurrent": concurrent,
    })
    assert r.status_code == 200, r.text
    return r.json()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


async def _emit_run(session_id: str, n_events: int, terminate: str | None = "completed", concurrent_tasks: int = 1):
    """Emit a synthetic agent run.

    `concurrent_tasks` lets the test stress the per-session lock by
    fanning out the broadcast across multiple coroutines. The seq
    log must still order them strictly.
    """
    from backend.apps.agents.ws_manager import ws_manager

    async def emit_chunk(start: int, count: int):
        for i in range(count):
            await ws_manager.send_to_session(session_id, "agent:stream_delta", {
                "session_id": session_id,
                "message_id": "m1",
                "delta": f"chunk-{start + i}",
            })
            # Yield to the scheduler so other coroutines interleave —
            # this is what surfaces the seq race if locking is wrong.
            await asyncio.sleep(0)

    if concurrent_tasks <= 1:
        await emit_chunk(0, n_events)
    else:
        per = n_events // concurrent_tasks
        tasks = [
            asyncio.create_task(emit_chunk(i * per, per))
            for i in range(concurrent_tasks)
        ]
        await asyncio.gather(*tasks)
        # Mop up the remainder so total event count is exact.
        rem = n_events - per * concurrent_tasks
        if rem > 0:
            await emit_chunk(per * concurrent_tasks, rem)

    if terminate is not None:
        await ws_manager.send_to_session(session_id, "agent:status", {
            "session_id": session_id,
            "status": terminate,
        })


# ---------------------------------------------------------------------------
# Unit-level: seq log fundamentals
# ---------------------------------------------------------------------------


def test_seq_monotonic_under_concurrency(_patch_persist_dir):
    """200 concurrent broadcasts must yield strictly monotonic seq."""
    app = _build_app(_patch_persist_dir)
    sid = "session-conc-1"
    with TestClient(app) as client:
        _emit(client, sid, n=200, terminate=None, concurrent=4)
    _, newest, events = _patch_persist_dir.replay(sid, 0)
    assert newest == 200
    seqs = [json.loads(s)["seq"] for s in events]
    assert seqs == sorted(seqs)
    assert len(set(seqs)) == len(seqs)


def test_terminal_event_persisted(_patch_persist_dir):
    app = _build_app(_patch_persist_dir)
    sid = "session-term-1"
    with TestClient(app) as client:
        _emit(client, sid, n=0, terminate="completed")
    raw = _patch_persist_dir.load_terminal(sid)
    assert raw is not None
    obj = json.loads(raw)
    assert obj["event"] == "agent:status"
    assert obj["data"]["status"] == "completed"


def test_replay_after_eviction_reports_gap(_patch_persist_dir):
    app = _build_app(_patch_persist_dir)
    sid = "session-evict-1"
    with TestClient(app) as client:
        _emit(client, sid, n=700, terminate=None)
    oldest, newest, events = _patch_persist_dir.replay(sid, last_seq=10)
    assert newest == 700
    assert oldest is not None and oldest > 10
    # Replay only includes seqs > 10 that survived eviction.
    assert all(json.loads(s)["seq"] > 10 for s in events)


# ---------------------------------------------------------------------------
# Integration: full WS connect / disconnect / resume cycle
# ---------------------------------------------------------------------------


def test_resume_after_disconnect_recovers_all_events(_patch_persist_dir):
    """Simulate a single disconnect mid-run, then a clean resume."""
    app = _build_app(_patch_persist_dir)

    sid = "session-res-1"
    received: list[dict] = []

    with TestClient(app) as client:
        # Phase 1: connect, hello, see N events, then close abruptly.
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 0, "connection_uuid": "c1"}}))
            hello = json.loads(ws.receive_text())
            assert hello["event"] == "server:hello"
            # Inject a few events via the /test/emit endpoint.
            _emit(client, sid, n=10, terminate=None)
            for _ in range(10):
                received.append(json.loads(ws.receive_text()))
            assert len(received) == 10
            assert received[-1]["seq"] == 10

        # Phase 2: between connections, the server keeps emitting. The
        # agent task is alive; only the WS is gone.
        _emit(client, sid, n=10, terminate="completed")

        # Phase 3: reconnect with last_seq=10, expect replay of seq 11..21
        # (10 deltas + 1 status), then the server:hello ack.
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 10, "connection_uuid": "c2"}}))
            replay: list[dict] = []
            while True:
                msg = json.loads(ws.receive_text())
                if msg["event"] == "server:hello":
                    break
                replay.append(msg)
            seqs = [m["seq"] for m in replay]
            assert seqs == list(range(11, 22)), f"unexpected replay seqs: {seqs}"
            statuses = [m for m in replay if m["event"] == "agent:status"]
            assert len(statuses) == 1
            assert statuses[0]["data"]["status"] == "completed"


def test_terminal_event_visible_after_full_eviction(_patch_persist_dir):
    """If the in-memory log is wiped (process restart simulation),
    a reconnecting client should still see the terminal event from
    disk — never a phantom 'running' spinner."""
    app = _build_app(_patch_persist_dir)
    sid = "session-evict-term-1"

    seq_log = _patch_persist_dir

    with TestClient(app) as client:
        _emit(client, sid, n=5, terminate="completed")

        # Simulate a process restart: clear the in-memory ring buffer
        # but keep the persisted terminal file.
        seq_log._per_session.pop(sid, None)

        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 0, "connection_uuid": "c1"}}))
            received = []
            while True:
                msg = json.loads(ws.receive_text())
                if msg["event"] == "server:hello":
                    received.append(msg)
                    break
                received.append(msg)
            terminals = [m for m in received if m["event"] == "agent:status"]
            assert len(terminals) == 1
            assert terminals[0]["data"]["status"] == "completed"


def test_gap_detected_when_buffer_evicted(_patch_persist_dir):
    """A client whose lastSeq is older than the oldest buffered seq
    should receive `agent:gap_detected` so it can REST-refresh,
    rather than silently miss events."""
    app = _build_app(_patch_persist_dir)
    sid = "session-gap-1"

    with TestClient(app) as client:
        # Fill the buffer past its limit so seq 1..200 are evicted.
        _emit(client, sid, n=700, terminate=None)
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 5, "connection_uuid": "c1"}}))
            saw_gap = False
            saw_hello = False
            while not saw_hello:
                msg = json.loads(ws.receive_text())
                if msg["event"] == "agent:gap_detected":
                    saw_gap = True
                elif msg["event"] == "server:hello":
                    saw_hello = True
                    assert msg["data"]["ack"]["ok"] is False
                    assert msg["data"]["ack"]["reason"] == "gap"
            assert saw_gap


def test_ping_pong_round_trip(_patch_persist_dir):
    app = _build_app(_patch_persist_dir)
    sid = "session-ping-1"

    with TestClient(app) as client:
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 0, "connection_uuid": "c1"}}))
            assert json.loads(ws.receive_text())["event"] == "server:hello"
            ws.send_text(json.dumps({"event": "client:ping", "data": {"nonce": "abc"}}))
            pong = json.loads(ws.receive_text())
            assert pong["event"] == "server:pong"
            assert pong["data"]["nonce"] == "abc"


# ---------------------------------------------------------------------------
# The big one: hundreds of randomized disconnect scenarios.
# ---------------------------------------------------------------------------


N_STRESS_ITERATIONS = int(os.environ.get("DISCONNECT_STRESS_N", "500"))


@pytest.mark.parametrize("iteration", range(N_STRESS_ITERATIONS))
def test_stress_random_disconnect(iteration, _patch_persist_dir):
    """Each iteration: a random number of events, a random number of
    disconnects at random points, optionally ending in a terminal
    status. After all reconnects, the client must have observed
    every event exactly once, in seq order, and the terminal event
    if one was emitted."""
    rng = random.Random(iteration)  # deterministic per iteration
    app = _build_app(_patch_persist_dir)

    sid = f"session-stress-{iteration}"
    total_events = rng.randint(5, 80)
    n_disconnects = rng.randint(1, min(5, total_events // 2 or 1))
    will_terminate = rng.random() < 0.7  # 70% of runs reach a terminal

    # Disconnect points: each is a count of events emitted *before*
    # the WS drops. We deliberately exclude `total_events` itself
    # so the breakpoint list never collides with the appended final
    # iteration (which is when the optional terminal status fires).
    if total_events > 1:
        breakpoints = sorted(rng.sample(range(1, total_events), min(n_disconnects, total_events - 1)))
    else:
        breakpoints = []
    seen: dict[int, dict] = {}  # seq -> event payload
    last_seq = 0

    with TestClient(app) as client:
        emitted_so_far = 0
        for bp in breakpoints + [total_events]:
            # Open a fresh socket, hello with our last_seq.
            with client.websocket_connect(f"/ws/agents/{sid}") as ws:
                ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": last_seq, "connection_uuid": f"c-{rng.random()}"}}))

                # Drain until server:hello, recording any replayed events.
                while True:
                    msg = json.loads(ws.receive_text())
                    if msg["event"] == "server:hello":
                        break
                    if "seq" in msg:
                        seen[msg["seq"]] = msg
                        last_seq = max(last_seq, msg["seq"])

                to_emit = bp - emitted_so_far
                emitted_so_far = bp
                terminate = "completed" if (bp == total_events and will_terminate) else None

                # Drive the emit through the test app's HTTP endpoint so
                # the broadcast happens on the same event loop as the WS
                # handler. Using asyncio.run() here would create an
                # isolated loop and re-bind the per-session asyncio.Lock
                # to a different loop, which is hostile to anyio's
                # blocking-portal pattern.
                _emit(client, sid, n=to_emit, terminate=terminate)

                expected = to_emit + (1 if terminate else 0)
                for _ in range(expected):
                    msg = json.loads(ws.receive_text())
                    seen[msg["seq"]] = msg
                    last_seq = max(last_seq, msg["seq"])

            # Closing the with-block disconnects the WS. The loop
            # opens a fresh socket on the next iteration.

    # ----- Assertions: completeness, ordering, no dups, terminal -----
    expected_total = total_events + (1 if will_terminate else 0)
    assert len(seen) == expected_total, f"missing events: expected {expected_total}, got {len(seen)}"
    seqs = sorted(seen.keys())
    assert seqs == list(range(1, expected_total + 1)), f"non-contiguous seqs: {seqs[:5]}...{seqs[-5:]}"
    if will_terminate:
        last = seen[expected_total]
        assert last["event"] == "agent:status"
        assert last["data"]["status"] == "completed"


# ---------------------------------------------------------------------------
# Concurrent broadcast: many fan-out coroutines must preserve seq order
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("trial", range(30))
def test_concurrent_broadcast_preserves_order(trial, _patch_persist_dir):
    """8 coroutines fanning out 400 events under the per-session lock.
    Drives the emit through the TestClient's portal so we use the
    real event loop the rest of the WS layer runs on."""
    app = _build_app(_patch_persist_dir)
    sid = f"session-conc-{trial}"

    with TestClient(app) as client:
        _emit(client, sid, n=400, terminate="completed", concurrent=8)

    oldest, newest, events = _patch_persist_dir.replay(sid, last_seq=0)
    assert newest == 401  # 400 deltas + 1 status
    seqs = [json.loads(s)["seq"] for s in events]
    assert seqs == sorted(seqs)
    # Each seq appears exactly once in the buffer.
    assert len(seqs) == len(set(seqs))


# ---------------------------------------------------------------------------
# Auth/security smoke: the WS endpoint here is unauth'd by design (test
# scaffolding) — but main.py's _ws_auth_ok must remain in place. This
# test pins that contract so a future refactor can't accidentally
# strip it.
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# Extra stress: terminate happens INSIDE a disconnect window. The
# client must see the terminal event on its next reconnect (whether
# from ring buffer or persisted disk record).
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("trial", range(50))
def test_terminate_during_disconnect_is_observable(trial, _patch_persist_dir):
    rng = random.Random(1000 + trial)
    app = _build_app(_patch_persist_dir)
    sid = f"session-mid-term-{trial}"
    n_pre = rng.randint(0, 40)
    n_post = rng.randint(0, 40)
    seen: dict[int, dict] = {}
    last_seq = 0

    with TestClient(app) as client:
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": 0, "connection_uuid": "c1"}}))
            assert json.loads(ws.receive_text())["event"] == "server:hello"
            if n_pre:
                _emit(client, sid, n=n_pre, terminate=None)
                for _ in range(n_pre):
                    msg = json.loads(ws.receive_text())
                    seen[msg["seq"]] = msg
                    last_seq = max(last_seq, msg["seq"])
        # Disconnected. Emit the rest + terminate while WS is gone.
        _emit(client, sid, n=n_post, terminate="completed")
        # Reconnect. We expect to receive everything from last_seq+1
        # through to the terminal — possibly via disk if the buffer
        # rolled (it won't here; numbers are small).
        with client.websocket_connect(f"/ws/agents/{sid}") as ws:
            ws.send_text(json.dumps({"event": "client:hello", "data": {"last_seq": last_seq, "connection_uuid": "c2"}}))
            while True:
                msg = json.loads(ws.receive_text())
                if msg["event"] == "server:hello":
                    break
                if "seq" in msg:
                    seen[msg["seq"]] = msg

    expected = n_pre + n_post + 1
    assert len(seen) == expected
    seqs = sorted(seen.keys())
    assert seqs == list(range(1, expected + 1))
    last = seen[expected]
    assert last["event"] == "agent:status"
    assert last["data"]["status"] == "completed"


# ---------------------------------------------------------------------------
# Sanity: an explicit `WebSocketDisconnect` MUST NOT cancel the
# underlying agent task. We don't have a real agent here, but we can
# at least assert that the ws_manager's disconnect path doesn't touch
# any task registry.
# ---------------------------------------------------------------------------


def test_disconnect_does_not_touch_agent_task(_patch_persist_dir):
    """If a future refactor adds task cancellation to disconnect_session,
    this test will catch it. We import agent_manager lazily so the
    `tasks` dict starts empty; we register a sentinel task and confirm
    disconnect_session doesn't poke it."""
    from backend.apps.agents.ws_manager import ws_manager
    # Insert a real Future into a parallel registry to mimic
    # `agent_manager.tasks[session_id]` and confirm ws_manager
    # never reaches into it. We don't import agent_manager (heavy);
    # we just inspect the source.
    import inspect
    src = inspect.getsource(ws_manager.disconnect_session)
    assert "cancel" not in src.lower()
    assert "agent_manager" not in src
    assert "tasks" not in src


def test_main_ws_endpoints_still_gated_by_auth(_patch_persist_dir):
    src = open(os.path.join(os.path.dirname(__file__), "..", "main.py")).read()
    assert "_ws_auth_ok(websocket)" in src, (
        "main.py WS endpoints must still call _ws_auth_ok before accepting "
        "the connection — otherwise any local web page can read agent traffic."
    )
    # And the disconnect handler must NOT call any task-cancel helper
    # — that's the regression we're guarding against.
    assert "stop_agent" not in src.split("WebSocketDisconnect")[1].split("def ")[0], (
        "WebSocketDisconnect handler must not cancel the agent task."
    )