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

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 { .. }));