spikard 0.5.0 → 0.6.1

data/vendor/crates/spikard-http/src/websocket.rs

@@ -1,1375 +1,1404 @@
-//! WebSocket support for Spikard
-//!
-//! Provides WebSocket connection handling with message validation and routing.
-
-use axum::{
-    extract::{
-        State,
-        ws::{Message, WebSocket, WebSocketUpgrade},
-    },
-    response::IntoResponse,
-};
-use serde_json::Value;
-use std::sync::Arc;
-use tracing::{debug, error, info, warn};
-
-fn trace_ws(message: &str) {
-    if std::env::var("SPIKARD_WS_TRACE").ok().as_deref() == Some("1") {
-        eprintln!("[spikard-ws] {message}");
-    }
-}
-
-/// WebSocket message handler trait
-///
-/// Implement this trait to create custom WebSocket message handlers for your application.
-/// The handler processes JSON messages received from WebSocket clients and can optionally
-/// send responses back.
-///
-/// # Implementing the Trait
-///
-/// You must implement the `handle_message` method. The `on_connect` and `on_disconnect`
-/// methods are optional and provide lifecycle hooks.
-///
-/// # Example
-///
-/// ```ignore
-/// use spikard_http::websocket::WebSocketHandler;
-/// use serde_json::{json, Value};
-///
-/// struct EchoHandler;
-///
-/// #[async_trait]
-/// impl WebSocketHandler for EchoHandler {
-///     async fn handle_message(&self, message: Value) -> Option<Value> {
-///         // Echo the message back to the client
-///         Some(message)
-///     }
-///
-///     async fn on_connect(&self) {
-///         println!("Client connected");
-///     }
-///
-///     async fn on_disconnect(&self) {
-///         println!("Client disconnected");
-///     }
-/// }
-/// ```
-pub trait WebSocketHandler: Send + Sync {
-    /// Handle incoming WebSocket message
-    ///
-    /// Called whenever a text message is received from a WebSocket client.
-    /// Messages are automatically parsed as JSON.
-    ///
-    /// # Arguments
-    /// * `message` - JSON value received from the client
-    ///
-    /// # Returns
-    /// * `Some(value)` - JSON value to send back to the client
-    /// * `None` - No response to send
-    fn handle_message(&self, message: Value) -> impl std::future::Future<Output = Option<Value>> + Send;
-
-    /// Called when a client connects to the WebSocket
-    ///
-    /// Optional lifecycle hook invoked when a new WebSocket connection is established.
-    /// Default implementation does nothing.
-    fn on_connect(&self) -> impl std::future::Future<Output = ()> + Send {
-        async {}
-    }
-
-    /// Called when a client disconnects from the WebSocket
-    ///
-    /// Optional lifecycle hook invoked when a WebSocket connection is closed
-    /// (either by the client or due to an error). Default implementation does nothing.
-    fn on_disconnect(&self) -> impl std::future::Future<Output = ()> + Send {
-        async {}
-    }
-}
-
-/// WebSocket state shared across connections
-///
-/// Contains the message handler and optional JSON schemas for validating
-/// incoming and outgoing messages. This state is shared among all connections
-/// to the same WebSocket endpoint.
-#[derive(Debug)]
-pub struct WebSocketState<H: WebSocketHandler> {
-    /// The message handler implementation
-    handler: Arc<H>,
-    /// Optional JSON Schema for validating incoming messages
-    message_schema: Option<Arc<jsonschema::Validator>>,
-    /// Optional JSON Schema for validating outgoing responses
-    response_schema: Option<Arc<jsonschema::Validator>>,
-}
-
-impl<H: WebSocketHandler> Clone for WebSocketState<H> {
-    fn clone(&self) -> Self {
-        Self {
-            handler: Arc::clone(&self.handler),
-            message_schema: self.message_schema.clone(),
-            response_schema: self.response_schema.clone(),
-        }
-    }
-}
-
-impl<H: WebSocketHandler + 'static> WebSocketState<H> {
-    /// Create new WebSocket state with a handler
-    ///
-    /// Creates a new state without message or response validation schemas.
-    /// Messages and responses are not validated.
-    ///
-    /// # Arguments
-    /// * `handler` - The message handler implementation
-    ///
-    /// # Example
-    ///
-    /// ```ignore
-    /// let state = WebSocketState::new(MyHandler);
-    /// ```
-    pub fn new(handler: H) -> Self {
-        Self {
-            handler: Arc::new(handler),
-            message_schema: None,
-            response_schema: None,
-        }
-    }
-
-    /// Create new WebSocket state with a handler and optional validation schemas
-    ///
-    /// Creates a new state with optional JSON schemas for validating incoming messages
-    /// and outgoing responses. If a schema is provided and validation fails, the message
-    /// or response is rejected.
-    ///
-    /// # Arguments
-    /// * `handler` - The message handler implementation
-    /// * `message_schema` - Optional JSON schema for validating client messages
-    /// * `response_schema` - Optional JSON schema for validating handler responses
-    ///
-    /// # Returns
-    /// * `Ok(state)` - Successfully created state
-    /// * `Err(msg)` - Invalid schema provided
-    ///
-    /// # Example
-    ///
-    /// ```ignore
-    /// use serde_json::json;
-    ///
-    /// let message_schema = json!({
-    ///     "type": "object",
-    ///     "properties": {
-    ///         "type": {"type": "string"},
-    ///         "data": {"type": "string"}
-    ///     }
-    /// });
-    ///
-    /// let state = WebSocketState::with_schemas(
-    ///     MyHandler,
-    ///     Some(message_schema),
-    ///     None,
-    /// )?;
-    /// ```
-    pub fn with_schemas(
-        handler: H,
-        message_schema: Option<serde_json::Value>,
-        response_schema: Option<serde_json::Value>,
-    ) -> Result<Self, String> {
-        let message_validator = if let Some(schema) = message_schema {
-            Some(Arc::new(
-                jsonschema::validator_for(&schema).map_err(|e| format!("Invalid message schema: {}", e))?,
-            ))
-        } else {
-            None
-        };
-
-        let response_validator = if let Some(schema) = response_schema {
-            Some(Arc::new(
-                jsonschema::validator_for(&schema).map_err(|e| format!("Invalid response schema: {}", e))?,
-            ))
-        } else {
-            None
-        };
-
-        Ok(Self {
-            handler: Arc::new(handler),
-            message_schema: message_validator,
-            response_schema: response_validator,
-        })
-    }
-}
-
-/// WebSocket upgrade handler
-///
-/// This is the main entry point for WebSocket connections. Use this as an Axum route
-/// handler by passing it to an Axum router's `.route()` method with `get()`.
-///
-/// # Arguments
-/// * `ws` - WebSocket upgrade from Axum
-/// * `State(state)` - Application state containing the handler and optional schemas
-///
-/// # Returns
-/// An Axum response that upgrades the connection to WebSocket
-///
-/// # Example
-///
-/// ```ignore
-/// use axum::{Router, routing::get, extract::State};
-///
-/// let state = WebSocketState::new(MyHandler);
-/// let router = Router::new()
-///     .route("/ws", get(websocket_handler::<MyHandler>))
-///     .with_state(state);
-/// ```
-pub async fn websocket_handler<H: WebSocketHandler + 'static>(
-    ws: WebSocketUpgrade,
-    State(state): State<WebSocketState<H>>,
-) -> impl IntoResponse {
-    ws.on_upgrade(move |socket| handle_socket(socket, state))
-}
-
-/// Handle an individual WebSocket connection
-async fn handle_socket<H: WebSocketHandler>(mut socket: WebSocket, state: WebSocketState<H>) {
-    info!("WebSocket client connected");
-    trace_ws("socket:connected");
-
-    state.handler.on_connect().await;
-    trace_ws("socket:on_connect:done");
-
-    while let Some(msg) = socket.recv().await {
-        match msg {
-            Ok(Message::Text(text)) => {
-                debug!("Received text message: {}", text);
-                trace_ws(&format!("recv:text len={}", text.len()));
-
-                match serde_json::from_str::<Value>(&text) {
-                    Ok(json_msg) => {
-                        trace_ws("recv:text:json-ok");
-                        if let Some(validator) = &state.message_schema
-                            && !validator.is_valid(&json_msg)
-                        {
-                            error!("Message validation failed");
-                            trace_ws("recv:text:validation-failed");
-                            let error_response = serde_json::json!({
-                                "error": "Message validation failed"
-                            });
-                            if let Ok(error_text) = serde_json::to_string(&error_response) {
-                                trace_ws(&format!("send:validation-error len={}", error_text.len()));
-                                let _ = socket.send(Message::Text(error_text.into())).await;
-                            }
-                            continue;
-                        }
-
-                        if let Some(response) = state.handler.handle_message(json_msg).await {
-                            trace_ws("handler:response:some");
-                            if let Some(validator) = &state.response_schema
-                                && !validator.is_valid(&response)
-                            {
-                                error!("Response validation failed");
-                                trace_ws("send:response:validation-failed");
-                                continue;
-                            }
-
-                            let response_text = serde_json::to_string(&response).unwrap_or_else(|_| "{}".to_string());
-                            let response_len = response_text.len();
-
-                            if let Err(e) = socket.send(Message::Text(response_text.into())).await {
-                                error!("Failed to send response: {}", e);
-                                trace_ws("send:response:error");
-                                break;
-                            }
-                            trace_ws(&format!("send:response len={}", response_len));
-                        } else {
-                            trace_ws("handler:response:none");
-                        }
-                    }
-                    Err(e) => {
-                        warn!("Failed to parse JSON message: {}", e);
-                        trace_ws("recv:text:json-error");
-                        let error_msg = serde_json::json!({
-                            "type": "error",
-                            "message": "Invalid JSON"
-                        });
-                        let error_text = serde_json::to_string(&error_msg).unwrap_or_else(|_| "{}".to_string());
-                        trace_ws(&format!("send:json-error len={}", error_text.len()));
-                        let _ = socket.send(Message::Text(error_text.into())).await;
-                    }
-                }
-            }
-            Ok(Message::Binary(data)) => {
-                debug!("Received binary message: {} bytes", data.len());
-                trace_ws(&format!("recv:binary len={}", data.len()));
-                if let Err(e) = socket.send(Message::Binary(data)).await {
-                    error!("Failed to send binary response: {}", e);
-                    trace_ws("send:binary:error");
-                    break;
-                }
-                trace_ws("send:binary:ok");
-            }
-            Ok(Message::Ping(data)) => {
-                debug!("Received ping");
-                trace_ws(&format!("recv:ping len={}", data.len()));
-                if let Err(e) = socket.send(Message::Pong(data)).await {
-                    error!("Failed to send pong: {}", e);
-                    trace_ws("send:pong:error");
-                    break;
-                }
-                trace_ws("send:pong:ok");
-            }
-            Ok(Message::Pong(_)) => {
-                debug!("Received pong");
-                trace_ws("recv:pong");
-            }
-            Ok(Message::Close(_)) => {
-                info!("Client closed connection");
-                trace_ws("recv:close");
-                break;
-            }
-            Err(e) => {
-                error!("WebSocket error: {}", e);
-                trace_ws(&format!("recv:error {}", e));
-                break;
-            }
-        }
-    }
-
-    state.handler.on_disconnect().await;
-    trace_ws("socket:on_disconnect:done");
-    info!("WebSocket client disconnected");
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use std::sync::Mutex;
-    use std::sync::atomic::{AtomicUsize, Ordering};
-
-    #[derive(Debug)]
-    struct EchoHandler;
-
-    impl WebSocketHandler for EchoHandler {
-        async fn handle_message(&self, message: Value) -> Option<Value> {
-            Some(message)
-        }
-    }
-
-    #[derive(Debug)]
-    struct TrackingHandler {
-        connect_count: Arc<AtomicUsize>,
-        disconnect_count: Arc<AtomicUsize>,
-        message_count: Arc<AtomicUsize>,
-        messages: Arc<Mutex<Vec<Value>>>,
-    }
-
-    impl TrackingHandler {
-        fn new() -> Self {
-            Self {
-                connect_count: Arc::new(AtomicUsize::new(0)),
-                disconnect_count: Arc::new(AtomicUsize::new(0)),
-                message_count: Arc::new(AtomicUsize::new(0)),
-                messages: Arc::new(Mutex::new(Vec::new())),
-            }
-        }
-    }
-
-    impl WebSocketHandler for TrackingHandler {
-        async fn handle_message(&self, message: Value) -> Option<Value> {
-            self.message_count.fetch_add(1, Ordering::SeqCst);
-            self.messages.lock().unwrap().push(message.clone());
-            Some(message)
-        }
-
-        async fn on_connect(&self) {
-            self.connect_count.fetch_add(1, Ordering::SeqCst);
-        }
-
-        async fn on_disconnect(&self) {
-            self.disconnect_count.fetch_add(1, Ordering::SeqCst);
-        }
-    }
-
-    #[derive(Debug)]
-    struct SelectiveHandler;
-
-    impl WebSocketHandler for SelectiveHandler {
-        async fn handle_message(&self, message: Value) -> Option<Value> {
-            if message.get("respond").is_some_and(|v| v.as_bool().unwrap_or(false)) {
-                Some(serde_json::json!({"response": "acknowledged"}))
-            } else {
-                None
-            }
-        }
-    }
-
-    #[derive(Debug)]
-    struct TransformHandler;
-
-    impl WebSocketHandler for TransformHandler {
-        async fn handle_message(&self, message: Value) -> Option<Value> {
-            message.as_object().map_or(None, |obj| {
-                let mut resp = obj.clone();
-                resp.insert("processed".to_string(), Value::Bool(true));
-                Some(Value::Object(resp))
-            })
-        }
-    }
-
-    #[test]
-    fn test_websocket_state_creation() {
-        let handler: EchoHandler = EchoHandler;
-        let state: WebSocketState<EchoHandler> = WebSocketState::new(handler);
-        let cloned: WebSocketState<EchoHandler> = state.clone();
-        assert!(Arc::ptr_eq(&state.handler, &cloned.handler));
-    }
-
-    #[test]
-    fn test_websocket_state_with_valid_schema() {
-        let handler: EchoHandler = EchoHandler;
-        let schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "type": {"type": "string"}
-            }
-        });
-
-        let result: Result<WebSocketState<EchoHandler>, String> =
-            WebSocketState::with_schemas(handler, Some(schema), None);
-        assert!(result.is_ok());
-    }
-
-    #[test]
-    fn test_websocket_state_with_invalid_schema() {
-        let handler: EchoHandler = EchoHandler;
-        let invalid_schema: serde_json::Value = serde_json::json!({
-            "type": "not_a_real_type",
-            "invalid": "schema"
-        });
-
-        let result: Result<WebSocketState<EchoHandler>, String> =
-            WebSocketState::with_schemas(handler, Some(invalid_schema), None);
-        assert!(result.is_err());
-        if let Err(error_msg) = result {
-            assert!(error_msg.contains("Invalid message schema"));
-        }
-    }
-
-    #[test]
-    fn test_websocket_state_with_both_schemas() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"action": {"type": "string"}}
-        });
-        let response_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"result": {"type": "string"}}
-        });
-
-        let result: Result<WebSocketState<EchoHandler>, String> =
-            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema));
-        assert!(result.is_ok());
-        let state: WebSocketState<EchoHandler> = result.unwrap();
-        assert!(state.message_schema.is_some());
-        assert!(state.response_schema.is_some());
-    }
-
-    #[test]
-    fn test_websocket_state_cloning_preserves_schemas() {
-        let handler: EchoHandler = EchoHandler;
-        let schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"id": {"type": "integer"}}
-        });
-
-        let state: WebSocketState<EchoHandler> = WebSocketState::with_schemas(handler, Some(schema), None).unwrap();
-        let cloned: WebSocketState<EchoHandler> = state.clone();
-
-        assert!(cloned.message_schema.is_some());
-        assert!(cloned.response_schema.is_none());
-        assert!(Arc::ptr_eq(&state.handler, &cloned.handler));
-    }
-
-    #[tokio::test]
-    async fn test_tracking_handler_lifecycle() {
-        let handler: TrackingHandler = TrackingHandler::new();
-        handler.on_connect().await;
-        assert_eq!(handler.connect_count.load(Ordering::SeqCst), 1);
-
-        let msg: Value = serde_json::json!({"test": "data"});
-        let _response: Option<Value> = handler.handle_message(msg).await;
-        assert_eq!(handler.message_count.load(Ordering::SeqCst), 1);
-
-        handler.on_disconnect().await;
-        assert_eq!(handler.disconnect_count.load(Ordering::SeqCst), 1);
-    }
-
-    #[tokio::test]
-    async fn test_selective_handler_responds_conditionally() {
-        let handler: SelectiveHandler = SelectiveHandler;
-
-        let respond_msg: Value = serde_json::json!({"respond": true});
-        let response1: Option<Value> = handler.handle_message(respond_msg).await;
-        assert!(response1.is_some());
-        assert_eq!(response1.unwrap(), serde_json::json!({"response": "acknowledged"}));
-
-        let no_respond_msg: Value = serde_json::json!({"respond": false});
-        let response2: Option<Value> = handler.handle_message(no_respond_msg).await;
-        assert!(response2.is_none());
-    }
-
-    #[tokio::test]
-    async fn test_transform_handler_modifies_message() {
-        let handler: TransformHandler = TransformHandler;
-        let original: Value = serde_json::json!({"name": "test"});
-        let transformed: Option<Value> = handler.handle_message(original).await;
-
-        assert!(transformed.is_some());
-        let resp: Value = transformed.unwrap();
-        assert_eq!(resp.get("name").unwrap(), "test");
-        assert_eq!(resp.get("processed").unwrap(), true);
-    }
-
-    #[tokio::test]
-    async fn test_echo_handler_preserves_json_types() {
-        let handler: EchoHandler = EchoHandler;
-
-        let messages: Vec<Value> = vec![
-            serde_json::json!({"string": "value"}),
-            serde_json::json!({"number": 42}),
-            serde_json::json!({"float": 3.14}),
-            serde_json::json!({"bool": true}),
-            serde_json::json!({"null": null}),
-            serde_json::json!({"array": [1, 2, 3]}),
-        ];
-
-        for msg in messages {
-            let response: Option<Value> = handler.handle_message(msg.clone()).await;
-            assert!(response.is_some());
-            assert_eq!(response.unwrap(), msg);
-        }
-    }
-
-    #[tokio::test]
-    async fn test_tracking_handler_accumulates_messages() {
-        let handler: TrackingHandler = TrackingHandler::new();
-
-        let messages: Vec<Value> = vec![
-            serde_json::json!({"id": 1}),
-            serde_json::json!({"id": 2}),
-            serde_json::json!({"id": 3}),
-        ];
-
-        for msg in messages {
-            let _: Option<Value> = handler.handle_message(msg).await;
-        }
-
-        assert_eq!(handler.message_count.load(Ordering::SeqCst), 3);
-        let stored: Vec<Value> = handler.messages.lock().unwrap().clone();
-        assert_eq!(stored.len(), 3);
-        assert_eq!(stored[0].get("id").unwrap(), 1);
-        assert_eq!(stored[1].get("id").unwrap(), 2);
-        assert_eq!(stored[2].get("id").unwrap(), 3);
-    }
-
-    #[tokio::test]
-    async fn test_echo_handler_with_nested_json() {
-        let handler: EchoHandler = EchoHandler;
-        let nested: Value = serde_json::json!({
-            "level1": {
-                "level2": {
-                    "level3": {
-                        "value": "deeply nested"
-                    }
-                }
-            }
-        });
-
-        let response: Option<Value> = handler.handle_message(nested.clone()).await;
-        assert!(response.is_some());
-        assert_eq!(response.unwrap(), nested);
-    }
-
-    #[tokio::test]
-    async fn test_echo_handler_with_large_array() {
-        let handler: EchoHandler = EchoHandler;
-        let large_array: Value = serde_json::json!({
-            "items": (0..1000).collect::<Vec<i32>>()
-        });
-
-        let response: Option<Value> = handler.handle_message(large_array.clone()).await;
-        assert!(response.is_some());
-        assert_eq!(response.unwrap(), large_array);
-    }
-
-    #[tokio::test]
-    async fn test_echo_handler_with_unicode() {
-        let handler: EchoHandler = EchoHandler;
-        let unicode_msg: Value = serde_json::json!({
-            "emoji": "🚀",
-            "chinese": "你好",
-            "arabic": "مرحبا",
-            "mixed": "Hello 世界 🌍"
-        });
-
-        let response: Option<Value> = handler.handle_message(unicode_msg.clone()).await;
-        assert!(response.is_some());
-        assert_eq!(response.unwrap(), unicode_msg);
-    }
-
-    #[test]
-    fn test_websocket_state_schemas_are_independent() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({"type": "object"});
-        let response_schema: serde_json::Value = serde_json::json!({"type": "array"});
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema)).unwrap();
-
-        let cloned: WebSocketState<EchoHandler> = state.clone();
-
-        assert!(state.message_schema.is_some());
-        assert!(state.response_schema.is_some());
-        assert!(cloned.message_schema.is_some());
-        assert!(cloned.response_schema.is_some());
-    }
-
-    #[test]
-    fn test_message_schema_validation_with_required_field() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"type": {"type": "string"}},
-            "required": ["type"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        assert!(state.message_schema.is_some());
-        assert!(state.response_schema.is_none());
-
-        let valid_msg: Value = serde_json::json!({"type": "test"});
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"other": "field"});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    #[test]
-    fn test_response_schema_validation_with_required_field() {
-        let handler: EchoHandler = EchoHandler;
-        let response_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"status": {"type": "string"}},
-            "required": ["status"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
-
-        assert!(state.message_schema.is_none());
-        assert!(state.response_schema.is_some());
-
-        let valid_response: Value = serde_json::json!({"status": "ok"});
-        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
-        assert!(validator.is_valid(&valid_response));
-
-        let invalid_response: Value = serde_json::json!({"other": "field"});
-        assert!(!validator.is_valid(&invalid_response));
-    }
-
-    #[test]
-    fn test_invalid_message_schema_returns_error() {
-        let handler: EchoHandler = EchoHandler;
-        let invalid_schema: serde_json::Value = serde_json::json!({
-            "type": "invalid_type_value",
-            "properties": {}
-        });
-
-        let result: Result<WebSocketState<EchoHandler>, String> =
-            WebSocketState::with_schemas(handler, Some(invalid_schema), None);
-
-        assert!(result.is_err());
-        match result {
-            Err(error_msg) => assert!(error_msg.contains("Invalid message schema")),
-            Ok(_) => panic!("Expected error but got Ok"),
-        }
-    }
-
-    #[test]
-    fn test_invalid_response_schema_returns_error() {
-        let handler: EchoHandler = EchoHandler;
-        let invalid_schema: serde_json::Value = serde_json::json!({
-            "type": "definitely_not_valid"
-        });
-
-        let result: Result<WebSocketState<EchoHandler>, String> =
-            WebSocketState::with_schemas(handler, None, Some(invalid_schema));
-
-        assert!(result.is_err());
-        match result {
-            Err(error_msg) => assert!(error_msg.contains("Invalid response schema")),
-            Ok(_) => panic!("Expected error but got Ok"),
-        }
-    }
-
-    #[tokio::test]
-    async fn test_handler_returning_none_response() {
-        let handler: SelectiveHandler = SelectiveHandler;
-
-        let no_response_msg: Value = serde_json::json!({"respond": false});
-        let result: Option<Value> = handler.handle_message(no_response_msg).await;
-
-        assert!(result.is_none());
-    }
-
-    #[tokio::test]
-    async fn test_handler_with_complex_schema_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "user": {
-                    "type": "object",
-                    "properties": {
-                        "id": {"type": "integer"},
-                        "name": {"type": "string"}
-                    },
-                    "required": ["id", "name"]
-                },
-                "action": {"type": "string"}
-            },
-            "required": ["user", "action"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let valid_msg: Value = serde_json::json!({
-            "user": {"id": 123, "name": "Alice"},
-            "action": "create"
-        });
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({
-            "user": {"id": "not_an_int", "name": "Bob"},
-            "action": "create"
-        });
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    #[tokio::test]
-    async fn test_tracking_handler_with_multiple_message_types() {
-        let handler: TrackingHandler = TrackingHandler::new();
-
-        let messages: Vec<Value> = vec![
-            serde_json::json!({"type": "text", "content": "hello"}),
-            serde_json::json!({"type": "image", "url": "http://example.com/image.png"}),
-            serde_json::json!({"type": "video", "duration": 120}),
-        ];
-
-        for msg in messages {
-            let _: Option<Value> = handler.handle_message(msg).await;
-        }
-
-        assert_eq!(handler.message_count.load(Ordering::SeqCst), 3);
-        let stored: Vec<Value> = handler.messages.lock().unwrap().clone();
-        assert_eq!(stored.len(), 3);
-        assert_eq!(stored[0].get("type").unwrap(), "text");
-        assert_eq!(stored[1].get("type").unwrap(), "image");
-        assert_eq!(stored[2].get("type").unwrap(), "video");
-    }
-
-    #[tokio::test]
-    async fn test_selective_handler_with_explicit_false() {
-        let handler: SelectiveHandler = SelectiveHandler;
-
-        let msg: Value = serde_json::json!({"respond": false, "data": "test"});
-        let response: Option<Value> = handler.handle_message(msg).await;
-
-        assert!(response.is_none());
-    }
-
-    #[tokio::test]
-    async fn test_selective_handler_without_respond_field() {
-        let handler: SelectiveHandler = SelectiveHandler;
-
-        let msg: Value = serde_json::json!({"data": "test"});
-        let response: Option<Value> = handler.handle_message(msg).await;
-
-        assert!(response.is_none());
-    }
-
-    #[tokio::test]
-    async fn test_transform_handler_with_empty_object() {
-        let handler: TransformHandler = TransformHandler;
-        let original: Value = serde_json::json!({});
-        let transformed: Option<Value> = handler.handle_message(original).await;
-
-        assert!(transformed.is_some());
-        let resp: Value = transformed.unwrap();
-        assert_eq!(resp.get("processed").unwrap(), true);
-        assert_eq!(resp.as_object().unwrap().len(), 1);
-    }
-
-    #[tokio::test]
-    async fn test_transform_handler_preserves_all_fields() {
-        let handler: TransformHandler = TransformHandler;
-        let original: Value = serde_json::json!({
-            "field1": "value1",
-            "field2": 42,
-            "field3": true,
-            "nested": {"key": "value"}
-        });
-        let transformed: Option<Value> = handler.handle_message(original.clone()).await;
-
-        assert!(transformed.is_some());
-        let resp: Value = transformed.unwrap();
-        assert_eq!(resp.get("field1").unwrap(), "value1");
-        assert_eq!(resp.get("field2").unwrap(), 42);
-        assert_eq!(resp.get("field3").unwrap(), true);
-        assert_eq!(resp.get("nested").unwrap(), &serde_json::json!({"key": "value"}));
-        assert_eq!(resp.get("processed").unwrap(), true);
-    }
-
-    #[tokio::test]
-    async fn test_transform_handler_with_non_object_input() {
-        let handler: TransformHandler = TransformHandler;
-
-        let array: Value = serde_json::json!([1, 2, 3]);
-        let response1: Option<Value> = handler.handle_message(array).await;
-        assert!(response1.is_none());
-
-        let string: Value = serde_json::json!("not an object");
-        let response2: Option<Value> = handler.handle_message(string).await;
-        assert!(response2.is_none());
-
-        let number: Value = serde_json::json!(42);
-        let response3: Option<Value> = handler.handle_message(number).await;
-        assert!(response3.is_none());
-    }
-
-    /// Test message validation failure with schema constraint
-    #[test]
-    fn test_message_schema_rejects_wrong_type() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"id": {"type": "integer"}},
-            "required": ["id"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let invalid_msg: Value = serde_json::json!({"id": "not_an_integer"});
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test response schema validation failure
-    #[test]
-    fn test_response_schema_rejects_invalid_type() {
-        let handler: EchoHandler = EchoHandler;
-        let response_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"count": {"type": "integer"}},
-            "required": ["count"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
-
-        let invalid_response: Value = serde_json::json!([1, 2, 3]);
-        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
-        assert!(!validator.is_valid(&invalid_response));
-    }
-
-    /// Test message with multiple required fields missing
-    #[test]
-    fn test_message_missing_multiple_required_fields() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "user_id": {"type": "integer"},
-                "action": {"type": "string"},
-                "timestamp": {"type": "string"}
-            },
-            "required": ["user_id", "action", "timestamp"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let invalid_msg: Value = serde_json::json!({"other": "value"});
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        assert!(!validator.is_valid(&invalid_msg));
-
-        let partial_msg: Value = serde_json::json!({"user_id": 123});
-        assert!(!validator.is_valid(&partial_msg));
-    }
-
-    /// Test deeply nested schema validation with required nested properties
-    #[test]
-    fn test_deeply_nested_schema_validation_failure() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "metadata": {
-                    "type": "object",
-                    "properties": {
-                        "request": {
-                            "type": "object",
-                            "properties": {
-                                "id": {"type": "string"}
-                            },
-                            "required": ["id"]
-                        }
-                    },
-                    "required": ["request"]
-                }
-            },
-            "required": ["metadata"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let invalid_msg: Value = serde_json::json!({
-            "metadata": {
-                "request": {}
-            }
-        });
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test array property validation with items constraint
-    #[test]
-    fn test_array_property_type_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "ids": {
-                    "type": "array",
-                    "items": {"type": "integer"}
-                }
-            }
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"ids": [1, 2, 3]});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"ids": [1, "two", 3]});
-        assert!(!validator.is_valid(&invalid_msg));
-
-        let invalid_msg2: Value = serde_json::json!({"ids": "not_an_array"});
-        assert!(!validator.is_valid(&invalid_msg2));
-    }
-
-    /// Test enum/const property validation
-    #[test]
-    fn test_enum_property_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "status": {
-                    "type": "string",
-                    "enum": ["pending", "active", "completed"]
-                }
-            },
-            "required": ["status"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"status": "active"});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"status": "unknown"});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test minimum/maximum constraints on numbers
-    #[test]
-    fn test_number_range_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "age": {
-                    "type": "integer",
-                    "minimum": 0,
-                    "maximum": 150
-                }
-            },
-            "required": ["age"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"age": 25});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"age": -1});
-        assert!(!validator.is_valid(&invalid_msg));
-
-        let invalid_msg2: Value = serde_json::json!({"age": 200});
-        assert!(!validator.is_valid(&invalid_msg2));
-    }
-
-    /// Test string length constraints
-    #[test]
-    fn test_string_length_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "username": {
-                    "type": "string",
-                    "minLength": 3,
-                    "maxLength": 20
-                }
-            },
-            "required": ["username"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"username": "alice"});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"username": "ab"});
-        assert!(!validator.is_valid(&invalid_msg));
-
-        let invalid_msg2: Value =
-            serde_json::json!({"username": "this_is_a_very_long_username_over_twenty_characters"});
-        assert!(!validator.is_valid(&invalid_msg2));
-    }
-
-    /// Test pattern (regex) validation
-    #[test]
-    fn test_pattern_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "email": {
-                    "type": "string",
-                    "pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$"
-                }
-            },
-            "required": ["email"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"email": "user@example.com"});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"email": "user@example"});
-        assert!(!validator.is_valid(&invalid_msg));
-
-        let invalid_msg2: Value = serde_json::json!({"email": "userexample.com"});
-        assert!(!validator.is_valid(&invalid_msg2));
-    }
-
-    /// Test additionalProperties constraint
-    #[test]
-    fn test_additional_properties_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "name": {"type": "string"}
-            },
-            "additionalProperties": false
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"name": "Alice"});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"name": "Bob", "age": 30});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test oneOf constraint (mutually exclusive properties)
-    #[test]
-    fn test_one_of_constraint() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "oneOf": [
-                {
-                    "properties": {"type": {"const": "text"}},
-                    "required": ["type"]
-                },
-                {
-                    "properties": {"type": {"const": "number"}},
-                    "required": ["type"]
-                }
-            ]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"type": "text"});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"type": "unknown"});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test anyOf constraint (at least one match)
-    #[test]
-    fn test_any_of_constraint() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "value": {"type": ["string", "integer"]}
-            },
-            "required": ["value"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let msg1: Value = serde_json::json!({"value": "text"});
-        assert!(validator.is_valid(&msg1));
-
-        let msg2: Value = serde_json::json!({"value": 42});
-        assert!(validator.is_valid(&msg2));
-
-        let invalid_msg: Value = serde_json::json!({"value": true});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test response validation with complex constraints
-    #[test]
-    fn test_response_schema_with_multiple_constraints() {
-        let handler: EchoHandler = EchoHandler;
-        let response_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "success": {"type": "boolean"},
-                "data": {
-                    "type": "object",
-                    "properties": {
-                        "items": {
-                            "type": "array",
-                            "items": {"type": "object"},
-                            "minItems": 1
-                        }
-                    },
-                    "required": ["items"]
-                }
-            },
-            "required": ["success", "data"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
-
-        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
-
-        let valid_response: Value = serde_json::json!({
-            "success": true,
-            "data": {
-                "items": [{"id": 1}]
-            }
-        });
-        assert!(validator.is_valid(&valid_response));
-
-        let invalid_response: Value = serde_json::json!({
-            "success": true,
-            "data": {
-                "items": []
-            }
-        });
-        assert!(!validator.is_valid(&invalid_response));
-
-        let invalid_response2: Value = serde_json::json!({
-            "success": true
-        });
-        assert!(!validator.is_valid(&invalid_response2));
-    }
-
-    /// Test null type validation
-    #[test]
-    fn test_null_value_validation() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "optional_field": {"type": ["string", "null"]},
-                "required_field": {"type": "string"}
-            },
-            "required": ["required_field"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let msg1: Value = serde_json::json!({
-            "optional_field": null,
-            "required_field": "value"
-        });
-        assert!(validator.is_valid(&msg1));
-
-        let msg2: Value = serde_json::json!({"required_field": "value"});
-        assert!(validator.is_valid(&msg2));
-
-        let invalid_msg: Value = serde_json::json!({"required_field": null});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-
-    /// Test schema with default values (they don't change validation)
-    #[test]
-    fn test_schema_with_defaults_still_validates() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "status": {
-                    "type": "string",
-                    "default": "pending"
-                }
-            }
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let msg: Value = serde_json::json!({});
-        assert!(validator.is_valid(&msg));
-    }
-
-    /// Test both message and response schema validation together
-    #[test]
-    fn test_both_schemas_validate_independently() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"action": {"type": "string"}},
-            "required": ["action"]
-        });
-        let response_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {"result": {"type": "string"}},
-            "required": ["result"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema)).unwrap();
-
-        let msg_validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-        let resp_validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"action": "test"});
-        let invalid_response: Value = serde_json::json!({"data": "oops"});
-
-        assert!(msg_validator.is_valid(&valid_msg));
-        assert!(!resp_validator.is_valid(&invalid_response));
-
-        let invalid_msg: Value = serde_json::json!({"data": "oops"});
-        let valid_response: Value = serde_json::json!({"result": "ok"});
-
-        assert!(!msg_validator.is_valid(&invalid_msg));
-        assert!(resp_validator.is_valid(&valid_response));
-    }
-
-    /// Test validation with very long/large payload
-    #[test]
-    fn test_validation_with_large_payload() {
-        let handler: EchoHandler = EchoHandler;
-        let message_schema: serde_json::Value = serde_json::json!({
-            "type": "object",
-            "properties": {
-                "items": {
-                    "type": "array",
-                    "items": {"type": "integer"}
-                }
-            },
-            "required": ["items"]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let mut items = Vec::new();
-        for i in 0..10_000 {
-            items.push(i);
-        }
-        let large_msg: Value = serde_json::json!({"items": items});
-
-        assert!(validator.is_valid(&large_msg));
-    }
-
-    /// Test validation error doesn't panic with invalid schema combinations
-    #[test]
-    fn test_mutually_exclusive_schema_properties() {
-        let handler: EchoHandler = EchoHandler;
-
-        let message_schema: serde_json::Value = serde_json::json!({
-            "allOf": [
-                {
-                    "type": "object",
-                    "properties": {"a": {"type": "string"}},
-                    "required": ["a"]
-                },
-                {
-                    "type": "object",
-                    "properties": {"b": {"type": "integer"}},
-                    "required": ["b"]
-                }
-            ]
-        });
-
-        let state: WebSocketState<EchoHandler> =
-            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
-
-        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
-
-        let valid_msg: Value = serde_json::json!({"a": "text", "b": 42});
-        assert!(validator.is_valid(&valid_msg));
-
-        let invalid_msg: Value = serde_json::json!({"a": "text"});
-        assert!(!validator.is_valid(&invalid_msg));
-    }
-}
+//! WebSocket support for Spikard
+//!
+//! Provides WebSocket connection handling with message validation and routing.
+
+use axum::{
+    extract::{
+        State,
+        ws::{Message, WebSocket, WebSocketUpgrade},
+    },
+    response::IntoResponse,
+};
+use serde_json::Value;
+use std::sync::Arc;
+use tracing::{debug, error, info, warn};
+
+fn trace_ws(message: &str) {
+    if std::env::var("SPIKARD_WS_TRACE").ok().as_deref() == Some("1") {
+        eprintln!("[spikard-ws] {message}");
+    }
+}
+
+/// WebSocket message handler trait
+///
+/// Implement this trait to create custom WebSocket message handlers for your application.
+/// The handler processes JSON messages received from WebSocket clients and can optionally
+/// send responses back.
+///
+/// # Implementing the Trait
+///
+/// You must implement the `handle_message` method. The `on_connect` and `on_disconnect`
+/// methods are optional and provide lifecycle hooks.
+///
+/// # Example
+///
+/// ```ignore
+/// use spikard_http::websocket::WebSocketHandler;
+/// use serde_json::{json, Value};
+///
+/// struct EchoHandler;
+///
+/// #[async_trait]
+/// impl WebSocketHandler for EchoHandler {
+///     async fn handle_message(&self, message: Value) -> Option<Value> {
+///         // Echo the message back to the client
+///         Some(message)
+///     }
+///
+///     async fn on_connect(&self) {
+///         println!("Client connected");
+///     }
+///
+///     async fn on_disconnect(&self) {
+///         println!("Client disconnected");
+///     }
+/// }
+/// ```
+pub trait WebSocketHandler: Send + Sync {
+    /// Handle incoming WebSocket message
+    ///
+    /// Called whenever a text message is received from a WebSocket client.
+    /// Messages are automatically parsed as JSON.
+    ///
+    /// # Arguments
+    /// * `message` - JSON value received from the client
+    ///
+    /// # Returns
+    /// * `Some(value)` - JSON value to send back to the client
+    /// * `None` - No response to send
+    fn handle_message(&self, message: Value) -> impl std::future::Future<Output = Option<Value>> + Send;
+
+    /// Called when a client connects to the WebSocket
+    ///
+    /// Optional lifecycle hook invoked when a new WebSocket connection is established.
+    /// Default implementation does nothing.
+    fn on_connect(&self) -> impl std::future::Future<Output = ()> + Send {
+        async {}
+    }
+
+    /// Called when a client disconnects from the WebSocket
+    ///
+    /// Optional lifecycle hook invoked when a WebSocket connection is closed
+    /// (either by the client or due to an error). Default implementation does nothing.
+    fn on_disconnect(&self) -> impl std::future::Future<Output = ()> + Send {
+        async {}
+    }
+}
+
+/// WebSocket state shared across connections
+///
+/// Contains the message handler and optional JSON schemas for validating
+/// incoming and outgoing messages. This state is shared among all connections
+/// to the same WebSocket endpoint.
+#[derive(Debug)]
+pub struct WebSocketState<H: WebSocketHandler> {
+    /// The message handler implementation
+    handler: Arc<H>,
+    /// Optional JSON Schema for validating incoming messages
+    message_schema: Option<Arc<jsonschema::Validator>>,
+    /// Optional JSON Schema for validating outgoing responses
+    response_schema: Option<Arc<jsonschema::Validator>>,
+}
+
+impl<H: WebSocketHandler> Clone for WebSocketState<H> {
+    fn clone(&self) -> Self {
+        Self {
+            handler: Arc::clone(&self.handler),
+            message_schema: self.message_schema.clone(),
+            response_schema: self.response_schema.clone(),
+        }
+    }
+}
+
+impl<H: WebSocketHandler + 'static> WebSocketState<H> {
+    /// Create new WebSocket state with a handler
+    ///
+    /// Creates a new state without message or response validation schemas.
+    /// Messages and responses are not validated.
+    ///
+    /// # Arguments
+    /// * `handler` - The message handler implementation
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// let state = WebSocketState::new(MyHandler);
+    /// ```
+    pub fn new(handler: H) -> Self {
+        Self {
+            handler: Arc::new(handler),
+            message_schema: None,
+            response_schema: None,
+        }
+    }
+
+    /// Create new WebSocket state with a handler and optional validation schemas
+    ///
+    /// Creates a new state with optional JSON schemas for validating incoming messages
+    /// and outgoing responses. If a schema is provided and validation fails, the message
+    /// or response is rejected.
+    ///
+    /// # Arguments
+    /// * `handler` - The message handler implementation
+    /// * `message_schema` - Optional JSON schema for validating client messages
+    /// * `response_schema` - Optional JSON schema for validating handler responses
+    ///
+    /// # Returns
+    /// * `Ok(state)` - Successfully created state
+    /// * `Err(msg)` - Invalid schema provided
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// use serde_json::json;
+    ///
+    /// let message_schema = json!({
+    ///     "type": "object",
+    ///     "properties": {
+    ///         "type": {"type": "string"},
+    ///         "data": {"type": "string"}
+    ///     }
+    /// });
+    ///
+    /// let state = WebSocketState::with_schemas(
+    ///     MyHandler,
+    ///     Some(message_schema),
+    ///     None,
+    /// )?;
+    /// ```
+    pub fn with_schemas(
+        handler: H,
+        message_schema: Option<serde_json::Value>,
+        response_schema: Option<serde_json::Value>,
+    ) -> Result<Self, String> {
+        let message_validator = if let Some(schema) = message_schema {
+            Some(Arc::new(
+                jsonschema::validator_for(&schema).map_err(|e| format!("Invalid message schema: {}", e))?,
+            ))
+        } else {
+            None
+        };
+
+        let response_validator = if let Some(schema) = response_schema {
+            Some(Arc::new(
+                jsonschema::validator_for(&schema).map_err(|e| format!("Invalid response schema: {}", e))?,
+            ))
+        } else {
+            None
+        };
+
+        Ok(Self {
+            handler: Arc::new(handler),
+            message_schema: message_validator,
+            response_schema: response_validator,
+        })
+    }
+
+    /// Invoke the connection hook for testing.
+    pub async fn on_connect(&self) {
+        self.handler.on_connect().await;
+    }
+
+    /// Invoke the disconnect hook for testing.
+    pub async fn on_disconnect(&self) {
+        self.handler.on_disconnect().await;
+    }
+
+    /// Validate and handle an incoming message without a socket transport.
+    pub async fn handle_message_validated(&self, message: Value) -> Result<Option<Value>, String> {
+        if let Some(validator) = &self.message_schema
+            && !validator.is_valid(&message)
+        {
+            return Err("Message validation failed".to_string());
+        }
+
+        let response = self.handler.handle_message(message).await;
+        if let Some(ref value) = response
+            && let Some(validator) = &self.response_schema
+            && !validator.is_valid(value)
+        {
+            return Ok(None);
+        }
+
+        Ok(response)
+    }
+}
+
+/// WebSocket upgrade handler
+///
+/// This is the main entry point for WebSocket connections. Use this as an Axum route
+/// handler by passing it to an Axum router's `.route()` method with `get()`.
+///
+/// # Arguments
+/// * `ws` - WebSocket upgrade from Axum
+/// * `State(state)` - Application state containing the handler and optional schemas
+///
+/// # Returns
+/// An Axum response that upgrades the connection to WebSocket
+///
+/// # Example
+///
+/// ```ignore
+/// use axum::{Router, routing::get, extract::State};
+///
+/// let state = WebSocketState::new(MyHandler);
+/// let router = Router::new()
+///     .route("/ws", get(websocket_handler::<MyHandler>))
+///     .with_state(state);
+/// ```
+pub async fn websocket_handler<H: WebSocketHandler + 'static>(
+    ws: WebSocketUpgrade,
+    State(state): State<WebSocketState<H>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_socket(socket, state))
+}
+
+/// Handle an individual WebSocket connection
+async fn handle_socket<H: WebSocketHandler>(mut socket: WebSocket, state: WebSocketState<H>) {
+    info!("WebSocket client connected");
+    trace_ws("socket:connected");
+
+    state.handler.on_connect().await;
+    trace_ws("socket:on_connect:done");
+
+    while let Some(msg) = socket.recv().await {
+        match msg {
+            Ok(Message::Text(text)) => {
+                debug!("Received text message: {}", text);
+                trace_ws(&format!("recv:text len={}", text.len()));
+
+                match serde_json::from_str::<Value>(&text) {
+                    Ok(json_msg) => {
+                        trace_ws("recv:text:json-ok");
+                        if let Some(validator) = &state.message_schema
+                            && !validator.is_valid(&json_msg)
+                        {
+                            error!("Message validation failed");
+                            trace_ws("recv:text:validation-failed");
+                            let error_response = serde_json::json!({
+                                "error": "Message validation failed"
+                            });
+                            if let Ok(error_text) = serde_json::to_string(&error_response) {
+                                trace_ws(&format!("send:validation-error len={}", error_text.len()));
+                                let _ = socket.send(Message::Text(error_text.into())).await;
+                            }
+                            continue;
+                        }
+
+                        if let Some(response) = state.handler.handle_message(json_msg).await {
+                            trace_ws("handler:response:some");
+                            if let Some(validator) = &state.response_schema
+                                && !validator.is_valid(&response)
+                            {
+                                error!("Response validation failed");
+                                trace_ws("send:response:validation-failed");
+                                continue;
+                            }
+
+                            let response_text = serde_json::to_string(&response).unwrap_or_else(|_| "{}".to_string());
+                            let response_len = response_text.len();
+
+                            if let Err(e) = socket.send(Message::Text(response_text.into())).await {
+                                error!("Failed to send response: {}", e);
+                                trace_ws("send:response:error");
+                                break;
+                            }
+                            trace_ws(&format!("send:response len={}", response_len));
+                        } else {
+                            trace_ws("handler:response:none");
+                        }
+                    }
+                    Err(e) => {
+                        warn!("Failed to parse JSON message: {}", e);
+                        trace_ws("recv:text:json-error");
+                        let error_msg = serde_json::json!({
+                            "type": "error",
+                            "message": "Invalid JSON"
+                        });
+                        let error_text = serde_json::to_string(&error_msg).unwrap_or_else(|_| "{}".to_string());
+                        trace_ws(&format!("send:json-error len={}", error_text.len()));
+                        let _ = socket.send(Message::Text(error_text.into())).await;
+                    }
+                }
+            }
+            Ok(Message::Binary(data)) => {
+                debug!("Received binary message: {} bytes", data.len());
+                trace_ws(&format!("recv:binary len={}", data.len()));
+                if let Err(e) = socket.send(Message::Binary(data)).await {
+                    error!("Failed to send binary response: {}", e);
+                    trace_ws("send:binary:error");
+                    break;
+                }
+                trace_ws("send:binary:ok");
+            }
+            Ok(Message::Ping(data)) => {
+                debug!("Received ping");
+                trace_ws(&format!("recv:ping len={}", data.len()));
+                if let Err(e) = socket.send(Message::Pong(data)).await {
+                    error!("Failed to send pong: {}", e);
+                    trace_ws("send:pong:error");
+                    break;
+                }
+                trace_ws("send:pong:ok");
+            }
+            Ok(Message::Pong(_)) => {
+                debug!("Received pong");
+                trace_ws("recv:pong");
+            }
+            Ok(Message::Close(_)) => {
+                info!("Client closed connection");
+                trace_ws("recv:close");
+                break;
+            }
+            Err(e) => {
+                error!("WebSocket error: {}", e);
+                trace_ws(&format!("recv:error {}", e));
+                break;
+            }
+        }
+    }
+
+    state.handler.on_disconnect().await;
+    trace_ws("socket:on_disconnect:done");
+    info!("WebSocket client disconnected");
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    #[derive(Debug)]
+    struct EchoHandler;
+
+    impl WebSocketHandler for EchoHandler {
+        async fn handle_message(&self, message: Value) -> Option<Value> {
+            Some(message)
+        }
+    }
+
+    #[derive(Debug)]
+    struct TrackingHandler {
+        connect_count: Arc<AtomicUsize>,
+        disconnect_count: Arc<AtomicUsize>,
+        message_count: Arc<AtomicUsize>,
+        messages: Arc<Mutex<Vec<Value>>>,
+    }
+
+    impl TrackingHandler {
+        fn new() -> Self {
+            Self {
+                connect_count: Arc::new(AtomicUsize::new(0)),
+                disconnect_count: Arc::new(AtomicUsize::new(0)),
+                message_count: Arc::new(AtomicUsize::new(0)),
+                messages: Arc::new(Mutex::new(Vec::new())),
+            }
+        }
+    }
+
+    impl WebSocketHandler for TrackingHandler {
+        async fn handle_message(&self, message: Value) -> Option<Value> {
+            self.message_count.fetch_add(1, Ordering::SeqCst);
+            self.messages.lock().unwrap().push(message.clone());
+            Some(message)
+        }
+
+        async fn on_connect(&self) {
+            self.connect_count.fetch_add(1, Ordering::SeqCst);
+        }
+
+        async fn on_disconnect(&self) {
+            self.disconnect_count.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    #[derive(Debug)]
+    struct SelectiveHandler;
+
+    impl WebSocketHandler for SelectiveHandler {
+        async fn handle_message(&self, message: Value) -> Option<Value> {
+            if message.get("respond").is_some_and(|v| v.as_bool().unwrap_or(false)) {
+                Some(serde_json::json!({"response": "acknowledged"}))
+            } else {
+                None
+            }
+        }
+    }
+
+    #[derive(Debug)]
+    struct TransformHandler;
+
+    impl WebSocketHandler for TransformHandler {
+        async fn handle_message(&self, message: Value) -> Option<Value> {
+            message.as_object().map_or(None, |obj| {
+                let mut resp = obj.clone();
+                resp.insert("processed".to_string(), Value::Bool(true));
+                Some(Value::Object(resp))
+            })
+        }
+    }
+
+    #[test]
+    fn test_websocket_state_creation() {
+        let handler: EchoHandler = EchoHandler;
+        let state: WebSocketState<EchoHandler> = WebSocketState::new(handler);
+        let cloned: WebSocketState<EchoHandler> = state.clone();
+        assert!(Arc::ptr_eq(&state.handler, &cloned.handler));
+    }
+
+    #[test]
+    fn test_websocket_state_with_valid_schema() {
+        let handler: EchoHandler = EchoHandler;
+        let schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "type": {"type": "string"}
+            }
+        });
+
+        let result: Result<WebSocketState<EchoHandler>, String> =
+            WebSocketState::with_schemas(handler, Some(schema), None);
+        assert!(result.is_ok());
+    }
+
+    #[test]
+    fn test_websocket_state_with_invalid_schema() {
+        let handler: EchoHandler = EchoHandler;
+        let invalid_schema: serde_json::Value = serde_json::json!({
+            "type": "not_a_real_type",
+            "invalid": "schema"
+        });
+
+        let result: Result<WebSocketState<EchoHandler>, String> =
+            WebSocketState::with_schemas(handler, Some(invalid_schema), None);
+        assert!(result.is_err());
+        if let Err(error_msg) = result {
+            assert!(error_msg.contains("Invalid message schema"));
+        }
+    }
+
+    #[test]
+    fn test_websocket_state_with_both_schemas() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"action": {"type": "string"}}
+        });
+        let response_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"result": {"type": "string"}}
+        });
+
+        let result: Result<WebSocketState<EchoHandler>, String> =
+            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema));
+        assert!(result.is_ok());
+        let state: WebSocketState<EchoHandler> = result.unwrap();
+        assert!(state.message_schema.is_some());
+        assert!(state.response_schema.is_some());
+    }
+
+    #[test]
+    fn test_websocket_state_cloning_preserves_schemas() {
+        let handler: EchoHandler = EchoHandler;
+        let schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"id": {"type": "integer"}}
+        });
+
+        let state: WebSocketState<EchoHandler> = WebSocketState::with_schemas(handler, Some(schema), None).unwrap();
+        let cloned: WebSocketState<EchoHandler> = state.clone();
+
+        assert!(cloned.message_schema.is_some());
+        assert!(cloned.response_schema.is_none());
+        assert!(Arc::ptr_eq(&state.handler, &cloned.handler));
+    }
+
+    #[tokio::test]
+    async fn test_tracking_handler_lifecycle() {
+        let handler: TrackingHandler = TrackingHandler::new();
+        handler.on_connect().await;
+        assert_eq!(handler.connect_count.load(Ordering::SeqCst), 1);
+
+        let msg: Value = serde_json::json!({"test": "data"});
+        let _response: Option<Value> = handler.handle_message(msg).await;
+        assert_eq!(handler.message_count.load(Ordering::SeqCst), 1);
+
+        handler.on_disconnect().await;
+        assert_eq!(handler.disconnect_count.load(Ordering::SeqCst), 1);
+    }
+
+    #[tokio::test]
+    async fn test_selective_handler_responds_conditionally() {
+        let handler: SelectiveHandler = SelectiveHandler;
+
+        let respond_msg: Value = serde_json::json!({"respond": true});
+        let response1: Option<Value> = handler.handle_message(respond_msg).await;
+        assert!(response1.is_some());
+        assert_eq!(response1.unwrap(), serde_json::json!({"response": "acknowledged"}));
+
+        let no_respond_msg: Value = serde_json::json!({"respond": false});
+        let response2: Option<Value> = handler.handle_message(no_respond_msg).await;
+        assert!(response2.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_transform_handler_modifies_message() {
+        let handler: TransformHandler = TransformHandler;
+        let original: Value = serde_json::json!({"name": "test"});
+        let transformed: Option<Value> = handler.handle_message(original).await;
+
+        assert!(transformed.is_some());
+        let resp: Value = transformed.unwrap();
+        assert_eq!(resp.get("name").unwrap(), "test");
+        assert_eq!(resp.get("processed").unwrap(), true);
+    }
+
+    #[tokio::test]
+    async fn test_echo_handler_preserves_json_types() {
+        let handler: EchoHandler = EchoHandler;
+
+        let messages: Vec<Value> = vec![
+            serde_json::json!({"string": "value"}),
+            serde_json::json!({"number": 42}),
+            serde_json::json!({"float": 3.14}),
+            serde_json::json!({"bool": true}),
+            serde_json::json!({"null": null}),
+            serde_json::json!({"array": [1, 2, 3]}),
+        ];
+
+        for msg in messages {
+            let response: Option<Value> = handler.handle_message(msg.clone()).await;
+            assert!(response.is_some());
+            assert_eq!(response.unwrap(), msg);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_tracking_handler_accumulates_messages() {
+        let handler: TrackingHandler = TrackingHandler::new();
+
+        let messages: Vec<Value> = vec![
+            serde_json::json!({"id": 1}),
+            serde_json::json!({"id": 2}),
+            serde_json::json!({"id": 3}),
+        ];
+
+        for msg in messages {
+            let _: Option<Value> = handler.handle_message(msg).await;
+        }
+
+        assert_eq!(handler.message_count.load(Ordering::SeqCst), 3);
+        let stored: Vec<Value> = handler.messages.lock().unwrap().clone();
+        assert_eq!(stored.len(), 3);
+        assert_eq!(stored[0].get("id").unwrap(), 1);
+        assert_eq!(stored[1].get("id").unwrap(), 2);
+        assert_eq!(stored[2].get("id").unwrap(), 3);
+    }
+
+    #[tokio::test]
+    async fn test_echo_handler_with_nested_json() {
+        let handler: EchoHandler = EchoHandler;
+        let nested: Value = serde_json::json!({
+            "level1": {
+                "level2": {
+                    "level3": {
+                        "value": "deeply nested"
+                    }
+                }
+            }
+        });
+
+        let response: Option<Value> = handler.handle_message(nested.clone()).await;
+        assert!(response.is_some());
+        assert_eq!(response.unwrap(), nested);
+    }
+
+    #[tokio::test]
+    async fn test_echo_handler_with_large_array() {
+        let handler: EchoHandler = EchoHandler;
+        let large_array: Value = serde_json::json!({
+            "items": (0..1000).collect::<Vec<i32>>()
+        });
+
+        let response: Option<Value> = handler.handle_message(large_array.clone()).await;
+        assert!(response.is_some());
+        assert_eq!(response.unwrap(), large_array);
+    }
+
+    #[tokio::test]
+    async fn test_echo_handler_with_unicode() {
+        let handler: EchoHandler = EchoHandler;
+        let unicode_msg: Value = serde_json::json!({
+            "emoji": "🚀",
+            "chinese": "你好",
+            "arabic": "مرحبا",
+            "mixed": "Hello 世界 🌍"
+        });
+
+        let response: Option<Value> = handler.handle_message(unicode_msg.clone()).await;
+        assert!(response.is_some());
+        assert_eq!(response.unwrap(), unicode_msg);
+    }
+
+    #[test]
+    fn test_websocket_state_schemas_are_independent() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({"type": "object"});
+        let response_schema: serde_json::Value = serde_json::json!({"type": "array"});
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema)).unwrap();
+
+        let cloned: WebSocketState<EchoHandler> = state.clone();
+
+        assert!(state.message_schema.is_some());
+        assert!(state.response_schema.is_some());
+        assert!(cloned.message_schema.is_some());
+        assert!(cloned.response_schema.is_some());
+    }
+
+    #[test]
+    fn test_message_schema_validation_with_required_field() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"type": {"type": "string"}},
+            "required": ["type"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        assert!(state.message_schema.is_some());
+        assert!(state.response_schema.is_none());
+
+        let valid_msg: Value = serde_json::json!({"type": "test"});
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"other": "field"});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    #[test]
+    fn test_response_schema_validation_with_required_field() {
+        let handler: EchoHandler = EchoHandler;
+        let response_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"status": {"type": "string"}},
+            "required": ["status"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
+
+        assert!(state.message_schema.is_none());
+        assert!(state.response_schema.is_some());
+
+        let valid_response: Value = serde_json::json!({"status": "ok"});
+        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
+        assert!(validator.is_valid(&valid_response));
+
+        let invalid_response: Value = serde_json::json!({"other": "field"});
+        assert!(!validator.is_valid(&invalid_response));
+    }
+
+    #[test]
+    fn test_invalid_message_schema_returns_error() {
+        let handler: EchoHandler = EchoHandler;
+        let invalid_schema: serde_json::Value = serde_json::json!({
+            "type": "invalid_type_value",
+            "properties": {}
+        });
+
+        let result: Result<WebSocketState<EchoHandler>, String> =
+            WebSocketState::with_schemas(handler, Some(invalid_schema), None);
+
+        assert!(result.is_err());
+        match result {
+            Err(error_msg) => assert!(error_msg.contains("Invalid message schema")),
+            Ok(_) => panic!("Expected error but got Ok"),
+        }
+    }
+
+    #[test]
+    fn test_invalid_response_schema_returns_error() {
+        let handler: EchoHandler = EchoHandler;
+        let invalid_schema: serde_json::Value = serde_json::json!({
+            "type": "definitely_not_valid"
+        });
+
+        let result: Result<WebSocketState<EchoHandler>, String> =
+            WebSocketState::with_schemas(handler, None, Some(invalid_schema));
+
+        assert!(result.is_err());
+        match result {
+            Err(error_msg) => assert!(error_msg.contains("Invalid response schema")),
+            Ok(_) => panic!("Expected error but got Ok"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_handler_returning_none_response() {
+        let handler: SelectiveHandler = SelectiveHandler;
+
+        let no_response_msg: Value = serde_json::json!({"respond": false});
+        let result: Option<Value> = handler.handle_message(no_response_msg).await;
+
+        assert!(result.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_handler_with_complex_schema_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "user": {
+                    "type": "object",
+                    "properties": {
+                        "id": {"type": "integer"},
+                        "name": {"type": "string"}
+                    },
+                    "required": ["id", "name"]
+                },
+                "action": {"type": "string"}
+            },
+            "required": ["user", "action"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let valid_msg: Value = serde_json::json!({
+            "user": {"id": 123, "name": "Alice"},
+            "action": "create"
+        });
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({
+            "user": {"id": "not_an_int", "name": "Bob"},
+            "action": "create"
+        });
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    #[tokio::test]
+    async fn test_tracking_handler_with_multiple_message_types() {
+        let handler: TrackingHandler = TrackingHandler::new();
+
+        let messages: Vec<Value> = vec![
+            serde_json::json!({"type": "text", "content": "hello"}),
+            serde_json::json!({"type": "image", "url": "http://example.com/image.png"}),
+            serde_json::json!({"type": "video", "duration": 120}),
+        ];
+
+        for msg in messages {
+            let _: Option<Value> = handler.handle_message(msg).await;
+        }
+
+        assert_eq!(handler.message_count.load(Ordering::SeqCst), 3);
+        let stored: Vec<Value> = handler.messages.lock().unwrap().clone();
+        assert_eq!(stored.len(), 3);
+        assert_eq!(stored[0].get("type").unwrap(), "text");
+        assert_eq!(stored[1].get("type").unwrap(), "image");
+        assert_eq!(stored[2].get("type").unwrap(), "video");
+    }
+
+    #[tokio::test]
+    async fn test_selective_handler_with_explicit_false() {
+        let handler: SelectiveHandler = SelectiveHandler;
+
+        let msg: Value = serde_json::json!({"respond": false, "data": "test"});
+        let response: Option<Value> = handler.handle_message(msg).await;
+
+        assert!(response.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_selective_handler_without_respond_field() {
+        let handler: SelectiveHandler = SelectiveHandler;
+
+        let msg: Value = serde_json::json!({"data": "test"});
+        let response: Option<Value> = handler.handle_message(msg).await;
+
+        assert!(response.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_transform_handler_with_empty_object() {
+        let handler: TransformHandler = TransformHandler;
+        let original: Value = serde_json::json!({});
+        let transformed: Option<Value> = handler.handle_message(original).await;
+
+        assert!(transformed.is_some());
+        let resp: Value = transformed.unwrap();
+        assert_eq!(resp.get("processed").unwrap(), true);
+        assert_eq!(resp.as_object().unwrap().len(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_transform_handler_preserves_all_fields() {
+        let handler: TransformHandler = TransformHandler;
+        let original: Value = serde_json::json!({
+            "field1": "value1",
+            "field2": 42,
+            "field3": true,
+            "nested": {"key": "value"}
+        });
+        let transformed: Option<Value> = handler.handle_message(original.clone()).await;
+
+        assert!(transformed.is_some());
+        let resp: Value = transformed.unwrap();
+        assert_eq!(resp.get("field1").unwrap(), "value1");
+        assert_eq!(resp.get("field2").unwrap(), 42);
+        assert_eq!(resp.get("field3").unwrap(), true);
+        assert_eq!(resp.get("nested").unwrap(), &serde_json::json!({"key": "value"}));
+        assert_eq!(resp.get("processed").unwrap(), true);
+    }
+
+    #[tokio::test]
+    async fn test_transform_handler_with_non_object_input() {
+        let handler: TransformHandler = TransformHandler;
+
+        let array: Value = serde_json::json!([1, 2, 3]);
+        let response1: Option<Value> = handler.handle_message(array).await;
+        assert!(response1.is_none());
+
+        let string: Value = serde_json::json!("not an object");
+        let response2: Option<Value> = handler.handle_message(string).await;
+        assert!(response2.is_none());
+
+        let number: Value = serde_json::json!(42);
+        let response3: Option<Value> = handler.handle_message(number).await;
+        assert!(response3.is_none());
+    }
+
+    /// Test message validation failure with schema constraint
+    #[test]
+    fn test_message_schema_rejects_wrong_type() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"id": {"type": "integer"}},
+            "required": ["id"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let invalid_msg: Value = serde_json::json!({"id": "not_an_integer"});
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test response schema validation failure
+    #[test]
+    fn test_response_schema_rejects_invalid_type() {
+        let handler: EchoHandler = EchoHandler;
+        let response_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"count": {"type": "integer"}},
+            "required": ["count"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
+
+        let invalid_response: Value = serde_json::json!([1, 2, 3]);
+        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
+        assert!(!validator.is_valid(&invalid_response));
+    }
+
+    /// Test message with multiple required fields missing
+    #[test]
+    fn test_message_missing_multiple_required_fields() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "user_id": {"type": "integer"},
+                "action": {"type": "string"},
+                "timestamp": {"type": "string"}
+            },
+            "required": ["user_id", "action", "timestamp"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let invalid_msg: Value = serde_json::json!({"other": "value"});
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        assert!(!validator.is_valid(&invalid_msg));
+
+        let partial_msg: Value = serde_json::json!({"user_id": 123});
+        assert!(!validator.is_valid(&partial_msg));
+    }
+
+    /// Test deeply nested schema validation with required nested properties
+    #[test]
+    fn test_deeply_nested_schema_validation_failure() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "metadata": {
+                    "type": "object",
+                    "properties": {
+                        "request": {
+                            "type": "object",
+                            "properties": {
+                                "id": {"type": "string"}
+                            },
+                            "required": ["id"]
+                        }
+                    },
+                    "required": ["request"]
+                }
+            },
+            "required": ["metadata"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let invalid_msg: Value = serde_json::json!({
+            "metadata": {
+                "request": {}
+            }
+        });
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test array property validation with items constraint
+    #[test]
+    fn test_array_property_type_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "ids": {
+                    "type": "array",
+                    "items": {"type": "integer"}
+                }
+            }
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"ids": [1, 2, 3]});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"ids": [1, "two", 3]});
+        assert!(!validator.is_valid(&invalid_msg));
+
+        let invalid_msg2: Value = serde_json::json!({"ids": "not_an_array"});
+        assert!(!validator.is_valid(&invalid_msg2));
+    }
+
+    /// Test enum/const property validation
+    #[test]
+    fn test_enum_property_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "status": {
+                    "type": "string",
+                    "enum": ["pending", "active", "completed"]
+                }
+            },
+            "required": ["status"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"status": "active"});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"status": "unknown"});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test minimum/maximum constraints on numbers
+    #[test]
+    fn test_number_range_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "age": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "maximum": 150
+                }
+            },
+            "required": ["age"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"age": 25});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"age": -1});
+        assert!(!validator.is_valid(&invalid_msg));
+
+        let invalid_msg2: Value = serde_json::json!({"age": 200});
+        assert!(!validator.is_valid(&invalid_msg2));
+    }
+
+    /// Test string length constraints
+    #[test]
+    fn test_string_length_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "username": {
+                    "type": "string",
+                    "minLength": 3,
+                    "maxLength": 20
+                }
+            },
+            "required": ["username"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"username": "alice"});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"username": "ab"});
+        assert!(!validator.is_valid(&invalid_msg));
+
+        let invalid_msg2: Value =
+            serde_json::json!({"username": "this_is_a_very_long_username_over_twenty_characters"});
+        assert!(!validator.is_valid(&invalid_msg2));
+    }
+
+    /// Test pattern (regex) validation
+    #[test]
+    fn test_pattern_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "email": {
+                    "type": "string",
+                    "pattern": "^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}$"
+                }
+            },
+            "required": ["email"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"email": "user@example.com"});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"email": "user@example"});
+        assert!(!validator.is_valid(&invalid_msg));
+
+        let invalid_msg2: Value = serde_json::json!({"email": "userexample.com"});
+        assert!(!validator.is_valid(&invalid_msg2));
+    }
+
+    /// Test additionalProperties constraint
+    #[test]
+    fn test_additional_properties_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "name": {"type": "string"}
+            },
+            "additionalProperties": false
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"name": "Alice"});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"name": "Bob", "age": 30});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test oneOf constraint (mutually exclusive properties)
+    #[test]
+    fn test_one_of_constraint() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "oneOf": [
+                {
+                    "properties": {"type": {"const": "text"}},
+                    "required": ["type"]
+                },
+                {
+                    "properties": {"type": {"const": "number"}},
+                    "required": ["type"]
+                }
+            ]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"type": "text"});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"type": "unknown"});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test anyOf constraint (at least one match)
+    #[test]
+    fn test_any_of_constraint() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "value": {"type": ["string", "integer"]}
+            },
+            "required": ["value"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let msg1: Value = serde_json::json!({"value": "text"});
+        assert!(validator.is_valid(&msg1));
+
+        let msg2: Value = serde_json::json!({"value": 42});
+        assert!(validator.is_valid(&msg2));
+
+        let invalid_msg: Value = serde_json::json!({"value": true});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test response validation with complex constraints
+    #[test]
+    fn test_response_schema_with_multiple_constraints() {
+        let handler: EchoHandler = EchoHandler;
+        let response_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "success": {"type": "boolean"},
+                "data": {
+                    "type": "object",
+                    "properties": {
+                        "items": {
+                            "type": "array",
+                            "items": {"type": "object"},
+                            "minItems": 1
+                        }
+                    },
+                    "required": ["items"]
+                }
+            },
+            "required": ["success", "data"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, None, Some(response_schema)).unwrap();
+
+        let validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
+
+        let valid_response: Value = serde_json::json!({
+            "success": true,
+            "data": {
+                "items": [{"id": 1}]
+            }
+        });
+        assert!(validator.is_valid(&valid_response));
+
+        let invalid_response: Value = serde_json::json!({
+            "success": true,
+            "data": {
+                "items": []
+            }
+        });
+        assert!(!validator.is_valid(&invalid_response));
+
+        let invalid_response2: Value = serde_json::json!({
+            "success": true
+        });
+        assert!(!validator.is_valid(&invalid_response2));
+    }
+
+    /// Test null type validation
+    #[test]
+    fn test_null_value_validation() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "optional_field": {"type": ["string", "null"]},
+                "required_field": {"type": "string"}
+            },
+            "required": ["required_field"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let msg1: Value = serde_json::json!({
+            "optional_field": null,
+            "required_field": "value"
+        });
+        assert!(validator.is_valid(&msg1));
+
+        let msg2: Value = serde_json::json!({"required_field": "value"});
+        assert!(validator.is_valid(&msg2));
+
+        let invalid_msg: Value = serde_json::json!({"required_field": null});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+
+    /// Test schema with default values (they don't change validation)
+    #[test]
+    fn test_schema_with_defaults_still_validates() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "status": {
+                    "type": "string",
+                    "default": "pending"
+                }
+            }
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let msg: Value = serde_json::json!({});
+        assert!(validator.is_valid(&msg));
+    }
+
+    /// Test both message and response schema validation together
+    #[test]
+    fn test_both_schemas_validate_independently() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"action": {"type": "string"}},
+            "required": ["action"]
+        });
+        let response_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {"result": {"type": "string"}},
+            "required": ["result"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), Some(response_schema)).unwrap();
+
+        let msg_validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+        let resp_validator: &jsonschema::Validator = state.response_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"action": "test"});
+        let invalid_response: Value = serde_json::json!({"data": "oops"});
+
+        assert!(msg_validator.is_valid(&valid_msg));
+        assert!(!resp_validator.is_valid(&invalid_response));
+
+        let invalid_msg: Value = serde_json::json!({"data": "oops"});
+        let valid_response: Value = serde_json::json!({"result": "ok"});
+
+        assert!(!msg_validator.is_valid(&invalid_msg));
+        assert!(resp_validator.is_valid(&valid_response));
+    }
+
+    /// Test validation with very long/large payload
+    #[test]
+    fn test_validation_with_large_payload() {
+        let handler: EchoHandler = EchoHandler;
+        let message_schema: serde_json::Value = serde_json::json!({
+            "type": "object",
+            "properties": {
+                "items": {
+                    "type": "array",
+                    "items": {"type": "integer"}
+                }
+            },
+            "required": ["items"]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let mut items = Vec::new();
+        for i in 0..10_000 {
+            items.push(i);
+        }
+        let large_msg: Value = serde_json::json!({"items": items});
+
+        assert!(validator.is_valid(&large_msg));
+    }
+
+    /// Test validation error doesn't panic with invalid schema combinations
+    #[test]
+    fn test_mutually_exclusive_schema_properties() {
+        let handler: EchoHandler = EchoHandler;
+
+        let message_schema: serde_json::Value = serde_json::json!({
+            "allOf": [
+                {
+                    "type": "object",
+                    "properties": {"a": {"type": "string"}},
+                    "required": ["a"]
+                },
+                {
+                    "type": "object",
+                    "properties": {"b": {"type": "integer"}},
+                    "required": ["b"]
+                }
+            ]
+        });
+
+        let state: WebSocketState<EchoHandler> =
+            WebSocketState::with_schemas(handler, Some(message_schema), None).unwrap();
+
+        let validator: &jsonschema::Validator = state.message_schema.as_ref().unwrap();
+
+        let valid_msg: Value = serde_json::json!({"a": "text", "b": 42});
+        assert!(validator.is_valid(&valid_msg));
+
+        let invalid_msg: Value = serde_json::json!({"a": "text"});
+        assert!(!validator.is_valid(&invalid_msg));
+    }
+}