Harden WebSocket transport: pending-Future leak + duplicate-ID handshake reject (#374)

* Harden WebSocket transport: pending-Future leak + duplicate-ID handshake reject

Two small audit-v2 fixes in src/godot_ai/transport/websocket.py
covering #343 sub-findings #2 and #5. (Sub-finding #4 — errno.EADDRINUSE
portability — landed separately via PR #373.)

- #5 (P2): wrap `send_command`'s `ws.send` + `wait_for` in a
  try/finally that always pops `_pending`. Pre-fix, a `ws.send`
  raise (ConnectionClosed mid-send, transport error) leaked the
  Future entry forever; under churn the dict grew unbounded.
  Happy path still pops via the response receiver, so the finally
  pop is a no-op there.

- #2 (P1): reject a second handshake whose session_id is already
  registered (close code 4001), instead of silently overwriting
  the routing map. session_id format is `<slug>@<4hex>` — 16 bits
  of suffix is locally guessable, so without this any local peer
  could hijack an active session by impersonating its ID.
  Legitimate plugin reconnect after `editor_reload_plugin` first
  triggers ConnectionClosed -> unregister, so the new connect
  still lands cleanly.

Tests:
- tests/integration/test_websocket.py — TestDuplicateHandshake
  pins the reject + reconnect-after-clean-disconnect paths;
  TestPendingFutureCleanup pins the timeout-pop and
  send-failure-pop behaviors.

Live-smoked against a real Godot 4.6 editor: full GDScript suite
1037/1040 passed, duplicate-handshake rejected with close 4001,
original session unaffected.

https://claude.ai/code/session_016ijmCD5S6QfwJGJcc5Wirp

* Apply pending ruff format to test_mcp_tools and test_gdscript_no_adjacent_string_concat

Drive-by: `ruff format --check` was failing on these two test files on beta.
Reformatting them keeps the format-check CI step green for the audit-v2
transport-hardening PR. No behavioral change — purely whitespace /
line-wrapping deltas produced by `ruff format`.

https://claude.ai/code/session_016ijmCD5S6QfwJGJcc5Wirp

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: dsarno

src/godot_ai/transport/websocket.py

@@ -20,6 +20,11 @@
 
 DEFAULT_PORT = 9500
 
+## RFC 6455 reserves 4000-4999 for application-defined close codes; we use
+## 4001 to flag a handshake rejected for duplicate session_id so a debugging
+## peer can distinguish it from a normal close.
+_CLOSE_CODE_DUPLICATE_SESSION = 4001
+
 
 class GodotWebSocketServer:
     """Accepts connections from Godot editor plugins and routes commands."""
@@ -62,6 +67,23 @@ async def _handle_connection(self, ws: ServerConnection):
             data = json.loads(raw)
             handshake = HandshakeMessage.model_validate(data)
 
+            ## Reject duplicate session_id while the first peer is live —
+            ## otherwise the second handshake silently overwrites the
+            ## routing map (duplicate-ID hijack).
+            existing = self.registry.get(handshake.session_id)
+            if existing is not None:
+                logger.warning(
+                    "Rejecting duplicate handshake for session %s (existing pid=%s, project=%s)",
+                    handshake.session_id,
+                    existing.editor_pid,
+                    existing.project_path,
+                )
+                await ws.close(
+                    code=_CLOSE_CODE_DUPLICATE_SESSION,
+                    reason="session id already registered",
+                )
+                return
+
             session_id = handshake.session_id
             session = Session(
                 session_id=handshake.session_id,
@@ -159,12 +181,16 @@ async def send_command(
         future: asyncio.Future[CommandResponse] = asyncio.get_running_loop().create_future()
         self._pending[request.request_id] = future
 
-        await ws.send(request.model_dump_json())
-
+        ## Always pop on exit — the response receiver in _handle_connection
+        ## pops on the happy path, so this is a no-op there; on `ws.send`
+        ## raise / TimeoutError / cancellation it prevents Futures leaking
+        ## into _pending forever.
         try:
+            await ws.send(request.model_dump_json())
             return await asyncio.wait_for(future, timeout=timeout)
         except asyncio.TimeoutError:
-            self._pending.pop(request.request_id, None)
             raise TimeoutError(
                 f"Command {command} timed out after {timeout}s on session {session_id}"
             )
+        finally:
+            self._pending.pop(request.request_id, None)

tests/integration/test_mcp_tools.py

@@ -2067,9 +2067,7 @@ async def respond():
         assert result.data["shape_created"] is True
         assert result.data["size"]["x"] == 2.0
 
-    async def test_ambiguous_visual_candidates_preserved_in_structured_error(
-        self, mcp_stack
-    ):
+    async def test_ambiguous_visual_candidates_preserved_in_structured_error(self, mcp_stack):
         client, plugin = mcp_stack
         candidates = ["/Main/VisualA", "/Main/VisualB"]
 
@@ -3328,6 +3326,7 @@ async def test_concurrent_read_and_write_route_to_target_sessions(self, mcp_stac
         client, plugin_a = mcp_stack
         plugin_b = await self._connect_second_plugin("proj-b@0002")
         try:
+
             async def respond_a():
                 cmd = await plugin_a.recv_command()
                 assert cmd["command"] == "get_open_scenes"

tests/integration/test_websocket.py

@@ -561,6 +561,127 @@ async def test_rejected_request_does_not_register_session(self, harness):
         assert len(harness.registry) == before
 
 
+# ---------------------------------------------------------------------------
+# Duplicate-ID handshake hardening (#343 finding #2)
+# ---------------------------------------------------------------------------
+
+
+class TestDuplicateHandshake:
+    async def test_duplicate_session_id_handshake_is_rejected(self, harness):
+        ## Without rejection, a second handshake with the same session_id
+        ## silently overwrites both `_connections[session_id]` and the
+        ## registry entry — routing every subsequent command to the
+        ## attacker. session_id is `<slug>@<4hex>` so 16 bits of suffix is
+        ## locally guessable. Reject keeps the first peer authoritative.
+        first = await harness.connect_plugin(session_id="dup-target")
+        original_session = harness.registry.get("dup-target")
+        assert original_session is not None
+        original_pid = original_session.editor_pid
+
+        ## Hand-roll the second handshake so we observe the close on the
+        ## wire — `connect_plugin()` would assert on the missing ack.
+        ws2 = await websockets.connect(f"ws://127.0.0.1:{harness.port}")
+        await ws2.send(
+            json.dumps(
+                {
+                    "type": "handshake",
+                    "session_id": "dup-target",
+                    "godot_version": "4.4.1",
+                    "project_path": "/tmp/attacker",
+                    "plugin_version": "0.0.1",
+                    "protocol_version": 1,
+                    "readiness": "ready",
+                    "editor_pid": 9999,
+                }
+            )
+        )
+
+        ## Server should close us before sending an ack. Drain until the WS
+        ## reports closed; recv() will raise ConnectionClosed.
+        with pytest.raises(websockets.ConnectionClosed):
+            await asyncio.wait_for(ws2.recv(), timeout=2.0)
+
+        ## Original session must still be live and unaffected.
+        live = harness.registry.get("dup-target")
+        assert live is original_session, "registry entry was overwritten by duplicate"
+        assert live.editor_pid == original_pid
+        assert live.project_path != "/tmp/attacker"
+
+        ## Round-trip a command through the original to prove its WS is
+        ## still wired to the routing map (regression: silent overwrite
+        ## also hijacks `_connections[session_id]`).
+        client = GodotClient(harness.server, harness.registry)
+
+        async def mock_handler():
+            cmd = await first.recv_command()
+            await first.send_response(cmd["request_id"], {"alive": True})
+
+        handler = asyncio.create_task(mock_handler())
+        result = await client.send("ping", session_id="dup-target")
+        await handler
+        assert result == {"alive": True}
+
+        await first.close()
+
+    async def test_reconnect_after_clean_disconnect_succeeds(self, harness):
+        ## The reject must not break the legitimate plugin reconnect path
+        ## (e.g. after `editor_reload_plugin`): close → unregister →
+        ## fresh connect with the same session_id should succeed because
+        ## the registry entry has already been removed.
+        first = await harness.connect_plugin(session_id="reconnect-1")
+        await first.close()
+        await asyncio.sleep(0.1)  # let server process disconnect
+        assert harness.registry.get("reconnect-1") is None
+
+        second = await harness.connect_plugin(session_id="reconnect-1")
+        assert harness.registry.get("reconnect-1") is not None
+        await second.close()
+
+
+# ---------------------------------------------------------------------------
+# send_command pending-Future leak (#343 finding #5)
+# ---------------------------------------------------------------------------
+
+
+class TestPendingFutureCleanup:
+    async def test_timeout_pops_pending_entry(self, harness):
+        ## TimeoutError path always cleared the pending dict; this test
+        ## pins that behavior so a future refactor doesn't regress it.
+        plugin = await harness.connect_plugin(session_id="leak-timeout")
+        client = GodotClient(harness.server, harness.registry)
+
+        with pytest.raises(TimeoutError):
+            await client.send("never_responded", timeout=0.1)
+
+        assert harness.server._pending == {}, "TimeoutError should not leave entries in _pending"
+        await plugin.close()
+
+    async def test_send_failure_pops_pending_entry(self, harness):
+        ## If `ws.send` raises (e.g. ConnectionClosed mid-send), the
+        ## pending Future was previously leaked into `_pending` forever.
+        ## Force the failure by replacing the connection's send with one
+        ## that raises after the pending entry has been registered.
+        plugin = await harness.connect_plugin(session_id="leak-send")
+
+        ws = harness.server._connections["leak-send"]
+        boom = ConnectionError("simulated mid-send transport error")
+
+        async def raising_send(_payload: str) -> None:
+            raise boom
+
+        ws.send = raising_send  # type: ignore[assignment]
+
+        with pytest.raises(ConnectionError):
+            await harness.server.send_command(
+                session_id="leak-send",
+                command="will_fail",
+                timeout=1.0,
+            )
+
+        assert harness.server._pending == {}, "send-time exception must not leak _pending entries"
+        await plugin.close()
+
+
 # --- Issue #262: editor_state self-heals a stale "playing" cache ---
 
 

tests/unit/test_gdscript_no_adjacent_string_concat.py

@@ -151,11 +151,7 @@ def test_no_python_style_adjacent_string_concat_in_gdscript() -> None:
 
 def test_detector_flags_canonical_multiline_bug_pattern() -> None:
     """The exact shape from PR #236 must be detected."""
-    src = (
-        "assert_false(cond,\n"
-        '    "first half - "\n'
-        '    "second half")\n'
-    )
+    src = 'assert_false(cond,\n    "first half - "\n    "second half")\n'
     hits = _find_adjacent_string_pairs(src)
     assert len(hits) == 1, f"expected 1 hit, got {hits}"
     (prev_line, _), (cur_line, _) = hits[0]