From 370fa6871edbe7290ee448d806de58da27c2af29 Mon Sep 17 00:00:00 2001 From: meh Date: Sun, 9 Aug 2026 07:15:48 +0700 Subject: [PATCH] feat: support MCP-only websocket servers --- README.md | 11 ++++++--- src/libs/ws/hooks.rs | 20 +++++++++++++++- src/libs/ws/server.rs | 52 +++++++++++++++++++++++++++++++++++++++++- src/libs/ws/session.rs | 15 ++++++++++++ 4 files changed, 93 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 4ecf916..6b7620d 100644 --- a/README.md +++ b/README.md @@ -205,7 +205,9 @@ The WebSocket server can optionally expose every registered endpoint as an **MCP tool** over JSON-RPC 2.0, alongside the legacy `{method, seq, params}` protocol. MCP is **off by default** and fully additive: with it disabled the server behaves exactly as before, and even with it enabled, legacy frames are -routed unchanged — both protocols work on the same connection. +routed unchanged, so both protocols work on the same connection. Set +`WsServerConfig::mcp_only` to `true` to reject non-MCP frames. MCP-only mode +requires `enable_mcp` and fails validation at startup without it. Supported MCP methods: `initialize`, `ping`, `tools/list` (filtered by the connection's roles), `tools/call`, and `notifications/*`. Tool metadata @@ -428,12 +430,15 @@ defines the vocabulary; the platform implementations live in a sibling crate. ### Hooks `BeforeRequest` (may reject and may attach claims to `ctx.extensions`), `AfterRequest` -(observes outcomes), and `OnConnect` (refuses a peer once, rather than per request). -All three run on both the legacy and MCP dispatch paths. +(observes outcomes), `OnConnect` (refuses a peer once, rather than per request), and +`OnDisconnect` (cleans up connection-owned work after the peer leaves). Request hooks +run on both the legacy and MCP dispatch paths unless MCP-only mode rejects legacy +frames first. ```rust server.add_before_hook(MyMissionTokenCheck); server.add_on_connect_hook(RefuseUnattestedPeers); +server.add_on_disconnect_hook(CleanUpConnectionWork); ``` > **Note:** `MessageStream`'s futures are not `Send`, so `serve_connection`, diff --git a/src/libs/ws/hooks.rs b/src/libs/ws/hooks.rs index 369ca88..b01303d 100644 --- a/src/libs/ws/hooks.rs +++ b/src/libs/ws/hooks.rs @@ -26,6 +26,7 @@ use serde_json::Value; use crate::libs::peer::{Extensions, PeerIdentity}; use crate::libs::toolbox::{CustomError, RequestContext}; +use crate::libs::ws::ConnectionId; use crate::model::EndpointSchema; /// How a request finished. Passed to [`AfterRequest`]. @@ -86,17 +87,27 @@ pub trait OnConnect: Send + Sync { ) -> Result<(), CustomError>; } +/// Runs once after a connection's session loop ends. +#[async_trait(?Send)] +pub trait OnDisconnect: Send + Sync { + async fn on_disconnect(&self, connection_id: ConnectionId, peer: &PeerIdentity); +} + /// The registered hooks, snapshotted into each spawned dispatch task. #[derive(Clone, Default)] pub struct Hooks { pub(crate) before: Vec>, pub(crate) after: Vec>, pub(crate) on_connect: Vec>, + pub(crate) on_disconnect: Vec>, } impl Hooks { pub fn is_empty(&self) -> bool { - self.before.is_empty() && self.after.is_empty() && self.on_connect.is_empty() + self.before.is_empty() + && self.after.is_empty() + && self.on_connect.is_empty() + && self.on_disconnect.is_empty() } /// Run every `BeforeRequest` in registration order, stopping at the first error. @@ -135,6 +146,12 @@ impl Hooks { } Ok(()) } + + pub(crate) async fn run_on_disconnect(&self, connection_id: ConnectionId, peer: &PeerIdentity) { + for hook in &self.on_disconnect { + hook.on_disconnect(connection_id, peer).await; + } + } } impl std::fmt::Debug for Hooks { @@ -143,6 +160,7 @@ impl std::fmt::Debug for Hooks { .field("before", &self.before.len()) .field("after", &self.after.len()) .field("on_connect", &self.on_connect.len()) + .field("on_disconnect", &self.on_disconnect.len()) .finish() } } diff --git a/src/libs/ws/server.rs b/src/libs/ws/server.rs index 9d4160c..ce1f2fd 100644 --- a/src/libs/ws/server.rs +++ b/src/libs/ws/server.rs @@ -30,7 +30,8 @@ use crate::libs::ws::mcp::{McpServerInfo, McpState}; use crate::libs::ws::tungstenite::upgrader::create_ws_stream; use crate::libs::ws::{ AfterRequest, BeforeRequest, BoxedStream, ConnectionListener, Hooks, MessageStream, OnConnect, - SessionListener, TcpListener, WsClientSession, WsConnection, WsRequest, WsUpgrader, + OnDisconnect, SessionListener, TcpListener, WsClientSession, WsConnection, WsRequest, + WsUpgrader, }; use crate::model::{EndpointSchema, TypeRegistry}; @@ -88,6 +89,13 @@ impl WebsocketServer { self.mcp = Some(Arc::new(state)); Ok(()) } + + fn validate_protocol_mode(&self) -> Result<()> { + if self.config.mcp_only && self.mcp.is_none() { + bail!("mcp_only requires enable_mcp before serving connections"); + } + Ok(()) + } /// Register a hook that runs before every request, on both the legacy and MCP /// paths. Hooks run in registration order; the first error rejects the request. pub fn add_before_hook(&mut self, hook: impl BeforeRequest + 'static) { @@ -105,6 +113,11 @@ impl WebsocketServer { self.hooks.on_connect.push(Arc::new(hook)); } + /// Register a hook that runs once when a connection's session loop ends. + pub fn add_on_disconnect_hook(&mut self, hook: impl OnDisconnect + 'static) { + self.hooks.on_disconnect.push(Arc::new(hook)); + } + pub fn set_auth_controller(&mut self, controller: impl AuthController + 'static) { self.auth_controller = Arc::new(controller); } @@ -309,6 +322,7 @@ impl WebsocketServer { let addr = conn.peer.display(); let context = RequestContext::from_conn(&conn); let conn_id = context.connection_id; + let disconnect_hooks = self.hooks.clone(); debug!( ws_server = true, @@ -320,6 +334,9 @@ impl WebsocketServer { session.run().await; states.remove(context.connection_id); + disconnect_hooks + .run_on_disconnect(context.connection_id, &context.peer) + .await; info!( ws_server = true, ?addr, @@ -340,6 +357,7 @@ impl WebsocketServer { where L: SessionListener + 'static, { + self.validate_protocol_mode()?; let this = Arc::new(self); let states = Arc::new(WebsocketStates::new()); this.toolbox.set_ws_states( @@ -369,6 +387,7 @@ impl WebsocketServer { } pub async fn listen(self) -> Result<()> { + self.validate_protocol_mode()?; debug!(ws_server = true, "Listening on {}", self.config.address); // Resolve the address and get the socket address @@ -640,6 +659,9 @@ pub struct WsServerConfig { pub allow_cors_urls: Arc>>, #[serde(default = "WsServerConfig::default_server_name")] pub server_name: String, + /// Reject every non-JSON-RPC frame. Requires [`WebsocketServer::enable_mcp`]. + #[serde(default)] + pub mcp_only: bool, } impl Default for WsServerConfig { @@ -656,6 +678,7 @@ impl Default for WsServerConfig { header_only: false, allow_cors_urls: Arc::new(None), server_name: Self::default_server_name(), + mcp_only: false, } } } @@ -679,3 +702,30 @@ fn default_upgrader() -> Option> { None } } + +#[cfg(test)] +mod protocol_mode_tests { + use super::*; + + #[test] + fn mcp_only_requires_the_mcp_router() { + let mut server = WebsocketServer::new(WsServerConfig { + mcp_only: true, + ..Default::default() + }); + assert!(server.validate_protocol_mode().is_err()); + + server + .enable_mcp( + &TypeRegistry::new(), + McpServerInfo { + name: "test".into(), + version: "0".into(), + }, + ) + .expect("empty MCP surface should initialize"); + server + .validate_protocol_mode() + .expect("MCP-only server should validate after enable_mcp"); + } +} diff --git a/src/libs/ws/session.rs b/src/libs/ws/session.rs index 2839c53..24209f0 100644 --- a/src/libs/ws/session.rs +++ b/src/libs/ws/session.rs @@ -78,6 +78,21 @@ impl WsClientSession { } } + if self.server.config.mcp_only { + self.server.toolbox.send_raw( + context.connection_id, + jsonrpc_error( + &None, + JsonRpcError::new( + mcp::INVALID_REQUEST, + "This WebSocket accepts MCP JSON-RPC 2.0 frames only", + ), + ) + .to_string(), + ); + return Ok(true); + } + #[allow(unreachable_patterns)] let obj: Result = match msg { Message::Text(t) => {