diff --git a/src/textual_webterm/local_server.py b/src/textual_webterm/local_server.py index fba0305..b4a0bb7 100644 --- a/src/textual_webterm/local_server.py +++ b/src/textual_webterm/local_server.py @@ -366,11 +366,13 @@ class LocalServer: self._websocket_connections[route_key] = ws session_id = self.session_manager.routes.get(RouteKey(route_key)) + session = None if session_id is not None: session = self.session_manager.get_session(session_id) if session is None or not session.is_running(): self.session_manager.on_session_end(session_id) session_id = None + session = None else: # Force terminal redraw on reconnect to avoid blank screen if hasattr(session, 'force_redraw'): @@ -378,8 +380,14 @@ class LocalServer: if hasattr(session, 'send_bytes'): await session.send_bytes(CLEAR_AND_REDRAW_SEQ.encode('utf-8')) + session_created = session_id is not None + if session_created and session is not None and hasattr(session, 'get_replay_buffer'): + replay = await session.get_replay_buffer() + if replay: + await ws.send_bytes(replay) + try: async for msg in ws: if msg.type == WSMsgType.TEXT: @@ -853,12 +861,14 @@ class LocalServer: return web.Response(text=html_content, content_type="text/html") async def handle_session_data(self, route_key: RouteKey, data: bytes) -> None: + self.mark_route_activity(str(route_key)) ws = self._websocket_connections.get(route_key) if ws is None: return await ws.send_bytes(data) async def handle_binary_message(self, route_key: RouteKey, payload: bytes) -> None: + self.mark_route_activity(str(route_key)) ws = self._websocket_connections.get(route_key) if ws is None: return diff --git a/tests/test_local_server_unit.py b/tests/test_local_server_unit.py index adc1af0..1bd64e8 100644 --- a/tests/test_local_server_unit.py +++ b/tests/test_local_server_unit.py @@ -694,7 +694,17 @@ class TestLocalServerMoreCoverage: assert queue.get_nowait() == "existing" @pytest.mark.asyncio - async def test_mark_route_activity_triggers_notification(self, server_with_no_apps): + async def test_handle_session_data_marks_activity(self, server_with_no_apps, monkeypatch): + ws = MagicMock() + ws.send_bytes = AsyncMock() + server_with_no_apps._websocket_connections["rk"] = ws + server_with_no_apps._route_last_activity["rk"] = 0.0 + + await server_with_no_apps.handle_session_data("rk", b"data") + assert server_with_no_apps._route_last_activity["rk"] > 0.0 + ws.send_bytes.assert_awaited_once_with(b"data") + + def test_mark_route_activity_triggers_notification(self, server_with_no_apps): """Test that mark_route_activity triggers SSE notification.""" import asyncio diff --git a/tests/test_local_server_websocket_integration.py b/tests/test_local_server_websocket_integration.py index 53a9ef4..d574941 100644 --- a/tests/test_local_server_websocket_integration.py +++ b/tests/test_local_server_websocket_integration.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from unittest.mock import AsyncMock import pytest from aiohttp import WSMsgType, web @@ -59,9 +60,16 @@ async def test_websocket_creates_session_on_resize(tmp_path): server.session_manager.routes["test"] = "sid" server.session_manager.sessions["sid"] = DummySession() + # Replay buffer should be sent on reconnect + replay_session = server.session_manager.sessions["sid"] + replay_session.get_replay_buffer = AsyncMock(return_value=b"replay") + client = await _make_client(server) try: ws = await client.ws_connect("/ws/test") + msg = await ws.receive(timeout=1) + assert msg.type == WSMsgType.BINARY + assert msg.data == b"replay" await ws.close() finally: await client.close()