Uh oh!
There was an error while loading. Please reload this page.
Forward backend RMCP notifications - #23
Conversation
1f47355 to
28ecae4Compare4988829 to
75c49a9CompareThere was a problem hiding this comment.
Right, it is a big and complex code refactor and this PR addresses two things at the same time:
- handling of sse events
- 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:
- we already have a layer to handle session_ids so not sure why a new layer is necessary
- 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.
- on delete, i think all
75c49a9 to
25f52a8CompareSplit is done per review.
|
e5dbcb3 to
7db1d3cCompare25f52a8 to
f8b47a0Comparef8b47a0 to
90e6511Comparefab520d to
0515f35Compare6e353a1 to
308a4e6Comparec145a6d to
7cc3e18CompareServe each backend MCP connection through a GatewayBackendClient that holds the downstream peer and an in-flight tool-call registry keyed by progress token. While call_tool is in flight, backend progress notifications are routed to the matching call by token, run through the tool post hooks, and forwarded to the downstream client; notifications with unknown tokens are dropped, as MCP requires backends to echo the caller's token. The downstream token is kept on the backend request via PeerRequestOptions meta, which Peer::call_tool would otherwise overwrite with an auto-generated token at serialization. Signed-off-by: lucarlig <luca.carlig@ibm.com>
Route every backend tool call through send_cancellable_request and race the response against the downstream request's cancellation token. When the downstream client cancels, cancel the upstream backend request and return a cancelled error instead of leaving it running. Signed-off-by: lucarlig <luca.carlig@ibm.com>
7cc3e18 to
d24f8ebCompareProgress-tracking regressions (review note — deferred)Two regressions in 1. Shared progress-token eviction
Fix direction: only remove the entry the completing call actually owns (e.g. an 2. Cleanup not robust on dropped requestRegistration is cleaned up by a manual This one also can't be expressed as a deterministic black-box test: every reachable path already calls |
e3e1075 to
d24f8ebCompareUh oh!
There was an error while loading. Please reload this page.
🚀 What This PR Does
This PR forwards backend MCP
notifications/progressto the downstream client whilecall_toolis still in flight, and keeps downstream cancellation wired through to the backend request.🧭 End-to-End Flow
✨ Key Changes
GatewayBackendClientowns the downstream peer and an in-flight tool-call registry keyed by progress token.on_progressroutes each backend progress notification to the matching in-flight call, runs it through the tool post hooks, and forwards it downstream.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.PeerRequestOptionsmeta so the backend receives the client's progress token (Peer::call_toolwould overwrite it with an auto-generated token at serialization).tool_call_idand context table); a plugin deny drops the notification instead of failing the call.🔒 Guardrails
✅ Verification
cargo test --workspace --all-targetscargo clippy --workspace --all-targetsINVALID_PARAMStool_call_idacross progress and final result hooks🔬 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.progressTokenand 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 onmainincludes the fix. The gateway routes progress strictly by the downstream token and drops unknown tokens.Test Shape
servers_progress_demoStreamable HTTP example.Client request sequence:
initializenotifications/initializedtools/listtools/callwith a client progress tokenThe backend tool was exposed through the gateway namespace as:
Observed Stream
The tool-call response stream contained all backend progress notifications before the final tool result:
notifications/progress progress=1 message=11111notifications/progress progress=2 message=22222notifications/progress progress=3 message=33333notifications/progress progress=4 message=44444notifications/progress progress=5 message=55555Processed 5 records successfully