"""Phase 06 task 05 (US-09): heartbeats, drops, the seat grace and the rejoin — end to end. Two parts, one contract (docs/PROTOCOL.md "Heartbeats, drops, and seat grace"): * **the room level** — a `FakeClock`-driven `Room` (no test sleeps): `on_disconnect` keeps the seat and ghosts the sim's cook (one `GHOST_ON`), the park at `PARK_AFTER_S` (3 s) releases the carried item exactly once (a station set-down or `lost_items`), `on_reconnect` restores the seat (one `GHOST_OFF`, the carry kept), a reconnect of a live seat emits no `GHOST_OFF`, the `SeatAbandoned` at `guest_idle_timeout_s` empties the seat with a *forced* release (even before the 3 s park deadline), and a closed room runs the `on_close` hook once. * **the socket level** — a live uvicorn server (real time, the 20 Hz tick task) and `websockets` clients: an abrupt drop keeps the seat in the grace and parks the carry, a rejoin inside the grace restores the seat (`welcome.rejoined = true`, one `GHOST_OFF`, the snapshot stream resumes with a full frame within one tick's budget), a rejoin after the grace spectates (or gets `err room_gone` + close `4401` when the room has closed in the meantime), a 12 s silent client is closed `1000` and the room sees the standard drop, and the test-only seam (`POST /api/test/kill_socket/{code}/{token}`, `UNDERCOOKED_TEST_MODE=1` only) simulates a hard drop — `404` when the seam is off. """ from __future__ import annotations import asyncio import contextlib import copy import json import pathlib from collections.abc import Callable, Mapping from dataclasses import replace as dc_replace from typing import Any import httpx import pytest import websockets from undercooked.config import Settings from undercooked.server.clock import FakeClock from undercooked.server.room import Room, RoomState from undercooked.sim.content import ContentBundle from undercooked.sim.level import Level, LevelDef from undercooked.sim.model import ( EventKind, Item, ItemForm, ItemKind, PlayerState, StationKind, ) from undercooked.sim.recipe import Ingredient, Recipe, RecipeBook, Step pytestmark = pytest.mark.integration PLAYABLE_LEVELS_DIR = pathlib.Path(__file__).resolve().parents[1] / "fixtures" / "levels_playable" # --------------------------------------------------------------------- helpers def _uri(url: str, code: str) -> str: return url.replace("http://", "ws://", 1) + f"/ws/{code}" async def _connect( url: str, code: str, token: str | None = None, *, name: str | None = None, ) -> Any: """A connected client with the `hello` already sent.""" ws = await websockets.connect(_uri(url, code), open_timeout=10.0) if token is not None or name is not None: hello: dict[str, Any] = {"type": "hello", "p": 1} if token is not None: hello["token"] = token if name is not None: hello["name"] = name await ws.send(json.dumps(hello)) return ws async def _next_json(ws: Any, wait_s: float = 15.0) -> dict[str, Any]: return json.loads(await asyncio.wait_for(ws.recv(), wait_s)) async def _wait_for( predicate: Callable[[], bool], what: str, wait_s: float = 20.0 ) -> None: """Poll across threads (the server's loop does the work; this thread only observes it).""" loop = asyncio.get_running_loop() deadline = loop.time() + wait_s while not predicate(): if loop.time() > deadline: pytest.fail(f"timed out waiting for {what}") await asyncio.sleep(0.02) def _create_room(url: str, token: str | None = None) -> str: body: dict[str, Any] = {} if token is not None: body["token"] = token response = httpx.post(f"{url}/api/rooms", json=body, timeout=5.0) assert response.status_code == 201, response.text return response.json()["room_code"] def _join(url: str, code: str, token: str) -> None: response = httpx.post( f"{url}/api/rooms/{code}/join", json={"token": token}, timeout=5.0 ) assert response.status_code == 200, response.text def _kill_socket(url: str, code: str, token: str) -> httpx.Response: """The test-only seam (`UNDERCOOKED_TEST_MODE=1`): a hard drop.""" return httpx.post(f"{url}/api/test/kill_socket/{code}/{token}", timeout=5.0) def _ghost_count( frames: list[dict[str, Any]], kind: str ) -> int: """How many `kind` events a frame stream carries (the `snap` frame nests the PLAN §4 shape under `raw`, events in `raw.e` under `k`; the standalone `ev` frame under `kind`).""" count = 0 for frame in frames: if frame["type"] == "snap": count += sum( 1 for event in frame.get("raw", {}).get("e", ()) if event.get("k") == kind ) elif frame["type"] == "ev": count += sum( 1 for event in frame.get("events", ()) if event.get("kind") == kind ) return count def _sim_player(room: Room, token: str) -> Any: sim = room.sim_state assert sim is not None, "the round is not built" player = next((p for p in sim.players if p.id == token), None) assert player is not None, f"no sim player for {token!r}" return player def _counters(room: Room) -> Mapping[str, int]: sim = room.sim_state assert sim is not None, "the round is not built" return sim.counters def _seat_connected(room: Room, token: str) -> bool: seat = room.seat_for_token(token) return seat is not None and seat.connected is True def _seat_disconnected(room: Room, token: str) -> bool: seat = room.seat_for_token(token) return seat is not None and seat.connected is False def _give_carry(room: Room, token: str) -> None: """A test-side hand item: the park's release needs something to release (a live round's cooking is too flaky to script).""" sim = room.sim_state assert sim is not None sim = dc_replace( sim, players=tuple( dc_replace( p, carrying=Item(kind=ItemKind.TOMATO, form=ItemForm.RAW) ) if p.id == token else p for p in sim.players ), ) room.sim_state = sim async def _playing_room( url: str, app: Any, host_token: str, guest_token: str ) -> tuple[str, Room, Any, Any]: """Two seated players in `PLAYING` (the 3 s real-time countdown on the 20 Hz tick task).""" code = _create_room(url, host_token) _join(url, code, guest_token) host = await _connect(url, code, host_token, name="Host") await _next_json(host) # welcome await _next_json(host) # lobby guest = await _connect(url, code, guest_token, name="Guest") await _next_json(guest) # welcome await _next_json(guest) # lobby await guest.send(json.dumps({"type": "ready"})) await host.send(json.dumps({"type": "start"})) room = app.state.supervisor.get(code) assert room is not None await _wait_for( lambda: room.state is RoomState.PLAYING and room.sim_state is not None, f"room {code} to reach PLAYING", wait_s=30.0, ) return code, room, host, guest async def _close_clean(ws: Any) -> None: with contextlib.suppress(Exception): await ws.close() def _keepalive(ws: Any) -> asyncio.Task[None]: """The documented keep-alive (≥ 5 Hz budget; 0.5 Hz is plenty for a test client): a silent socket for `HEARTBEAT_DEAD_S` (12 s) is a dead client, so a test client that sits for a while must ping.""" state = {"n": 0} async def _loop() -> None: while True: await asyncio.sleep(2.0) state["n"] += 1 with contextlib.suppress(Exception): await ws.send(json.dumps({"type": "ping", "t": state["n"]})) return asyncio.create_task(_loop()) # ------------------------------------------------- the room level (FakeClock) def make_book() -> RecipeBook: return RecipeBook( recipes={ "soup": Recipe( id="soup", name="Soup", base_score=40, ingredients=[ Ingredient(kind=ItemKind.TOMATO, form=ItemForm.RAW, count=1) ], steps=[ Step(kind="serve", station=StationKind.SERVE, seconds=0.0) ], ) }, ingredients=[ItemKind.TOMATO], ) def make_level(level_id: str) -> Level: """A 7x4 kitchen; `time_limit` 60 s — a `PLAYING` round that will not end while the grace/park windows run.""" grid = [ "#######", "#S...V#", "#.....#", "#######", ] tiles: dict[str, dict[str, str]] = { ".": {"kind": "floor"}, "#": {"kind": "wall"}, "S": {"kind": "floor", "station": "spawn_shelf", "item": "tomato"}, "V": {"kind": "floor", "station": "serve"}, } data: dict[str, object] = { "id": level_id, "name": f"Ghost kitchen {level_id}", "order": 1, "time_limit": 60.0, "target_score": 1000, "grid": copy.deepcopy(grid), "tiles": copy.deepcopy(tiles), "spawns": [[2, 2]], "orders": { "pool": ["soup"], "concurrent": 1, "spawn_interval": [4, 6], "deadline": 30, }, } return Level.build(LevelDef.model_validate(data)) def make_bundle() -> ContentBundle: return ContentBundle(recipe_book=make_book(), levels=(make_level("ghost"),)) def make_room( **settings: Any, ) -> tuple[Room, FakeClock]: clock = FakeClock() bundle = make_bundle() room = Room( code="TEST1", settings=Settings(**settings), clock=clock, content=bundle, level_id=bundle.levels[0].id, ) return room, clock def _playing(room: Room) -> float: """A two-seat `PLAYING` room; returns the `now` the round starts at.""" room.add_seat("tok-a", "Ada", True) room.add_seat("tok-b", "Bob", False) room.set_ready("tok-b", True) room.request_start() room.tick_once(now=0.0) # the dt-0 prime (the countdown's first frame) for offset in (1.0, 2.0, 3.0): room.tick_once(now=offset) assert room.state is RoomState.PLAYING return 3.0 def test_on_disconnect_keeps_the_seat_and_ghosts_the_cook() -> None: room, _ = make_room() now = _playing(room) room.on_disconnect("tok-b", now=now) seat = room.seat_for_token("tok-b") assert seat is not None # the seat is KEPT for the rejoin path assert seat.connected is False player = _sim_player(room, "tok-b") assert player.connected is False assert player.state is PlayerState.DISCONNECTED events = [e for e in room.take_events() if e.kind is EventKind.GHOST_ON] assert len(events) == 1 # one `GHOST_ON` per drop # A second drop mark is a no-op: no second `GHOST_ON`. assert room.on_disconnect("tok-b", now=now) is True assert room.take_events() == () def test_park_releases_the_carried_item_after_three_seconds() -> None: room, _ = make_room() now = _playing(room) _give_carry(room, "tok-b") room.on_disconnect("tok-b", now=now) events = [e for e in room.take_events() if e.kind is EventKind.GHOST_ON] assert len(events) == 1 # one `GHOST_ON` per drop # One second before the deadline: no release yet. room.tick_once(now=now + 2.9) assert _sim_player(room, "tok-b").carrying is not None assert room.take_events() == () # The 3 s deadline (the room's tick, off `seat.last_seen`): the item # is released into `lost_items` (the tile has no station under the # spawn) with one `GHOST_PARK`. room.tick_once(now=now + 3.0) assert _sim_player(room, "tok-b").carrying is None assert _counters(room)["lost_items"] == 1 kinds = [e.kind for e in room.take_events()] assert EventKind.GHOST_PARK in kinds # A later tick cannot double-release. room.tick_once(now=now + 6.0) assert _counters(room)["lost_items"] == 1 assert room.take_events() == () def test_on_reconnect_restores_the_seat_and_keeps_the_carry() -> None: room, _ = make_room() now = _playing(room) _give_carry(room, "tok-b") room.on_disconnect("tok-b", now=now) room.on_reconnect("tok-b", now=now + 1.0) seat = room.seat_for_token("tok-b") assert seat is not None and seat.connected is True player = _sim_player(room, "tok-b") assert player.connected is True assert player.state is PlayerState.IDLE assert player.carrying is not None # the rejoin keeps the carry events = [e for e in room.take_events() if e.kind is EventKind.GHOST_OFF] assert len(events) == 1 # one `GHOST_OFF` (the drop's `GHOST_ON` drained) def test_reconnect_without_a_disconnect_emits_no_ghost_off() -> None: room, _ = make_room() _playing(room) assert room.on_reconnect("tok-b") is True # the seat was never a ghost assert room.take_events() == () # no `GHOST_OFF` def test_seat_abandoned_after_the_grace_releases_and_empties() -> None: room, _ = make_room(guest_idle_timeout_s=2.0) now = _playing(room) _give_carry(room, "tok-b") room.on_disconnect("tok-b", now=now) # Inside the grace: the seat stays (the rejoin window is open). room.tick_once(now=now + 1.5) assert room.seat_for_token("tok-b") is not None # Past the grace: the seat is emptied, the carry released at once, # one `GHOST_OFF`, the sim's player stays (parked). room.tick_once(now=now + 2.1) assert room.seat_for_token("tok-b") is None assert _sim_player(room, "tok-b").carrying is None assert _sim_player(room, "tok-b").connected is False assert _counters(room)["lost_items"] == 1 kinds = [e.kind for e in room.take_events()] assert EventKind.GHOST_OFF in kinds # The room may close once the last seat is gone too: the close # rules apply (the 30 s seat grace / the supervisor's GC). room.on_disconnect("tok-a", now=now + 2.1) room.tick_once(now=now + 4.2) assert room.seat_for_token("tok-a") is None assert room.is_empty def test_abandon_before_the_park_deadline_forces_the_release() -> None: """A 1 s grace runs out before the 3 s park deadline: the release must not wait for the park — the seat is going away.""" room, _ = make_room(guest_idle_timeout_s=1.0) now = _playing(room) _give_carry(room, "tok-b") room.on_disconnect("tok-b", now=now) room.tick_once(now=now + 1.1) # the grace expires, 1.9 s before the park assert room.seat_for_token("tok-b") is None assert _sim_player(room, "tok-b").carrying is None assert _counters(room)["lost_items"] == 1 def test_a_closed_room_runs_the_on_close_hook_once() -> None: calls: list[str] = [] async def _hook(room: Room) -> None: calls.append(room.code) clock = FakeClock() bundle = make_bundle() room = Room( code="TEST1", settings=Settings(), clock=clock, content=bundle, level_id=bundle.levels[0].id, on_close=_hook, ) room.add_seat("tok-a", "Ada", True) loop = asyncio.new_event_loop() try: loop.run_until_complete(room.close()) assert calls == ["TEST1"] # once, not twice loop.run_until_complete(room.close()) # idempotent: no second call assert calls == ["TEST1"] finally: loop.close() # ------------------------------------------------- the socket level (live) async def test_an_abrupt_drop_keeps_the_seat_and_parks_the_carry( ws_server: Any, ) -> None: with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR)) as (url, app): _code, room, host, guest = await _playing_room(url, app, "tok-a", "tok-b") _keep = _keepalive(host) _give_carry(room, "tok-b") # The drop: the client closes abruptly (no clean close). with contextlib.suppress(Exception): await guest.close() await _wait_for( lambda: _seat_disconnected(room, "tok-b"), "the seat to be marked disconnected", wait_s=10.0, ) seat = room.seat_for_token("tok-b") assert seat is not None # the seat is kept for the grace assert seat.connected is False player = _sim_player(room, "tok-b") assert player.connected is False assert player.state is PlayerState.DISCONNECTED # The park at 3 s (the room's tick, real time): the carry lands # in `lost_items`; the wire saw one `GHOST_ON` and one # `GHOST_PARK` (collected on the other player's stream). frames: list[dict[str, Any]] = [] await _wait_for( lambda: room.sim_state is not None and room.sim_state.counters.get("lost_items") == 1, "the park release", wait_s=15.0, ) # Drain the host's stream and count the ghost events. Bounded # by design: a live room's `snap` stream never goes quiet # while the round is `PLAYING` (this level's 150 s time limit # outlives the test), and both ghost events land within one # tick of the drop and of the park — a 10 s window after the # park release is ten times the margin; the `ConnectionClosed` # break still catches the round ending early. deadline = asyncio.get_running_loop().time() + 10.0 while asyncio.get_running_loop().time() < deadline: try: frames.append( json.loads(await asyncio.wait_for(host.recv(), 1.0)) ) except TimeoutError: continue # a quiet second is fine; the window bounds except websockets.exceptions.ConnectionClosed: break assert _ghost_count(frames, "ghost_on") == 1 assert _ghost_count(frames, "ghost_park") == 1 _keep.cancel() with contextlib.suppress(asyncio.CancelledError): await _keep async def test_rejoin_in_grace_restores_the_seat(ws_server: Any) -> None: with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR)) as (url, app): code, room, host, guest = await _playing_room(url, app, "tok-a", "tok-b") _keep = _keepalive(host) await _close_clean(guest) await _wait_for( lambda: _seat_disconnected(room, "tok-b"), "the seat to be marked disconnected", wait_s=10.0, ) # The rejoin inside the grace (a fresh socket, the same token). back = await _connect(url, code, "tok-b", name="Guest") welcome = await _next_json(back) assert welcome["type"] == "welcome" assert welcome["rejoined"] is True # the rejoin flag (task 05) await _wait_for( lambda: _seat_connected(room, "tok-b"), "the seat to be restored", wait_s=10.0, ) player = _sim_player(room, "tok-b") assert player.connected is True # The snapshot stream resumes with a full frame within one # tick's budget (the rejoining socket's `needs_full` resync). snap = None while snap is None: frame = json.loads(await asyncio.wait_for(back.recv(), 5.0)) if frame["type"] == "snap": snap = frame assert "d" not in snap.get("raw", {}) # a full frame, not a delta # The wire saw one `GHOST_OFF` (and the drop's one `GHOST_ON`). frames: list[dict[str, Any]] = [] for _ in range(30): try: frames.append(json.loads(await asyncio.wait_for(host.recv(), 1.0))) except TimeoutError: break except websockets.exceptions.ConnectionClosed: break assert _ghost_count(frames, "ghost_on") == 1 assert _ghost_count(frames, "ghost_off") >= 1 _keep.cancel() with contextlib.suppress(asyncio.CancelledError): await _keep async def test_rejoin_after_the_grace_spectates(ws_server: Any) -> None: with ws_server( levels_dir=str(PLAYABLE_LEVELS_DIR), guest_idle_timeout_s=2.0 ) as (url, app): code, room, _host, guest = await _playing_room(url, app, "tok-a", "tok-b") await _close_clean(guest) await _wait_for( lambda: room.seat_for_token("tok-b") is None, "the seat to be abandoned (the grace ran out)", wait_s=15.0, ) # The rejoin after the grace: the seat is gone — a late join in # `PLAYING` spectates (not a rejoin). back = await _connect(url, code, "tok-b", name="Guest") welcome = await _next_json(back) assert welcome["type"] == "welcome" assert welcome["rejoined"] is False assert room.seat_for_token("tok-b") is None assert room.spectator_for_token("tok-b") is not None async def test_a_rejoin_into_a_gone_room_says_room_gone(ws_server: Any) -> None: with ws_server( levels_dir=str(PLAYABLE_LEVELS_DIR), guest_idle_timeout_s=2.0, lobby_idle_timeout_s=3.0, ) as (url, app): # A solo host room in `LOBBY`: the drop empties it, and the # supervisor's GC closes it (the `LOBBY` idle rule, the 5 s # pass). code = _create_room(url, "tok-a") solo = await _connect(url, code, "tok-a", name="Solo") await _next_json(solo) # welcome await _next_json(solo) # lobby room = app.state.supervisor.get(code) assert room is not None await _close_clean(solo) await _wait_for( lambda: room.seat_for_token("tok-a") is None, "the seat to be abandoned", wait_s=15.0, ) await _wait_for( lambda: room.state is RoomState.CLOSED, "the room to close (the GC's idle rule)", wait_s=30.0, ) # A stale-token rejoin: `err room_gone` + close `4401`, before # any seat is taken — the client's "the room died" signal. The # close may race the open handshake (the accept went out before # the refusal): either the `err` frame arrives, or the `4401` # close does — both carry the same signal. closed: websockets.exceptions.ConnectionClosed | None = None late: Any = None try: late = await websockets.connect(_uri(url, code), open_timeout=10.0) await late.send(json.dumps({"type": "hello", "p": 1, "token": "tok-a"})) for _ in range(5): # the refusal close follows the `err` await asyncio.wait_for(late.recv(), 10.0) except websockets.exceptions.ConnectionClosed as exc: closed = exc assert closed is not None and closed.rcvd is not None assert closed.rcvd.code == 4401 async def test_a_dead_client_is_closed_by_the_heartbeat_watchdog(ws_server: Any) -> None: with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR)) as (url, app): code = _create_room(url, "tok-a") silent = await _connect(url, code, "tok-a", name="Silent") await _next_json(silent) # welcome await _next_json(silent) # lobby room = app.state.supervisor.get(code) assert room is not None # 12 s of silence (no `ping`): the watchdog closes `1000` and # the room sees the standard drop — indistinguishable from a # hard drop (one mark, one ghost). with pytest.raises(websockets.exceptions.ConnectionClosed) as exc: while True: await asyncio.wait_for(silent.recv(), 20.0) assert exc.value.rcvd is not None assert exc.value.rcvd.code == 1000 await _wait_for( lambda: _seat_disconnected(room, "tok-a"), "the seat to be marked disconnected", wait_s=10.0, ) async def test_the_test_seam_kills_a_socket(ws_server: Any) -> None: with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR), test_mode=True) as ( url, app, ): code, room, _host, guest = await _playing_room(url, app, "tok-a", "tok-b") # The seam force-closes the socket server-side: a hard drop for # the client, the standard `on_disconnect` for the room. response = _kill_socket(url, code, "tok-b") assert response.status_code == 200, response.text # The client sees the hard drop (past any buffered frames). closed = False for _ in range(50): try: await asyncio.wait_for(guest.recv(), 10.0) except websockets.exceptions.ConnectionClosed: closed = True break assert closed, "the killed socket did not close" await _wait_for( lambda: _seat_disconnected(room, "tok-b"), "the seat to be marked disconnected", wait_s=10.0, ) # An unknown token: `404 no_connection`. response = _kill_socket(url, code, "tok-nobody") assert response.status_code == 404 assert response.json()["code"] == "no_connection" async def test_the_test_seam_is_gated_by_test_mode(ws_server: Any) -> None: with ws_server(levels_dir=str(PLAYABLE_LEVELS_DIR)) as (url, _app): # `test_mode` unset: the seam is not registered at all. code = _create_room(url, "tok-a") solo = await _connect(url, code, "tok-a", name="Solo") await _next_json(solo) # welcome await _next_json(solo) # lobby response = _kill_socket(url, code, "tok-a") assert response.status_code == 404 # The socket is still alive (the seam never ran). await _close_clean(solo)