Skip to content

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

Merged
jonpspri merged 1 commit into
mainfrom
fix/re-register-over-fresh-connection
Aug 13, 2026
Merged

fix: re-register over a fresh gateway connection during session recovery#5
jonpspri merged 1 commit into
mainfrom
fix/re-register-over-fresh-connection

Conversation

@jonpspri

Copy link
Copy Markdown
Collaborator

Problem

ContextForge gateways on the current reverse-proxy lifecycle contract (the Phase 3 service work on IBM/mcp-context-forge) accept exactly one register frame per WebSocket connection: a duplicate register is answered with an error frame (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:

  1. Tool-call session expiry (_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.
  2. MCP health recovery (_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).
  • The supervisor loop in 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 — clean
  • uv run mypy src/mcp_reverse_proxy — clean

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>
@jonpspri
jonpspri merged commit 05c1cb3 into main Aug 13, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant