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

src/gateway/routes.rs

@@ -857,6 +857,7 @@ mod tests {
             identity_provider: Arc::clone(&provider),
             decoy: crate::server::DecoyConfig::NotFound,
             openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&registry),
+            ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
         };
         let auth_state = Arc::clone(&provider);
         gateway_router()

src/server/adapter.rs

@@ -87,6 +87,7 @@ pub struct HttpAdapter {
     alpn: &'static [u8],
     router: Router,
     openapi_doc: CachedOpenAPIDoc,
+    ws_sessions: Arc<crate::websocket::WsSessions>,
 }
 
 impl HttpAdapter {
@@ -111,11 +112,13 @@ impl HttpAdapter {
     ) -> Self {
         let decoy = DecoyConfig::default();
         let openapi_doc = CachedOpenAPIDoc::new(&registry);
+        let ws_sessions = Arc::new(crate::websocket::WsSessions::new());
         let state = RouterState {
             registry: Arc::clone(&registry),
             identity_provider: Arc::clone(&identity_provider),
             decoy: decoy.clone(),
             openapi_doc: openapi_doc.clone(),
+            ws_sessions: Arc::clone(&ws_sessions),
         };
         let router = build_router(state, None);
         Self {
@@ -126,6 +129,7 @@ impl HttpAdapter {
             alpn,
             router,
             openapi_doc,
+            ws_sessions,
         }
     }
 
@@ -136,6 +140,7 @@ impl HttpAdapter {
             identity_provider: Arc::clone(&self.identity_provider),
             decoy,
             openapi_doc: self.openapi_doc.clone(),
+            ws_sessions: Arc::clone(&self.ws_sessions),
         };
         // `extra_routes` is borrowed, not consumed (SRV-05): a builder
         // call after `with_extra_routes` must keep the custom routes in
@@ -151,12 +156,21 @@ impl HttpAdapter {
             identity_provider: Arc::clone(&self.identity_provider),
             decoy: self.decoy.clone(),
             openapi_doc: self.openapi_doc.clone(),
+            ws_sessions: Arc::clone(&self.ws_sessions),
         };
         self.router = build_router(state, Some(routes.clone()));
         self.extra_routes = Some(routes);
         self
     }
 
+    /// The shared WS session registry (WS-08): live sessions'
+    /// [`WsPumps`](crate::websocket::WsPumps) handles, evictable via
+    /// `WsSessions::abort`. The upgrade handler registers against this
+    /// instance.
+    pub fn ws_sessions(&self) -> Arc<crate::websocket::WsSessions> {
+        Arc::clone(&self.ws_sessions)
+    }
+
     pub fn decoy(&self) -> &DecoyConfig {
         &self.decoy
     }
@@ -941,6 +955,7 @@ mod tests {
             identity_provider: idp,
             decoy: DecoyConfig::default(),
             openapi_doc: CachedOpenAPIDoc::new(&OperationRegistry::new()),
+            ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
         }
     }
 

src/server/state.rs

@@ -25,14 +25,25 @@ pub enum DecoyConfig {
 
 /// State embedded in the axum `Router`: the registry and identity
 /// provider every request handler reaches through the router state, plus
-/// the decoy config for the fallback and the pre-serialized
-/// `/openapi.json` projection cache (SRV-09).
+/// the decoy config for the fallback, the pre-serialized
+/// `/openapi.json` projection cache (SRV-09), and the shared WS session
+/// registry the WS upgrade retains pump handles in (WS-08).
 #[derive(Clone)]
 pub(crate) struct RouterState {
     pub(crate) registry: Arc<OperationRegistry>,
     pub(crate) identity_provider: Arc<dyn IdentityProvider>,
     pub(crate) decoy: DecoyConfig,
     pub(crate) openapi_doc: crate::server::adapter::CachedOpenAPIDoc,
+    pub(crate) ws_sessions: Arc<crate::websocket::WsSessions>,
+}
+
+impl axum::extract::FromRef<RouterState> for crate::websocket::SessionState {
+    fn from_ref(state: &RouterState) -> Self {
+        crate::websocket::SessionState::new(
+            Arc::clone(&state.registry),
+            Arc::clone(&state.ws_sessions),
+        )
+    }
 }
 
 impl axum::extract::FromRef<RouterState> for DecoyConfig {
@@ -77,6 +88,7 @@ mod tests {
                 to: "https://example.com".to_string(),
             },
             openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&OperationRegistry::new()),
+            ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
         };
         let extracted: DecoyConfig = axum::extract::FromRef::from_ref(&state);
         assert!(matches!(extracted, DecoyConfig::Redirect { .. }));

src/websocket/mod.rs

@@ -22,7 +22,10 @@ pub use byte_adapter::split_tungstenite_to_bytes;
 #[cfg(any(test, feature = "wss"))]
 #[allow(unused_imports)]
 pub(crate) use upgrade::adapter_install_channel_zero;
-pub use upgrade::{run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy};
+pub use upgrade::{
+    run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy, SessionState,
+    WsSessions,
+};
 
 #[cfg(any(test, feature = "test-support"))]
 pub use upgrade::test_support::{frame_channel0_chunk, ChunkAssembler, FrameAssembler, WsClient};

src/websocket/upgrade.rs

@@ -10,6 +10,8 @@
 //! `CallConnection` (the identity rides the channels-layer
 //! `Connection`) and runs the shared `Dispatcher::run_loop_single_stream`.
 
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use alkcall::channels::adapter::ChannelsAdapter;
@@ -20,19 +22,137 @@ use alkcall::registry::registration::OperationRegistry;
 use axum::extract::ws::WebSocketUpgrade;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
+use parking_lot::Mutex;
 
-use super::byte_adapter::{split_ws_to_bytes, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP};
+use super::byte_adapter::{
+    split_ws_to_bytes, WsPumps, INBOUND_WS_FRAME_CAP, INBOUND_WS_MESSAGE_CAP,
+};
+
+/// Registry of live WS session pump handles (WS-08). The upgrade
+/// handler retains each session's [`WsPumps`] handle here for the
+/// session's lifetime so a stuck session is evictable in-crate:
+/// [`WsSessions::abort`] forces the pumps' tasks to end (the socket
+/// halves close and the channels session unwinds). The handle for a
+/// session is removed when the session task finishes (self-removing
+/// guard), so the registry only holds live sessions. Assembly layers
+/// share one instance via [`crate::server::state::RouterState`] (the
+/// upgrade handler registers against it) or per-route request
+/// extensions (an extension clone takes precedence).
+///
+/// Un-registered (default) — the upgrade runs fine and simply keeps no
+/// eviction lever, matching the pre-WS-08 behavior.
+#[derive(Clone, Default)]
+pub struct WsSessions {
+    counter: Arc<AtomicU64>,
+    sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
+}
+
+/// The upgrade handler's state slice: what it needs beyond the request
+/// itself. Axum lifts it via `FromRef` from either full router state —
+/// the adapter's [`crate::server::state::RouterState`] (carrying the
+/// shared [`WsSessions`] instance) or a bare `Arc<OperationRegistry>`
+/// (custom upgrade routes / integration tests get a handler-private
+/// registry; eviction still works in-crate, just not shared).
+#[derive(Clone)]
+pub struct SessionState {
+    registry: Arc<OperationRegistry>,
+    sessions: Arc<WsSessions>,
+}
+
+impl SessionState {
+    /// From the plain registry state (custom upgrade routes /
+    /// integration tests): sessions default to a fresh [`WsSessions`]
+    /// private to the handler (eviction still functional in-crate but
+    /// not shared with the assembly layer).
+    pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
+        Self {
+            registry: Arc::clone(registry),
+            sessions: Arc::new(WsSessions::new()),
+        }
+    }
+
+    /// From the adapter's router state: the shared [`WsSessions`]
+    /// instance the assembly layer can hold for eviction.
+    pub(crate) fn new(registry: Arc<OperationRegistry>, sessions: Arc<WsSessions>) -> Self {
+        Self { registry, sessions }
+    }
+
+    pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
+        &self.registry
+    }
+
+    pub(crate) fn sessions(&self) -> &Arc<WsSessions> {
+        &self.sessions
+    }
+}
+
+impl axum::extract::FromRef<SessionState> for Arc<OperationRegistry> {
+    fn from_ref(state: &SessionState) -> Self {
+        Arc::clone(&state.registry)
+    }
+}
+
+/// `FromRef` chain: a bare `Arc<OperationRegistry>` router state lifts
+/// into the handler's [`SessionState`]; a full [`RouterState`] carries
+/// the shared [`WsSessions`] instance and lifts through its own impl.
+impl axum::extract::FromRef<Arc<OperationRegistry>> for SessionState {
+    fn from_ref(registry: &Arc<OperationRegistry>) -> Self {
+        SessionState::from_registry(registry)
+    }
+}
+
+impl WsSessions {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Abort every live session's pump tasks (forced teardown).
+    pub fn abort(&self) {
+        for (_, pumps) in self.sessions.lock().drain() {
+            pumps.abort();
+        }
+    }
+
+    /// Number of live sessions currently tracked.
+    pub fn len(&self) -> usize {
+        self.sessions.lock().len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.sessions.lock().is_empty()
+    }
+
+    fn insert(&self, pumps: Arc<WsPumps>) -> u64 {
+        let id = self.counter.fetch_add(1, Ordering::Relaxed);
+        self.sessions.lock().insert(id, pumps);
+        id
+    }
+
+    fn remove(&self, id: u64) {
+        self.sessions.lock().remove(&id);
+    }
+}
 
 /// The channels session for an upgraded socket: adapt → `Connection`
 /// (identity attached) → `ChannelsAdapter::handle`. `policy` gates
-/// data-channel opens (ADR-041).
+/// data-channel opens (ADR-041). When `sessions` is `Some`, the
+/// session's pump handle is registered for the session's lifetime —
+/// the WS-08 eviction lever ([`WsSessions::abort`]).
 pub async fn run_channels_session(
     socket: axum::extract::ws::WebSocket,
     registry: Arc<OperationRegistry>,
     identity: Identity,
     policy: Arc<dyn ChannelLifecyclePolicy>,
+    sessions: Option<WsSessions>,
 ) {
-    let (byte_stream, _pumps) = split_ws_to_bytes(socket);
+    let (byte_stream, pumps) = split_ws_to_bytes(socket);
+    let pumps = Arc::new(pumps);
+
+    let _guard = sessions.map(|sessions| {
+        let id = sessions.insert(Arc::clone(&pumps));
+        SessionGuard { sessions, id }
+    });
+
     let conn = Connection::from_bidi(byte_stream, b"alk/channels".to_vec(), None);
     let _ = conn.set_identity(identity.clone());
 
@@ -48,6 +168,17 @@ pub async fn run_channels_session(
     }
 }
 
+struct SessionGuard {
+    sessions: WsSessions,
+    id: u64,
+}
+
+impl Drop for SessionGuard {
+    fn drop(&mut self) {
+        self.sessions.remove(self.id);
+    }
+}
+
 /// The `install_channel_zero` hook: split channel 0's `BiStream` into
 /// the shared writer + reader, construct the single-stream
 /// `CallConnection`, and run the shared `Dispatcher`'s single-stream
@@ -131,20 +262,31 @@ pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
 ///
 /// The channel lifecycle policy comes from the
 /// [`ChannelsPolicy`] request extension when present, else `NoCap`.
+/// The session pump handles are retained in the [`WsSessions`]
+/// registry — the shared instance from the router state
+/// (`RouterState::ws_sessions`) unless a request extension carries
+/// one, so a stuck session stays evictable (WS-08).
 pub async fn ws_upgrade_handler(
-    axum::extract::State(registry): axum::extract::State<Arc<OperationRegistry>>,
+    sessions: Option<axum::Extension<WsSessions>>,
+    axum::extract::State(state): axum::extract::State<SessionState>,
     axum::Extension(identity): axum::Extension<Identity>,
     policy: Option<axum::Extension<ChannelsPolicy>>,
     ws_upgrade: WebSocketUpgrade,
 ) -> Response {
+    let sessions = Some(
+        sessions
+            .map(|axum::Extension(s)| s)
+            .unwrap_or_else(|| WsSessions::clone(state.sessions())),
+    );
     let policy = policy
         .map(|axum::Extension(p)| p.0)
         .unwrap_or_else(|| Arc::new(NoCap));
+    let registry = Arc::clone(state.registry());
     ws_upgrade
         .max_frame_size(INBOUND_WS_FRAME_CAP)
         .max_message_size(INBOUND_WS_MESSAGE_CAP)
         .on_upgrade(move |socket| async move {
-            run_channels_session(socket, registry, identity, policy).await
+            run_channels_session(socket, registry, identity, policy, sessions).await
         })
 }
 

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;
+}