Skip to content

fix(rmcp): drive the MCP handshake once at worker startup #110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
harsh04044 wants to merge 1 commit into
base: main
Choose a base branch
from
+586 −117
Open

fix(rmcp): drive the MCP handshake once at worker startup #110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,17 @@

### Fixed

- rmcp server worker: a single inbound message that was not an `initialize` request could
permanently stop the server for every client. All Nostr clients are multiplexed through one
rmcp service, and rmcp's pre-service handshake accepts only an `initialize` request as its
first message; anything else (a `notifications/initialized`, any other typed notification, a
response, or an error response) terminated the service, which cancelled the worker and closed
the transport. Nothing recovered it: every later client hung. Reaching it needed no
authentication, no allowlisted key, and no honest client, and an honest client could trigger it
on itself whenever the relay delivered its `notifications/initialized` before its `initialize`,
which Nostr does not order. The worker now satisfies the handshake itself at startup, before it
drains any relay-sourced message, so no inbound message can reach the pre-service handshake.

- CEP-41 open-stream: a deferred final response (one held back while a stream was
still open) went out with routing tags only, so it carried neither the CEP-35
discovery tags nor the CEP-8 effective-mode disclosure. A client whose first server
Expand All @@ -39,6 +50,17 @@

### Changed

- rmcp server worker: every client now receives the handler's declared `protocolVersion` in its
`InitializeResult`. Because the worker satisfies the pre-service handshake itself, no client's
`initialize` runs rmcp's version negotiation any more, so the first client is answered with the
declared version rather than an echo of the version it requested. Clients after the first
already behaved this way, so the change makes all clients uniform, in the less conformant
direction: the multiplexed adapter no longer negotiates the protocol version per client. A
server announcement is unaffected for a handler that does not override the version, and
reflects the declared version for one that does. The per-client synthetic bootstrap that used
to inject an `initialize` ahead of a new client's first request is removed, since the handshake
is now satisfied before any client is served.

- **Breaking:** `ClientSession` and `SessionSnapshot` are now `#[non_exhaustive]`. Both
gain fields as new CEPs land (CEP-8 adds three to each), so downstream code must
construct `ClientSession` via `ClientSession::new` and destructure either struct with
Expand Down
4 changes: 4 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,10 @@ required-features = ["rmcp", "test-utils"]
name = "payments_negotiation_e2e"
required-features = ["test-utils"]

[[test]]
name = "rmcp_handshake_survival"
required-features = ["rmcp", "test-utils"]

[dev-dependencies]
# `start_paused` deterministic-time tests (tokio's "full" does not include test-util).
tokio = { version = "1", features = ["test-util"] }
Expand Down
166 changes: 49 additions & 117 deletions src/rmcp_transport/worker.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ use crate::transport::server::{
};
use rmcp::model::GetExtensions;
use rmcp::transport::worker::{Worker, WorkerContext, WorkerQuitReason};
use std::collections::HashSet;

use super::convert::{
internal_to_rmcp_client_rx, internal_to_rmcp_server_rx, rmcp_client_tx_to_internal,
Expand Down Expand Up @@ -46,27 +45,6 @@ fn synthetic_initialized_notification() -> JsonRpcMessage {
})
}

fn should_inject_stateless_bootstrap(
initialized_clients: &HashSet<String>,
client_pubkey: &str,
message: &JsonRpcMessage,
) -> bool {
if initialized_clients.contains(client_pubkey) {
return false;
}

matches!(message, JsonRpcMessage::Request(req) if req.method != "initialize")
}

fn is_synthetic_initialize_message(message: &JsonRpcMessage) -> bool {
matches!(
message,
JsonRpcMessage::Request(req)
if req.method == "initialize"
&& req.id == serde_json::json!(STATELESS_SYNTHETIC_EVENT_ID)
)
}

/// rmcp server worker wrapper for ContextVM Nostr server transport.
///
/// Multiplexes all connected clients through a single rmcp service instance.
Expand Down Expand Up @@ -133,8 +111,48 @@ impl Worker for NostrServerWorker {
)
})?;

// Satisfy rmcp's pre-init handshake worker-locally, before any relay-sourced
// message can be drained. The handler channel has one producer (this task) and
// is FIFO, so a synthetic initialize + initialized sent here reaches the service
// ahead of every inbound event, for every peer and ordering.
let startup: std::result::Result<(), WorkerQuitReason<Self::Error>> = async {
let init = synthetic_initialize_message();
let rmcp_init = internal_to_rmcp_server_rx(&init).ok_or_else(|| {
WorkerQuitReason::fatal(
Self::Error::Validation(
"failed converting synthetic initialize request to rmcp format".to_string(),
),
"converting synthetic initialize request",
)
})?;
context.send_to_handler(rmcp_init).await?;

let initialized = synthetic_initialized_notification();
let rmcp_initialized = internal_to_rmcp_server_rx(&initialized).ok_or_else(|| {
WorkerQuitReason::fatal(
Self::Error::Validation(
"failed converting synthetic initialized notification to rmcp format"
.to_string(),
),
"converting synthetic initialized notification",
)
})?;
context.send_to_handler(rmcp_initialized).await
}
.await;
// Share the loop's close-on-exit path: a send failure here means the
// WorkerTransport was already dropped (e.g. a serve()-vs-timeout race). Without
// this close, the transport's own tasks (event loop, cleanup, relay
// subscription) leak for the process lifetime, because nothing else runs between
// `take_message_receiver()` and the loop.
if let Err(reason) = startup {
if let Err(e) = self.transport.close().await {
tracing::warn!(target: LOG_TARGET, error = %e, "Failed to close server transport cleanly");
}
return Err(reason);
}

let cancellation_token = context.cancellation_token.clone();
let mut initialized_clients = HashSet::new();

let quit_reason = loop {
tokio::select! {
Expand All @@ -154,58 +172,12 @@ impl Worker for NostrServerWorker {
..
} = incoming;

let should_inject_bootstrap = should_inject_stateless_bootstrap(
&initialized_clients,
&client_pubkey,
&message,
);

if should_inject_bootstrap {
let synthetic_init = synthetic_initialize_message();
let Some(rmcp_init) = internal_to_rmcp_server_rx(&synthetic_init) else {
break WorkerQuitReason::fatal(
Self::Error::Validation(
"failed converting synthetic initialize request to rmcp format".to_string(),
),
"converting synthetic initialize request",
);
};

if let Err(reason) = context.send_to_handler(rmcp_init).await {
break reason;
}

let initialized = synthetic_initialized_notification();
let Some(rmcp_initialized) = internal_to_rmcp_server_rx(&initialized) else {
break WorkerQuitReason::fatal(
Self::Error::Validation(
"failed converting synthetic initialized notification to rmcp format".to_string(),
),
"converting synthetic initialized notification",
);
};

if let Err(reason) = context.send_to_handler(rmcp_initialized).await {
break reason;
}

initialized_clients.insert(client_pubkey.clone());
}

if matches!(&message, JsonRpcMessage::Request(req) if req.method == "initialize")
|| matches!(&message, JsonRpcMessage::Notification(n) if n.method == "notifications/initialized")
{
initialized_clients.insert(client_pubkey.clone());
}

// Rewrite real wire requests to the Nostr event_id.
// Synthetic stateless bootstrap messages must retain their
// sentinel ID so their responses can be dropped before they
// ever touch transport correlation.
if !is_synthetic_initialize_message(&message) {
if let JsonRpcMessage::Request(ref mut req) = message {
// Rewrite wire requests to the Nostr event_id. Every message
// reaching this loop came off the relay; the worker's own
// synthetic handshake pair goes straight to the handler at
// startup and never traverses it.
if let JsonRpcMessage::Request(ref mut req) = message {
req.id = serde_json::json!(event_id);
}
}

if let Some(mut rmcp_msg) = internal_to_rmcp_server_rx(&message) {
Expand Down Expand Up @@ -484,40 +456,6 @@ mod tests {
use super::*;
use crate::core::types::JsonRpcResponse;

#[test]
fn test_should_inject_stateless_bootstrap_for_first_non_initialize_request() {
let initialized_clients = HashSet::new();
let message = JsonRpcMessage::Request(JsonRpcRequest {
jsonrpc: "2.0".to_string(),
id: serde_json::json!(1),
method: "tools/list".to_string(),
params: Some(serde_json::json!({})),
});

assert!(should_inject_stateless_bootstrap(
&initialized_clients,
"client-a",
&message,
));
}

#[test]
fn test_should_not_inject_stateless_bootstrap_for_real_initialize() {
let initialized_clients = HashSet::new();
let message = JsonRpcMessage::Request(JsonRpcRequest {
jsonrpc: "2.0".to_string(),
id: serde_json::json!(1),
method: "initialize".to_string(),
params: Some(serde_json::json!({})),
});

assert!(!should_inject_stateless_bootstrap(
&initialized_clients,
"client-a",
&message,
));
}

#[test]
fn test_synthetic_initialize_keeps_sentinel_id() {
let message = synthetic_initialize_message();
Expand Down Expand Up @@ -579,13 +517,6 @@ mod tests {
}
}

#[test]
fn test_is_synthetic_initialize_message_detects_sentinel() {
assert!(is_synthetic_initialize_message(
&synthetic_initialize_message()
));
}

#[test]
fn test_announcement_sentinel_differs_from_stateless_sentinel() {
assert_ne!(ANNOUNCEMENT_REQUEST_ID, STATELESS_SYNTHETIC_EVENT_ID);
Expand All @@ -606,8 +537,9 @@ mod tests {
if let JsonRpcMessage::Response(ref resp) = response {
let event_id = resp.id.as_str().unwrap();
assert_eq!(event_id, ANNOUNCEMENT_REQUEST_ID);
// Must not be confused with the stateless synthetic sentinel
assert!(!is_synthetic_initialize_message(&response));
// Must not be confused with the stateless synthetic sentinel, whose
// response takes the drop path in `forward_server_internal`.
assert_ne!(event_id, STATELESS_SYNTHETIC_EVENT_ID);
}
}
}
Loading
Loading