Forward backend RMCP notifications

[lucarlig]

🚀 What This PR Does

This PR forwards backend MCP notifications/progress to the downstream client while call_tool is still in flight, and keeps downstream cancellation wired through to the backend request.

Rebased onto main after MCP logging was removed; logging notification forwarding is dropped and progress tracking is a plain token-indexed registry.

🧭 End-to-End Flow

sequenceDiagram
    participant Client as Downstream MCP Client
    participant Gateway as ContextForge Gateway
    participant Backend as Backend MCP Server

    Client->>Gateway: tools/call + progress token
    Gateway->>Backend: tools/call + same progress token
    Backend-->>Gateway: notifications/progress
    Gateway-->>Client: notifications/progress
    Backend-->>Gateway: final tool result
    Gateway-->>Client: final tool result

Loading

✨ Key Changes

Area	Change
Backend client handler	GatewayBackendClient owns the downstream peer and an in-flight tool-call registry keyed by progress token.
Progress forwarding	on_progress routes each backend progress notification to the matching in-flight call, runs it through the tool post hooks, and forwards it downstream.
Strict token routing	The registry is a HashMap<ProgressToken, _>; notifications with unknown tokens are dropped, as MCP requires backends to echo the caller's token. Concurrent tool calls to the same backend run in parallel.
Progress token propagation	Tool calls use PeerRequestOptions meta so the backend receives the client's progress token (Peer::call_tool would overwrite it with an auto-generated token at serialization).
Plugin hook coverage	Progress notifications and the final result share the same per-call plugin state (tool_call_id and context table); a plugin deny drops the notification instead of failing the call.
Cancellation	Downstream cancellation cancels the upstream backend request and returns a cancelled error.

🔒 Guardrails

• Progress is only forwarded while a tool call is registered as in flight.
• Progress is routed strictly by the downstream progress token; unknown tokens are dropped per spec.
• Calls without a downstream progress token are not tracked (no progress can be routed to them).
• Plugin denial drops that notification without failing the whole tool call.

✅ Verification

Check	Result
cargo test --workspace --all-targets	✅ Passed, 46 tests
cargo clippy --workspace --all-targets	✅ 0 warnings
cpex progress-hook unit tests	✅ Pass-through without state, rewrite round-trip, deny drop, invalid rewrite rejected with INVALID_PARAMS
E2E forwarding tests	✅ Progress forwarded in SSE and JSON response modes
Concurrency tests	✅ Concurrent RMCP-client calls partition progress per token, concurrent raw streamable HTTP calls partition per-POST SSE streams with no foreign tokens, and backend-generated progress tokens are dropped
Plugin hook tests	✅ Shared tool_call_id across progress and final result hooks
Progress deny test	✅ Progress suppressed while the tool call still succeeds
Cancellation relay test	✅ Backend observes cancellation
Repeated concurrency/deny/cancellation tests	✅ No timing flake observed

🔬 Additional E2E Evidence

Verified against the upstream RMCP progress-demo server, not only against gateway test fixtures.

Note: older revisions of the RMCP progress demo ignored the caller's _meta.progressToken and minted their own counter tokens, which breaks token-based routing. Upstream fixed the demo to require and echo the caller's token (modelcontextprotocol/rust-sdk#898), and the rmcp revision pinned on main includes the fix. The gateway routes progress strictly by the downstream token and drops unknown tokens.

Test Shape

• Backend: RMCP servers_progress_demo Streamable HTTP example.
• Gateway config: one virtual host backend pointing at that RMCP server.
• Client flow: normal downstream MCP calls through the gateway.

Client request sequence:

1. initialize
2. notifications/initialized
3. tools/list
4. tools/call with a client progress token

The backend tool was exposed through the gateway namespace as:

gateway-one-stream_processor

Observed Stream

The tool-call response stream contained all backend progress notifications before the final tool result:

Event	Forwarded payload
1	notifications/progress progress=1 message=11111
2	notifications/progress progress=2 message=22222
3	notifications/progress progress=3 message=33333
4	notifications/progress progress=4 message=44444
5	notifications/progress progress=5 message=55555
Final	Processed 5 records successfully

[dawid-nowak]

Right, it is a big and complex code refactor and this PR addresses two things at the same time:

1. handling of sse events
2. handling clean up of sessions on delete
   which makes it even harder to review.

I would like to see a separate PR for session deletion. Let's get that first while looking at the SSE stuff and then we can create follow up PRs to address the SSE events.

For session deletion, I like the approach but couple of things I am uncertain of:

1. we already have a layer to handle session_ids so not sure why a new layer is necessary
2. at the moment we have a session store in MCP_gateway, ideally this store should be reused between mcp_gatway and layers/session_id in a uniform way. So for example we should use Axum native mechanism in mcp_gateway as well as in SessionOwnerState, if possible.
3. on delete, i think all

[lucarlig]

Split is done per review.

• New first PR: Clean up MCP sessions on delete #27, Clean up MCP sessions on delete, based on latest main. It contains the session owner/delete cleanup and the session-delete integration test.
• This PR is now stacked on Clean up MCP sessions on delete #27 and only carries the SSE/RMCP notification forwarding work.

done

[lucarlig]

Progress-tracking regressions (review note — deferred)

Two regressions in backend_client.rs surfaced during review. No regression test is committed for either, and we're intentionally ignoring both for the moment — capturing them here so they aren't lost. Fix and tests to follow later.

1. Shared progress-token eviction

stop_tracking_tool_call removes the registry entry unconditionally. MCP requires progress tokens to be unique across active requests, but a malformed client can reuse one. The first call to register a token owns it; a second call reusing the same token does not displace it. When that second call completes, the unconditional remove wipes the still-in-flight owner's registration, and the owner's remaining progress notifications are silently dropped.

Fix direction: only remove the entry the completing call actually owns (e.g. an Arc::ptr_eq check, naturally expressed via an RAII registration guard).

2. Cleanup not robust on dropped request

Registration is cleaned up by a manual stop_tracking_tool_call call rather than RAII. It runs on success, backend error, and graceful MCP cancellation — but not if the call_tool future is dropped mid-flight (e.g. a downstream connection abort), in which case the in-flight entry leaks.

This one also can't be expressed as a deterministic black-box test: every reachable path already calls stop_tracking_tool_call, the leak isn't observable through the public API, and a unit test is blocked because track_tool_call/InFlightToolCall need an rmcp Peer<RoleServer> whose constructor is pub(crate). A test only becomes feasible alongside the fix that makes registration RAII and keeps the per-call record Peer-free.