Add WebSocket gateway and control plane (#8)

* Add WebSocket gateway and control plane endpoint

Adds bidirectional WebSocket transport to the web gateway alongside
the existing SSE stream. Clients can send messages, approvals, and
pings over a single persistent connection at /api/chat/ws.

- Enable axum `ws` feature for built-in WebSocket support
- Add WsClientMessage/WsServerMessage types with tagged JSON protocol
- Add subscribe_raw() to SseManager for non-SSE consumers
- Create ws.rs with connection handler (split sender/receiver tasks)
- Add WsConnectionTracker for active connection counting
- Add /api/gateway/status control plane endpoint (SSE + WS counts)
- 35 new tests covering message types, broadcast, and handler logic

https://claude.ai/code/session_01KEaLN6Xq2j5EeV3SGHQT6b

* Add e2e WebSocket gateway integration tests

- Add tokio-tungstenite dev-dependency for WebSocket client in tests
- Update start_server to return actual bound SocketAddr (enables port 0)
- Add 10 e2e tests covering full HTTP upgrade → WebSocket → message flow:
  ping/pong, message routing to agent, broadcast event delivery,
  connection tracking, invalid message handling, auth rejection,
  gateway status endpoint, and multi-event sequencing

https://claude.ai/code/session_01KEaLN6Xq2j5EeV3SGHQT6b

---------

Co-authored-by: Claude <[email protected]>

src/channels/web/mod.rs

@@ -8,6 +8,7 @@
 //! ```text
 //! Browser ─── POST /api/chat/send ──► Agent Loop
 //!         ◄── GET  /api/chat/events ── SSE stream
+//!         ─── GET  /api/chat/ws ─────► WebSocket (bidirectional)
 //!         ─── GET  /api/memory/* ────► Workspace
 //!         ─── GET  /api/jobs/* ──────► ContextManager
 //!         ◄── GET  / ───────────────── Static HTML/CSS/JS
@@ -18,6 +19,7 @@ pub mod log_layer;
 pub mod server;
 pub mod sse;
 pub mod types;
+pub mod ws;
 
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -75,6 +77,7 @@ impl GatewayChannel {
             tool_registry: None,
             user_id: config.user_id.clone(),
             shutdown_tx: tokio::sync::RwLock::new(None),
+            ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())),
         });
 
         Self {
@@ -97,6 +100,7 @@ impl GatewayChannel {
             tool_registry: self.state.tool_registry.clone(),
             user_id: self.state.user_id.clone(),
             shutdown_tx: tokio::sync::RwLock::new(None),
+            ws_tracker: self.state.ws_tracker.clone(),
         };
         mutate(&mut new_state);
         self.state = Arc::new(new_state);
@@ -169,9 +173,10 @@ impl Channel for GatewayChannel {
                 ),
             })?;
 
-        server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
+        let bound_addr =
+            server::start_server(addr, self.state.clone(), self.auth_token.clone()).await?;
 
-        tracing::info!("Web gateway listening on http://{}", addr);
+        tracing::info!("Web gateway listening on http://{}", bound_addr);
         tracing::info!("Auth token: {}", self.auth_token);
 
         Ok(Box::pin(ReceiverStream::new(rx)))

src/channels/web/server.rs

@@ -8,7 +8,7 @@ use std::sync::Arc;
 
 use axum::{
     Json, Router,
-    extract::{Path, Query, State},
+    extract::{Path, Query, State, WebSocketUpgrade},
     http::{StatusCode, header},
     middleware,
     response::{
@@ -55,20 +55,31 @@ pub struct GatewayState {
     pub user_id: String,
     /// Shutdown signal sender.
     pub shutdown_tx: tokio::sync::RwLock<Option<oneshot::Sender<()>>>,
+    /// WebSocket connection tracker.
+    pub ws_tracker: Option<Arc<crate::channels::web::ws::WsConnectionTracker>>,
 }
 
 /// Start the gateway HTTP server.
+///
+/// Returns the actual bound `SocketAddr` (useful when binding to port 0).
 pub async fn start_server(
     addr: SocketAddr,
     state: Arc<GatewayState>,
     auth_token: String,
-) -> Result<(), crate::error::ChannelError> {
+) -> Result<SocketAddr, crate::error::ChannelError> {
     let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
         crate::error::ChannelError::StartupFailed {
             name: "gateway".to_string(),
             reason: format!("Failed to bind to {}: {}", addr, e),
         }
     })?;
+    let bound_addr =
+        listener
+            .local_addr()
+            .map_err(|e| crate::error::ChannelError::StartupFailed {
+                name: "gateway".to_string(),
+                reason: format!("Failed to get local addr: {}", e),
+            })?;
 
     // Public routes (no auth)
     let public = Router::new().route("/api/health", get(health_handler));
@@ -80,6 +91,7 @@ pub async fn start_server(
         .route("/api/chat/send", post(chat_send_handler))
         .route("/api/chat/approval", post(chat_approval_handler))
         .route("/api/chat/events", get(chat_events_handler))
+        .route("/api/chat/ws", get(chat_ws_handler))
         .route("/api/chat/history", get(chat_history_handler))
         .route("/api/chat/threads", get(chat_threads_handler))
         .route("/api/chat/thread/new", post(chat_new_thread_handler))
@@ -108,6 +120,8 @@ pub async fn start_server(
             "/api/extensions/{name}/remove",
             post(extensions_remove_handler),
         )
+        // Gateway control plane
+        .route("/api/gateway/status", get(gateway_status_handler))
         .route_layer(middleware::from_fn_with_state(auth_state, auth_middleware));
 
     // Static file routes (no auth, served from embedded strings)
@@ -137,7 +151,7 @@ pub async fn start_server(
         }
     });
 
-    Ok(())
+    Ok(bound_addr)
 }
 
 // --- Static file handlers ---
@@ -272,6 +286,13 @@ async fn chat_events_handler(State(state): State<Arc<GatewayState>>) -> impl Int
     state.sse.subscribe()
 }
 
+async fn chat_ws_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<Arc<GatewayState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| crate::channels::web::ws::handle_ws_connection(socket, state))
+}
+
 #[derive(Deserialize)]
 struct HistoryQuery {
     thread_id: Option<String>,
@@ -834,3 +855,29 @@ async fn extensions_remove_handler(
         Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))),
     }
 }
+
+// --- Gateway control plane handlers ---
+
+async fn gateway_status_handler(
+    State(state): State<Arc<GatewayState>>,
+) -> Json<GatewayStatusResponse> {
+    let sse_connections = state.sse.connection_count();
+    let ws_connections = state
+        .ws_tracker
+        .as_ref()
+        .map(|t| t.connection_count())
+        .unwrap_or(0);
+
+    Json(GatewayStatusResponse {
+        sse_connections,
+        ws_connections,
+        total_connections: sse_connections + ws_connections,
+    })
+}
+
+#[derive(serde::Serialize)]
+struct GatewayStatusResponse {
+    sse_connections: u64,
+    ws_connections: u64,
+    total_connections: u64,
+}

src/channels/web/sse.rs

@@ -41,6 +41,23 @@ impl SseManager {
         self.connection_count.load(Ordering::Relaxed)
     }
 
+    /// Create a raw broadcast subscription for non-SSE consumers (e.g. WebSocket).
+    ///
+    /// Returns a stream of `SseEvent` values and increments/decrements the
+    /// connection counter on creation/drop, just like `subscribe()` does for SSE.
+    pub fn subscribe_raw(&self) -> impl Stream<Item = SseEvent> + Send + 'static + use<> {
+        let counter = Arc::clone(&self.connection_count);
+        counter.fetch_add(1, Ordering::Relaxed);
+        let rx = self.tx.subscribe();
+
+        let stream = BroadcastStream::new(rx).filter_map(|result| result.ok());
+
+        CountedStream {
+            inner: stream,
+            counter,
+        }
+    }
+
     /// Create a new SSE stream for a client connection.
     pub fn subscribe(
         &self,
@@ -144,4 +161,53 @@ mod tests {
             _ => panic!("unexpected event type"),
         }
     }
+
+    #[tokio::test]
+    async fn test_subscribe_raw_receives_events() {
+        let manager = SseManager::new();
+        let mut stream = Box::pin(manager.subscribe_raw());
+
+        assert_eq!(manager.connection_count(), 1);
+
+        manager.broadcast(SseEvent::Thinking {
+            message: "working".to_string(),
+        });
+
+        let event = stream.next().await.unwrap();
+        match event {
+            SseEvent::Thinking { message } => assert_eq!(message, "working"),
+            _ => panic!("Expected Thinking event"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_subscribe_raw_decrements_on_drop() {
+        let manager = SseManager::new();
+        {
+            let _stream = Box::pin(manager.subscribe_raw());
+            assert_eq!(manager.connection_count(), 1);
+        }
+        // Stream dropped, counter should decrement
+        assert_eq!(manager.connection_count(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_subscribe_raw_multiple_subscribers() {
+        let manager = SseManager::new();
+        let mut s1 = Box::pin(manager.subscribe_raw());
+        let mut s2 = Box::pin(manager.subscribe_raw());
+        assert_eq!(manager.connection_count(), 2);
+
+        manager.broadcast(SseEvent::Heartbeat);
+
+        let e1 = s1.next().await.unwrap();
+        let e2 = s2.next().await.unwrap();
+        assert!(matches!(e1, SseEvent::Heartbeat));
+        assert!(matches!(e2, SseEvent::Heartbeat));
+
+        drop(s1);
+        assert_eq!(manager.connection_count(), 1);
+        drop(s2);
+        assert_eq!(manager.connection_count(), 0);
+    }
 }

src/channels/web/types.rs

@@ -260,6 +260,72 @@ impl ActionResponse {
     }
 }
 
+// --- WebSocket ---
+
+/// Message sent by a WebSocket client to the server.
+#[derive(Debug, Clone, Deserialize)]
+#[serde(tag = "type")]
+pub enum WsClientMessage {
+    /// Send a chat message to the agent.
+    #[serde(rename = "message")]
+    Message {
+        content: String,
+        thread_id: Option<String>,
+    },
+    /// Approve or deny a pending tool execution.
+    #[serde(rename = "approval")]
+    Approval {
+        request_id: String,
+        /// "approve", "always", or "deny"
+        action: String,
+    },
+    /// Client heartbeat ping.
+    #[serde(rename = "ping")]
+    Ping,
+}
+
+/// Message sent by the server to a WebSocket client.
+#[derive(Debug, Clone, Serialize)]
+#[serde(tag = "type")]
+pub enum WsServerMessage {
+    /// An SSE-style event forwarded over WebSocket.
+    #[serde(rename = "event")]
+    Event {
+        /// The event sub-type (response, thinking, tool_started, etc.)
+        event_type: String,
+        /// The event payload as a JSON value.
+        data: serde_json::Value,
+    },
+    /// Server heartbeat pong.
+    #[serde(rename = "pong")]
+    Pong,
+    /// Error message.
+    #[serde(rename = "error")]
+    Error { message: String },
+}
+
+impl WsServerMessage {
+    /// Create a WsServerMessage from an SseEvent.
+    pub fn from_sse_event(event: &SseEvent) -> Self {
+        let event_type = match event {
+            SseEvent::Response { .. } => "response",
+            SseEvent::Thinking { .. } => "thinking",
+            SseEvent::ToolStarted { .. } => "tool_started",
+            SseEvent::ToolCompleted { .. } => "tool_completed",
+            SseEvent::StreamChunk { .. } => "stream_chunk",
+            SseEvent::Status { .. } => "status",
+            SseEvent::ApprovalNeeded { .. } => "approval_needed",
+            SseEvent::Error { .. } => "error",
+            SseEvent::Heartbeat => "heartbeat",
+        };
+        let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
+        WsServerMessage::Event {
+            event_type: event_type.to_string(),
+            data,
+        }
+    }
+}
+
 // --- Health ---
 
 #[derive(Debug, Serialize)]
@@ -267,3 +333,145 @@ pub struct HealthResponse {
     pub status: &'static str,
     pub channel: &'static str,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ---- WsClientMessage deserialization tests ----
+
+    #[test]
+    fn test_ws_client_message_parse() {
+        let json = r#"{"type":"message","content":"hello","thread_id":"t1"}"#;
+        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
+        match msg {
+            WsClientMessage::Message { content, thread_id } => {
+                assert_eq!(content, "hello");
+                assert_eq!(thread_id.as_deref(), Some("t1"));
+            }
+            _ => panic!("Expected Message variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_client_message_no_thread() {
+        let json = r#"{"type":"message","content":"hi"}"#;
+        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
+        match msg {
+            WsClientMessage::Message { content, thread_id } => {
+                assert_eq!(content, "hi");
+                assert!(thread_id.is_none());
+            }
+            _ => panic!("Expected Message variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_client_approval_parse() {
+        let json = r#"{"type":"approval","request_id":"abc-123","action":"approve"}"#;
+        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
+        match msg {
+            WsClientMessage::Approval { request_id, action } => {
+                assert_eq!(request_id, "abc-123");
+                assert_eq!(action, "approve");
+            }
+            _ => panic!("Expected Approval variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_client_ping_parse() {
+        let json = r#"{"type":"ping"}"#;
+        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
+        assert!(matches!(msg, WsClientMessage::Ping));
+    }
+
+    #[test]
+    fn test_ws_client_unknown_type_fails() {
+        let json = r#"{"type":"unknown"}"#;
+        let result: Result<WsClientMessage, _> = serde_json::from_str(json);
+        assert!(result.is_err());
+    }
+
+    // ---- WsServerMessage serialization tests ----
+
+    #[test]
+    fn test_ws_server_pong_serialize() {
+        let msg = WsServerMessage::Pong;
+        let json = serde_json::to_string(&msg).unwrap();
+        assert_eq!(json, r#"{"type":"pong"}"#);
+    }
+
+    #[test]
+    fn test_ws_server_error_serialize() {
+        let msg = WsServerMessage::Error {
+            message: "bad request".to_string(),
+        };
+        let json = serde_json::to_string(&msg).unwrap();
+        let parsed: serde_json::Value = serde_json::from_str(&json).unwrap();
+        assert_eq!(parsed["type"], "error");
+        assert_eq!(parsed["message"], "bad request");
+    }
+
+    #[test]
+    fn test_ws_server_from_sse_response() {
+        let sse = SseEvent::Response {
+            content: "hello".to_string(),
+            thread_id: "t1".to_string(),
+        };
+        let ws = WsServerMessage::from_sse_event(&sse);
+        match ws {
+            WsServerMessage::Event { event_type, data } => {
+                assert_eq!(event_type, "response");
+                assert_eq!(data["content"], "hello");
+                assert_eq!(data["thread_id"], "t1");
+            }
+            _ => panic!("Expected Event variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_server_from_sse_thinking() {
+        let sse = SseEvent::Thinking {
+            message: "reasoning...".to_string(),
+        };
+        let ws = WsServerMessage::from_sse_event(&sse);
+        match ws {
+            WsServerMessage::Event { event_type, data } => {
+                assert_eq!(event_type, "thinking");
+                assert_eq!(data["message"], "reasoning...");
+            }
+            _ => panic!("Expected Event variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_server_from_sse_approval_needed() {
+        let sse = SseEvent::ApprovalNeeded {
+            request_id: "r1".to_string(),
+            tool_name: "shell".to_string(),
+            description: "Run ls".to_string(),
+            parameters: "{}".to_string(),
+        };
+        let ws = WsServerMessage::from_sse_event(&sse);
+        match ws {
+            WsServerMessage::Event { event_type, data } => {
+                assert_eq!(event_type, "approval_needed");
+                assert_eq!(data["tool_name"], "shell");
+            }
+            _ => panic!("Expected Event variant"),
+        }
+    }
+
+    #[test]
+    fn test_ws_server_from_sse_heartbeat() {
+        let sse = SseEvent::Heartbeat;
+        let ws = WsServerMessage::from_sse_event(&sse);
+        match ws {
+            WsServerMessage::Event { event_type, .. } => {
+                assert_eq!(event_type, "heartbeat");
+            }
+            _ => panic!("Expected Event variant"),
+        }
+    }
+}

src/channels/web/ws.rs

@@ -0,0 +1,411 @@
+//! WebSocket handler for bidirectional client communication.
+//!
+//! Provides the same event stream as SSE but also accepts incoming messages
+//! (chat, approvals) over a single persistent connection.
+//!
+//! ```text
+//! Client ──── WS frame: {"type":"message","content":"hello"} ──► Agent Loop
+//!        ◄─── WS frame: {"type":"event","event_type":"response","data":{...}} ── Broadcast
+//!        ──── WS frame: {"type":"ping"} ──────────────────────────────────────►
+//!        ◄─── WS frame: {"type":"pong"} ──────────────────────────────────────
+//! ```
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use axum::extract::ws::{Message, WebSocket};
+use futures::{SinkExt, StreamExt};
+use tokio::sync::mpsc;
+use uuid::Uuid;
+
+use crate::agent::submission::Submission;
+use crate::channels::IncomingMessage;
+use crate::channels::web::server::GatewayState;
+use crate::channels::web::types::{WsClientMessage, WsServerMessage};
+
+/// Tracks active WebSocket connections.
+pub struct WsConnectionTracker {
+    count: AtomicU64,
+}
+
+impl WsConnectionTracker {
+    pub fn new() -> Self {
+        Self {
+            count: AtomicU64::new(0),
+        }
+    }
+
+    pub fn connection_count(&self) -> u64 {
+        self.count.load(Ordering::Relaxed)
+    }
+
+    fn increment(&self) {
+        self.count.fetch_add(1, Ordering::Relaxed);
+    }
+
+    fn decrement(&self) {
+        self.count.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+impl Default for WsConnectionTracker {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Handle an upgraded WebSocket connection.
+///
+/// Spawns two tasks:
+/// - **sender**: forwards broadcast events to the WebSocket client
+/// - **receiver**: reads client frames and routes them to the agent
+///
+/// When either task ends (client disconnect or broadcast closed), both are
+/// cleaned up.
+pub async fn handle_ws_connection(socket: WebSocket, state: Arc<GatewayState>) {
+    let (mut ws_sink, mut ws_stream) = socket.split();
+
+    // Track connection
+    if let Some(ref tracker) = state.ws_tracker {
+        tracker.increment();
+    }
+    let tracker_for_drop = state.ws_tracker.clone();
+
+    // Subscribe to broadcast events (same source as SSE)
+    let mut event_stream = Box::pin(state.sse.subscribe_raw());
+
+    // Channel for the sender task to receive messages from both
+    // the broadcast stream and any direct sends (like Pong)
+    let (direct_tx, mut direct_rx) = mpsc::channel::<WsServerMessage>(64);
+
+    // Sender task: forward broadcast events + direct messages to WS client
+    let sender_handle = tokio::spawn(async move {
+        loop {
+            let msg = tokio::select! {
+                event = event_stream.next() => {
+                    match event {
+                        Some(sse_event) => WsServerMessage::from_sse_event(&sse_event),
+                        None => break, // Broadcast channel closed
+                    }
+                }
+                direct = direct_rx.recv() => {
+                    match direct {
+                        Some(msg) => msg,
+                        None => break, // Direct channel closed
+                    }
+                }
+            };
+
+            let json = match serde_json::to_string(&msg) {
+                Ok(j) => j,
+                Err(_) => continue,
+            };
+
+            if ws_sink.send(Message::Text(json.into())).await.is_err() {
+                break; // Client disconnected
+            }
+        }
+    });
+
+    // Receiver task: read client frames and route to agent
+    let user_id = state.user_id.clone();
+    while let Some(Ok(frame)) = ws_stream.next().await {
+        match frame {
+            Message::Text(text) => {
+                let parsed: Result<WsClientMessage, _> = serde_json::from_str(&text);
+                match parsed {
+                    Ok(client_msg) => {
+                        handle_client_message(client_msg, &state, &user_id, &direct_tx).await;
+                    }
+                    Err(e) => {
+                        let _ = direct_tx
+                            .send(WsServerMessage::Error {
+                                message: format!("Invalid message: {}", e),
+                            })
+                            .await;
+                    }
+                }
+            }
+            Message::Close(_) => break,
+            // Ignore binary, ping/pong (axum handles protocol-level pings)
+            _ => {}
+        }
+    }
+
+    // Clean up: abort sender, decrement counter
+    sender_handle.abort();
+    if let Some(ref tracker) = tracker_for_drop {
+        tracker.decrement();
+    }
+}
+
+/// Route a parsed client message to the appropriate handler.
+async fn handle_client_message(
+    msg: WsClientMessage,
+    state: &GatewayState,
+    user_id: &str,
+    direct_tx: &mpsc::Sender<WsServerMessage>,
+) {
+    match msg {
+        WsClientMessage::Message { content, thread_id } => {
+            let mut incoming = IncomingMessage::new("gateway", user_id, &content);
+            if let Some(ref tid) = thread_id {
+                incoming = incoming.with_thread(tid);
+            }
+
+            let tx_guard = state.msg_tx.read().await;
+            if let Some(ref tx) = *tx_guard {
+                if tx.send(incoming).await.is_err() {
+                    let _ = direct_tx
+                        .send(WsServerMessage::Error {
+                            message: "Channel closed".to_string(),
+                        })
+                        .await;
+                }
+            } else {
+                let _ = direct_tx
+                    .send(WsServerMessage::Error {
+                        message: "Channel not started".to_string(),
+                    })
+                    .await;
+            }
+        }
+        WsClientMessage::Approval { request_id, action } => {
+            let (approved, always) = match action.as_str() {
+                "approve" => (true, false),
+                "always" => (true, true),
+                "deny" => (false, false),
+                other => {
+                    let _ = direct_tx
+                        .send(WsServerMessage::Error {
+                            message: format!("Unknown approval action: {}", other),
+                        })
+                        .await;
+                    return;
+                }
+            };
+
+            let request_uuid = match Uuid::parse_str(&request_id) {
+                Ok(id) => id,
+                Err(_) => {
+                    let _ = direct_tx
+                        .send(WsServerMessage::Error {
+                            message: "Invalid request_id (expected UUID)".to_string(),
+                        })
+                        .await;
+                    return;
+                }
+            };
+
+            let approval = Submission::ExecApproval {
+                request_id: request_uuid,
+                approved,
+                always,
+            };
+            let content = match serde_json::to_string(&approval) {
+                Ok(c) => c,
+                Err(e) => {
+                    let _ = direct_tx
+                        .send(WsServerMessage::Error {
+                            message: format!("Failed to serialize approval: {}", e),
+                        })
+                        .await;
+                    return;
+                }
+            };
+
+            let msg = IncomingMessage::new("gateway", user_id, content);
+            let tx_guard = state.msg_tx.read().await;
+            if let Some(ref tx) = *tx_guard {
+                let _ = tx.send(msg).await;
+            }
+        }
+        WsClientMessage::Ping => {
+            let _ = direct_tx.send(WsServerMessage::Pong).await;
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_ws_connection_tracker() {
+        let tracker = WsConnectionTracker::new();
+        assert_eq!(tracker.connection_count(), 0);
+
+        tracker.increment();
+        assert_eq!(tracker.connection_count(), 1);
+
+        tracker.increment();
+        assert_eq!(tracker.connection_count(), 2);
+
+        tracker.decrement();
+        assert_eq!(tracker.connection_count(), 1);
+
+        tracker.decrement();
+        assert_eq!(tracker.connection_count(), 0);
+    }
+
+    #[test]
+    fn test_ws_connection_tracker_default() {
+        let tracker = WsConnectionTracker::default();
+        assert_eq!(tracker.connection_count(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_message_ping() {
+        // Ping should produce a Pong on the direct channel
+        let (direct_tx, mut direct_rx) = mpsc::channel(16);
+        let state = make_test_state(None).await;
+
+        handle_client_message(WsClientMessage::Ping, &state, "user1", &direct_tx).await;
+
+        let response = direct_rx.recv().await.unwrap();
+        assert!(matches!(response, WsServerMessage::Pong));
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_message_sends_to_agent() {
+        // A Message should be forwarded to the agent's msg_tx
+        let (agent_tx, mut agent_rx) = mpsc::channel(16);
+        let state = make_test_state(Some(agent_tx)).await;
+        let (direct_tx, _direct_rx) = mpsc::channel(16);
+
+        handle_client_message(
+            WsClientMessage::Message {
+                content: "hello agent".to_string(),
+                thread_id: Some("t1".to_string()),
+            },
+            &state,
+            "user1",
+            &direct_tx,
+        )
+        .await;
+
+        let incoming = agent_rx.recv().await.unwrap();
+        assert_eq!(incoming.content, "hello agent");
+        assert_eq!(incoming.thread_id.as_deref(), Some("t1"));
+        assert_eq!(incoming.channel, "gateway");
+        assert_eq!(incoming.user_id, "user1");
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_message_no_channel() {
+        // When msg_tx is None, should send an error back
+        let state = make_test_state(None).await;
+        let (direct_tx, mut direct_rx) = mpsc::channel(16);
+
+        handle_client_message(
+            WsClientMessage::Message {
+                content: "hello".to_string(),
+                thread_id: None,
+            },
+            &state,
+            "user1",
+            &direct_tx,
+        )
+        .await;
+
+        let response = direct_rx.recv().await.unwrap();
+        match response {
+            WsServerMessage::Error { message } => {
+                assert!(message.contains("not started"));
+            }
+            _ => panic!("Expected Error variant"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_approval_approve() {
+        let (agent_tx, mut agent_rx) = mpsc::channel(16);
+        let state = make_test_state(Some(agent_tx)).await;
+        let (direct_tx, _direct_rx) = mpsc::channel(16);
+
+        let request_id = Uuid::new_v4();
+        handle_client_message(
+            WsClientMessage::Approval {
+                request_id: request_id.to_string(),
+                action: "approve".to_string(),
+            },
+            &state,
+            "user1",
+            &direct_tx,
+        )
+        .await;
+
+        let incoming = agent_rx.recv().await.unwrap();
+        // The content should be a serialized ExecApproval
+        assert!(incoming.content.contains("ExecApproval"));
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_approval_invalid_action() {
+        let state = make_test_state(None).await;
+        let (direct_tx, mut direct_rx) = mpsc::channel(16);
+
+        handle_client_message(
+            WsClientMessage::Approval {
+                request_id: Uuid::new_v4().to_string(),
+                action: "maybe".to_string(),
+            },
+            &state,
+            "user1",
+            &direct_tx,
+        )
+        .await;
+
+        let response = direct_rx.recv().await.unwrap();
+        match response {
+            WsServerMessage::Error { message } => {
+                assert!(message.contains("Unknown approval action"));
+            }
+            _ => panic!("Expected Error variant"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_handle_client_approval_invalid_uuid() {
+        let state = make_test_state(None).await;
+        let (direct_tx, mut direct_rx) = mpsc::channel(16);
+
+        handle_client_message(
+            WsClientMessage::Approval {
+                request_id: "not-a-uuid".to_string(),
+                action: "approve".to_string(),
+            },
+            &state,
+            "user1",
+            &direct_tx,
+        )
+        .await;
+
+        let response = direct_rx.recv().await.unwrap();
+        match response {
+            WsServerMessage::Error { message } => {
+                assert!(message.contains("Invalid request_id"));
+            }
+            _ => panic!("Expected Error variant"),
+        }
+    }
+
+    /// Helper to create a GatewayState for testing.
+    async fn make_test_state(msg_tx: Option<mpsc::Sender<IncomingMessage>>) -> GatewayState {
+        use crate::channels::web::sse::SseManager;
+
+        GatewayState {
+            msg_tx: tokio::sync::RwLock::new(msg_tx),
+            sse: SseManager::new(),
+            workspace: None,
+            context_manager: None,
+            session_manager: None,
+            log_broadcaster: None,
+            extension_manager: None,
+            tool_registry: None,
+            user_id: "test".to_string(),
+            shutdown_tx: tokio::sync::RwLock::new(None),
+            ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
+        }
+    }
+}