fix(websocket): retain WsPumps handle on the server path (WS-08)

tests/ws_upgrade_session.rs

@@ -609,3 +609,82 @@ async fn inbound_message_over_cap_fails_the_connection() {
         "expected a close or stream end, got {close:?}"
     );
 }
+
+/// WS-08 acceptance: the upgrade retains each session's pump handle in
+/// the shared `WsSessions` registry; `abort()` evicts a stuck session
+/// (the peer's socket closes) and the entry is removed when the
+/// session task ends, so the registry only tracks live sessions.
+#[tokio::test]
+async fn ws_sessions_registry_tracks_and_aborts_live_sessions() {
+    use alkhttp::websocket::WsSessions;
+
+    let sessions = WsSessions::new();
+    let registry = echo_registry();
+    let provider = provider_with(vec![("tok-1", identity("alice", &[]))]);
+
+    let app = axum::Router::new()
+        .route(
+            "/alk/channels",
+            axum::routing::get(alkhttp::websocket::ws_upgrade_handler),
+        )
+        .layer(axum::middleware::from_fn_with_state(
+            Arc::clone(&provider),
+            alkhttp::websocket::ws_bearer_auth,
+        ))
+        .with_state(registry)
+        .layer(axum::Extension(sessions.clone()));
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = format!("ws://{}", listener.local_addr().unwrap());
+    tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    let mut ws = WsClient::connect_authorized(&format!("{addr}/alk/channels"), "tok-1")
+        .await
+        .unwrap();
+    let env = call_and_await(&mut ws, "req-1", "echo/run", serde_json::json!({})).await;
+    assert_eq!(env.r#type, EVENT_RESPONDED, "session established");
+
+    // The upgrade registered the session.
+    let saw_session = tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        loop {
+            if !sessions.is_empty() {
+                return true;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
+        }
+    })
+    .await
+    .expect("session registered in the shared registry");
+    assert!(saw_session);
+    assert_eq!(sessions.len(), 1);
+
+    // Forced teardown: abort() closes the peer's socket.
+    sessions.abort();
+    let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10);
+    let close = loop {
+        assert!(
+            tokio::time::Instant::now() < deadline,
+            "abort never reached the peer"
+        );
+        match ws.next_close(std::time::Duration::from_millis(250)).await {
+            Some(x) => break Some(x),
+            None => continue,
+        }
+    };
+    assert!(close.is_some(), "peer observes the teardown: {close:?}");
+
+    // The self-removing guard drops the entry once the session task ends.
+    let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async {
+        loop {
+            if sessions.is_empty() {
+                return true;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
+        }
+    })
+    .await
+    .expect("session entry removed after the session task ends");
+    assert!(drained);
+    ws.close().await;
+}