fix: re-register over a fresh gateway connection during session recovery - #5
Merged
Merged
Conversation
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 <jps@s390x.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
ContextForge gateways on the current reverse-proxy lifecycle contract (the Phase 3 service work on
IBM/mcp-context-forge) accept exactly oneregisterframe per WebSocket connection: a duplicate register is answered with anerrorframe (connection already registered) and a policy-violation close (1008). Register is a one-shot lifecycle event per connection.Two client recovery paths re-register in place on the existing connection:
_handle_gateway_message): when the downstream MCP server restarts and a forwarded request fails with a stale session, the client stores the request and calls_register()on the same socket._keepalive_loop): when a previously unhealthy MCP server recovers, the client re-registers in place to trigger a fresh initialize sequence.Against a current gateway both paths get the duplicate-register error and a closed socket. Recovery still converges via the outer supervisor reconnect loop and no stored request is lost, but every downstream restart costs a full TCP/TLS/WebSocket reconnect plus complete catalog re-discovery anyway — after a wasted, noisy rejected register round-trip (
Gateway error: connection already registered).Fix
Re-register over a fresh gateway connection instead of in place:
_recover_gateway_connection()— disconnect → connect → register, with concurrent triggers collapsed behind a re-entrancy guard._schedule_gateway_recovery()— background-task wrapper for the receive-handler context: gateway message handlers run inside the transport's receive task, which the reconnect cancels and replaces, so the recovery cannot be awaited there (the existing register-failure path already uses this schedule-don't-await pattern).run_with_reconnect()skips its gateway-lost check while a recovery reconnect is in flight, so it cannot race a double-connect.disconnect()cancels any in-flight recovery task.The pending-request retry after
register_complete(success)is unchanged — instance state survives the reconnect, so a stored tool call is still retried once the new session is initialized.Compatibility
A fresh-connection register is accepted by both legacy and current gateways, so this works against either generation. No wire-protocol changes.
Verification
uv run pytest tests/ -q— 133 passed (6 new/updated tests cover: reconnect-before-register even when connected, concurrent-trigger collapse, scheduled-task coalescing, supervisor no-double-connect, disconnect cancels in-flight recovery, pending-request storage through scheduled recovery)uv run ruff check src tests— cleanuv run mypy src/mcp_reverse_proxy— clean