From 91633997cc05383b54d4bd1ae95e71fec43720b1 Mon Sep 17 00:00:00 2001 From: Jonathan Springer Date: Thu, 13 Aug 2026 08:40:57 +0100 Subject: [PATCH] fix: re-register over a fresh gateway connection during session recovery ContextForge gateways on the current reverse-proxy lifecycle contract accept exactly one register frame per WebSocket connection: a duplicate register is answered with an error frame and a policy-violation close. Both client recovery paths (tool-call session expiry and MCP health recovery) previously re-registered in place, which against such gateways fails noisily and only recovers via the outer supervisor reconnect. Reconnect the gateway transport before re-registering instead: - add _recover_gateway_connection() (disconnect -> connect -> register, with concurrent triggers collapsed) and _schedule_gateway_recovery() for receive-handler context, which must not await the reconnect that cancels and replaces the receive task - keep the pending-request retry after register_complete unchanged - skip the supervisor loop's gateway-lost check while recovery runs so it cannot double-connect, and cancel any in-flight recovery on disconnect() A fresh-connection register is accepted by both legacy and current gateways, so this stays backward compatible. Signed-off-by: Jonathan Springer --- src/mcp_reverse_proxy/client.py | 81 ++++++++++++++---- tests/test_mcp_reverse_proxy_client.py | 113 ++++++++++++++++++++++++- 2 files changed, 175 insertions(+), 19 deletions(-) diff --git a/src/mcp_reverse_proxy/client.py b/src/mcp_reverse_proxy/client.py index 1e2b0de..99fc437 100644 --- a/src/mcp_reverse_proxy/client.py +++ b/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}") diff --git a/tests/test_mcp_reverse_proxy_client.py b/tests/test_mcp_reverse_proxy_client.py index 410238a..faefea1 100644 --- a/tests/test_mcp_reverse_proxy_client.py +++ b/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()