fix: re-register over a fresh gateway connection during session recovery

src/mcp_reverse_proxy/client.py

@@ -107,6 +107,8 @@ def __init__(
 
         self._keepalive_task: asyncio.Task[None] | None = None
         self._pending_requests: dict[Any, asyncio.Future[Any]] = {}
+        self._gateway_recovering = False
+        self._gateway_recovery_task: asyncio.Task[None] | None = None
 
         # Register message handlers
         self.mcp_transport.add_message_handler(self._handle_mcp_message)
@@ -174,6 +176,11 @@ async def disconnect(self) -> None:
             with suppress(asyncio.CancelledError):
                 await self._keepalive_task
 
+        if self._gateway_recovery_task:
+            self._gateway_recovery_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._gateway_recovery_task
+
         # Send unregister message
         if await self.gateway_transport.is_connected():
             try:
@@ -216,7 +223,7 @@ async def run_with_reconnect(self) -> None:
                             break
 
                     # Check if gateway is still connected
-                    if not await self.gateway_transport.is_connected():
+                    if not self._gateway_recovering and not await self.gateway_transport.is_connected():
                         LOGGER.warning("Gateway connection lost, triggering reconnection")
                         self.state = ConnectionState.DISCONNECTED
                         break
@@ -295,6 +302,48 @@ async def _register(self) -> None:
         await self.gateway_transport.send(orjson.dumps(register_msg).decode())
         LOGGER.debug(f"Registration message sent: {register_msg}")
 
+    async def _recover_gateway_connection(self) -> None:
+        """Reconnect the gateway transport and re-register on the fresh connection.
+
+        Gateways on the current ContextForge lifecycle contract accept exactly
+        one register frame per WebSocket connection; a duplicate register is
+        answered with an error frame and a policy-violation close. Session
+        recovery therefore re-establishes the connection before re-registering,
+        which stays compatible with both legacy and current gateways.
+        Concurrent triggers collapse into a single reconnect.
+        """
+        if self._gateway_recovering:
+            LOGGER.info("[GATEWAY_RECOVERY] Recovery already in progress, skipping duplicate trigger")
+            return
+        self._gateway_recovering = True
+        try:
+            await self.gateway_transport.disconnect()
+            await self.gateway_transport.connect()
+            self._registration_successful = False
+            await self._register()
+        finally:
+            self._gateway_recovering = False
+
+    def _schedule_gateway_recovery(self) -> None:
+        """Schedule gateway connection recovery as a background task.
+
+        Gateway message handlers run inside the gateway transport's receive
+        task, which the reconnect cancels and replaces, so the recovery cannot
+        be awaited from that context. A failure is logged and left to the
+        supervisor loop in run_with_reconnect(), which performs a full
+        reconnect on its next poll.
+        """
+        if self._gateway_recovery_task is not None and not self._gateway_recovery_task.done():
+            return
+
+        async def _recover() -> None:
+            try:
+                await self._recover_gateway_connection()
+            except Exception as e:
+                LOGGER.error(f"[GATEWAY_RECOVERY] Gateway reconnect failed: {e}")
+
+        self._gateway_recovery_task = asyncio.create_task(_recover())
+
     async def _handle_mcp_message(self, message: str) -> None:
         """Handle message from MCP server."""
         try:
@@ -419,10 +468,15 @@ async def _handle_gateway_message(self, message: str) -> None:
                         "authType": auth_type,
                     }
 
-                    LOGGER.info("[REVERSE_PROXY_CLIENT] Triggering re-registration with gateway...")
+                    LOGGER.info("[REVERSE_PROXY_CLIENT] Triggering gateway reconnect for re-registration...")
                     self._registration_successful = False
-                    await self._register()
-                    LOGGER.info("[REVERSE_PROXY_CLIENT] Re-registration triggered, returning from handler")
+                    # Re-register over a fresh connection: current gateways
+                    # accept one register frame per connection and answer a
+                    # duplicate with an error frame and a policy close. The
+                    # reconnect is scheduled, not awaited, because this handler
+                    # runs inside the receive task the reconnect must replace.
+                    self._schedule_gateway_recovery()
+                    LOGGER.info("[REVERSE_PROXY_CLIENT] Re-registration scheduled, returning from handler")
                     return
 
             elif msg_type == MessageType.HEARTBEAT.value:
@@ -650,20 +704,13 @@ async def _keepalive_loop(self) -> None:
                     self._mcp_server_healthy = True
                     self._consecutive_mcp_failures = 0
 
-                    # Always re-register when MCP recovers to trigger new initialization
-                    if not await self.gateway_transport.is_connected():
-                        LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Gateway disconnected, reconnecting before re-registration")
-                        try:
-                            await self.gateway_transport.connect()
-                        except Exception as e:
-                            LOGGER.error(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Failed to reconnect to gateway: {e}")
-                            continue
-
-                    # Re-register to trigger new MCP initialization sequence
+                    # Re-register over a fresh gateway connection to trigger the
+                    # new initialization sequence: current gateways accept one
+                    # register frame per connection, so an in-place re-register
+                    # would be rejected with an error frame and a policy close.
                     try:
-                        LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Sending re-registration to gateway")
-                        self._registration_successful = False
-                        await self._register()
+                        LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Reconnecting gateway for re-registration")
+                        await self._recover_gateway_connection()
                         LOGGER.info(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Re-registration sent, gateway will initialize MCP server")
                     except Exception as e:
                         LOGGER.error(f"[HEARTBEAT_RECOVERY] Session {self.session_id[:8]}... | Failed to re-register with gateway: {e}")

tests/test_mcp_reverse_proxy_client.py

@@ -244,7 +244,7 @@ async def test_handle_gateway_request_sends_error_when_transport_unavailable_and
 
 @pytest.mark.asyncio
 async def test_handle_gateway_request_stores_pending_and_reregisters_when_health_recovers(proxy_client, monkeypatch) -> None:
-    """Recoverable transport errors should save the request and trigger re-registration."""
+    """Recoverable transport errors should save the request and re-register over a fresh gateway connection."""
     proxy_client.mcp_transport.send.side_effect = SessionExpiredError("expired")
     register_mock = AsyncMock()
     monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True))
@@ -260,9 +260,93 @@ async def test_handle_gateway_request_stores_pending_and_reregisters_when_health
         "authType": "bearer",
     }
     assert proxy_client._registration_successful is False
+    assert proxy_client._gateway_recovery_task is not None
+    await proxy_client._gateway_recovery_task
+
+    proxy_client.gateway_transport.disconnect.assert_awaited_once()
+    proxy_client.gateway_transport.connect.assert_awaited_once()
+    register_mock.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_recover_gateway_connection_reconnects_even_when_connected(proxy_client, monkeypatch) -> None:
+    """Recovery must replace a live connection because register is one-shot per connection."""
+    register_mock = AsyncMock()
+    monkeypatch.setattr(proxy_client, "_register", register_mock)
+
+    await proxy_client._recover_gateway_connection()
+
+    proxy_client.gateway_transport.disconnect.assert_awaited_once()
+    proxy_client.gateway_transport.connect.assert_awaited_once()
+    register_mock.assert_awaited_once()
+    assert proxy_client._gateway_recovering is False
+
+
+@pytest.mark.asyncio
+async def test_recover_gateway_connection_collapses_concurrent_triggers(proxy_client, monkeypatch) -> None:
+    """A second recovery trigger while one runs must not reconnect twice."""
+    register_mock = AsyncMock()
+    monkeypatch.setattr(proxy_client, "_register", register_mock)
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    async def _blocking_disconnect() -> None:
+        started.set()
+        await release.wait()
+
+    proxy_client.gateway_transport.disconnect.side_effect = _blocking_disconnect
+
+    first = asyncio.create_task(proxy_client._recover_gateway_connection())
+    await started.wait()
+    await proxy_client._recover_gateway_connection()
+    release.set()
+    await first
+
+    proxy_client.gateway_transport.disconnect.assert_awaited_once()
+    proxy_client.gateway_transport.connect.assert_awaited_once()
     register_mock.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_schedule_gateway_recovery_coalesces_into_single_task(proxy_client, monkeypatch) -> None:
+    """Repeated scheduled recoveries should share one background reconnect."""
+    register_mock = AsyncMock()
+    monkeypatch.setattr(proxy_client, "_register", register_mock)
+
+    proxy_client._schedule_gateway_recovery()
+    proxy_client._schedule_gateway_recovery()
+    task = proxy_client._gateway_recovery_task
+    assert task is not None
+    await task
+
+    proxy_client.gateway_transport.disconnect.assert_awaited_once()
+    proxy_client.gateway_transport.connect.assert_awaited_once()
+    register_mock.assert_awaited_once()
+    assert proxy_client._gateway_recovery_task is not None
+    assert proxy_client._gateway_recovery_task.done()
+
+
+@pytest.mark.asyncio
+async def test_disconnect_cancels_pending_gateway_recovery(proxy_client) -> None:
+    """A full disconnect should cancel an in-flight recovery reconnect."""
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    async def _blocking_disconnect() -> None:
+        started.set()
+        await release.wait()
+
+    proxy_client.gateway_transport.disconnect.side_effect = _blocking_disconnect
+    proxy_client._schedule_gateway_recovery()
+    await started.wait()
+    proxy_client.gateway_transport.disconnect.side_effect = None
+
+    await proxy_client.disconnect()
+
+    assert proxy_client._gateway_recovery_task is not None
+    assert proxy_client._gateway_recovery_task.cancelled()
+
+
 @pytest.mark.asyncio
 async def test_handle_gateway_request_reraises_non_connection_runtime_error(proxy_client, monkeypatch) -> None:
     """Non-connection runtime errors should be re-raised then logged by the outer handler."""
@@ -495,6 +579,30 @@ async def _fake_connect() -> None:
     assert proxy_client.retry_count == 1
     assert call_count == 2
 
+@pytest.mark.asyncio
+async def test_run_with_reconnect_ignores_gateway_drop_during_recovery(proxy_client, monkeypatch) -> None:
+    """The supervisor loop must not double-connect while a recovery reconnect runs."""
+    call_count = 0
+
+    async def _fake_connect() -> None:
+        nonlocal call_count
+        call_count += 1
+        proxy_client.state = ConnectionState.CONNECTED
+        proxy_client._keepalive_task = None
+        proxy_client._gateway_recovering = True
+        proxy_client.gateway_transport.is_connected.return_value = False
+
+    async def _fake_sleep(_delay: float) -> None:
+        proxy_client.state = ConnectionState.SHUTTING_DOWN
+
+    monkeypatch.setattr(proxy_client, "connect", _fake_connect)
+    monkeypatch.setattr(proxy_client, "_check_mcp_server_health", AsyncMock(return_value=True))
+    monkeypatch.setattr(client_mod.asyncio, "sleep", _fake_sleep)
+
+    await proxy_client.run_with_reconnect()
+
+    assert call_count == 1
+
 
 @pytest.mark.asyncio
 async def test_run_with_reconnect_marks_mcp_unhealthy_when_transport_disconnects(proxy_client, monkeypatch) -> None:
@@ -719,7 +827,7 @@ async def _fake_sleep(_delay: float) -> None:
 
 @pytest.mark.asyncio
 async def test_keepalive_loop_recovers_and_reregisters_when_mcp_returns(proxy_client, monkeypatch) -> None:
-    """Recovered MCP server should reconnect gateway if needed and re-register."""
+    """Recovered MCP server should reconnect the gateway and re-register on the fresh connection."""
     proxy_client.state = ConnectionState.CONNECTED
     proxy_client.keepalive_interval = 0
     proxy_client._mcp_server_healthy = False
@@ -739,6 +847,7 @@ async def _fake_sleep(_delay: float) -> None:
 
     proxy_client.gateway_transport.connect.assert_awaited_once()
     register_mock.assert_awaited_once()
+    proxy_client.gateway_transport.disconnect.assert_awaited_once()
     assert proxy_client._mcp_server_healthy is True
     assert proxy_client._consecutive_mcp_failures == 0
     proxy_client.gateway_transport.send.assert_awaited_once()