Add disconnect hooks for MCP websocket servers

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`,

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<Arc<dyn BeforeRequest>>,
     pub(crate) after: Vec<Arc<dyn AfterRequest>>,
     pub(crate) on_connect: Vec<Arc<dyn OnConnect>>,
+    pub(crate) on_disconnect: Vec<Arc<dyn OnDisconnect>>,
 }
 
 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()
     }
 }

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<Option<Vec<String>>>,
     #[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<Arc<dyn WsUpgrader>> {
         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");
+    }
+}

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<WsRequestValue, _> = match msg {
             Message::Text(t) => {