fix: reconnect beaver terminal websocket before retrying turn

test_beaver_terminal_client.py

@@ -1,16 +1,32 @@
+import asyncio
 import json
+import sys
+from pathlib import Path
 
 import aiohttp
 import pytest
 from aiohttp import web
 
-from custom.beaver_terminal_client import (
-    BeaverTerminalClient,
-    BeaverTerminalError,
-    MessageIdGenerator,
-    build_connect_frame,
-    build_message_frame,
-)
+if __name__ == "__main__":
+    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+    raise SystemExit(pytest.main([__file__]))
+
+try:
+    from custom.beaver_terminal_client import (
+        BeaverTerminalClient,
+        BeaverTerminalError,
+        MessageIdGenerator,
+        build_connect_frame,
+        build_message_frame,
+    )
+except ModuleNotFoundError:
+    from beaver_terminal_client import (
+        BeaverTerminalClient,
+        BeaverTerminalError,
+        MessageIdGenerator,
+        build_connect_frame,
+        build_message_frame,
+    )
 
 
 def test_build_connect_frame_uses_stable_peer_id() -> None:
@@ -318,3 +334,76 @@ async def test_client_ping_sends_ping_and_waits_for_pong(unused_tcp_port: int) -
     finally:
         await client.close()
         await runner.cleanup()
+
+
+async def test_client_reconnects_with_same_peer_id_when_socket_closes_before_send(
+    unused_tcp_port: int,
+) -> None:
+    connect_peer_ids: list[str] = []
+    connection_count = 0
+
+    async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+        nonlocal connection_count
+        connection_count += 1
+        current_connection = connection_count
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+
+        async for message in ws:
+            frame = json.loads(message.data)
+            if frame["type"] == "connect":
+                connect_peer_ids.append(frame["peer_id"])
+                await ws.send_json(
+                    {
+                        "type": "connected",
+                        "channel_id": "terminal-dev",
+                        "session_id": "terminal-dev:local:device-001",
+                    }
+                )
+                if current_connection == 1:
+                    await ws.close()
+            elif frame["type"] == "message":
+                await ws.send_json(
+                    {
+                        "type": "ack",
+                        "message_id": frame["message_id"],
+                        "session_id": "terminal-dev:local:device-001",
+                        "accepted": True,
+                    }
+                )
+                await ws.send_json(
+                    {
+                        "type": "message",
+                        "role": "assistant",
+                        "message_id": frame["message_id"],
+                        "run_id": "run-2",
+                        "text": "reply after reconnect",
+                        "finish_reason": "stop",
+                    }
+                )
+
+        return ws
+
+    app = web.Application()
+    app.router.add_get("/api/channels/terminal-dev/ws", websocket_handler)
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, "127.0.0.1", unused_tcp_port)
+    await site.start()
+
+    client = BeaverTerminalClient(
+        url=f"http://127.0.0.1:{unused_tcp_port}/api/channels/terminal-dev/ws",
+        peer_id="device-001",
+        device_name="desk-terminal",
+    )
+
+    try:
+        await client.connect()
+        await asyncio.sleep(0.01)
+        reply = await client.send_text("hello")
+    finally:
+        await client.close()
+        await runner.cleanup()
+
+    assert reply == "reply after reconnect"
+    assert connect_peer_ids == ["device-001", "device-001"]