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]>
This commit is contained in:
firat.sertgoz
2026-02-09 01:00:35 +00:00
committed by GitHub
co-authored by Claude
parent 6831a54793
commit 642c320b13
8 changed files with 1137 additions and 6 deletions
+66
View File
@@ -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);
}
}