fix(websocket): configurable WS session cap (WS-09)

src/gateway/routes.rs

@@ -858,6 +858,9 @@ mod tests {
             decoy: crate::server::DecoyConfig::NotFound,
             openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&registry),
             ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
+            ws_session_slots: Arc::new(tokio::sync::Semaphore::new(
+                crate::websocket::DEFAULT_WS_MAX_SESSIONS,
+            )),
         };
         let auth_state = Arc::clone(&provider);
         gateway_router()

src/server/adapter.rs

@@ -88,6 +88,8 @@ pub struct HttpAdapter {
     router: Router,
     openapi_doc: CachedOpenAPIDoc,
     ws_sessions: Arc<crate::websocket::WsSessions>,
+    ws_max_sessions: usize,
+    ws_session_slots: Arc<tokio::sync::Semaphore>,
 }
 
 impl HttpAdapter {
@@ -113,12 +115,15 @@ impl HttpAdapter {
         let decoy = DecoyConfig::default();
         let openapi_doc = CachedOpenAPIDoc::new(&registry);
         let ws_sessions = Arc::new(crate::websocket::WsSessions::new());
+        let ws_max_sessions = crate::websocket::DEFAULT_WS_MAX_SESSIONS;
+        let ws_session_slots = Arc::new(tokio::sync::Semaphore::new(ws_max_sessions));
         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),
+            ws_session_slots: Arc::clone(&ws_session_slots),
         };
         let router = build_router(state, None);
         Self {
@@ -130,6 +135,8 @@ impl HttpAdapter {
             router,
             openapi_doc,
             ws_sessions,
+            ws_max_sessions,
+            ws_session_slots,
         }
     }
 
@@ -141,6 +148,7 @@ impl HttpAdapter {
             decoy,
             openapi_doc: self.openapi_doc.clone(),
             ws_sessions: Arc::clone(&self.ws_sessions),
+            ws_session_slots: Arc::clone(&self.ws_session_slots),
         };
         // `extra_routes` is borrowed, not consumed (SRV-05): a builder
         // call after `with_extra_routes` must keep the custom routes in
@@ -157,12 +165,34 @@ impl HttpAdapter {
             decoy: self.decoy.clone(),
             openapi_doc: self.openapi_doc.clone(),
             ws_sessions: Arc::clone(&self.ws_sessions),
+            ws_session_slots: Arc::clone(&self.ws_session_slots),
         };
         self.router = build_router(state, Some(routes.clone()));
         self.extra_routes = Some(routes);
         self
     }
 
+    /// The concurrent WS session cap (WS-09): the upgrade handler
+    /// acquires one semaphore permit per upgrade, post-auth and
+    /// pre-upgrade; a caller over the configured cap is rejected with
+    /// **503 Service Unavailable**, and the permit is held for the
+    /// session's lifetime (an ended session frees its slot).
+    ///
+    /// Default: [`crate::websocket::DEFAULT_WS_MAX_SESSIONS`] (64).
+    pub fn with_ws_max_sessions(mut self, max_sessions: usize) -> Self {
+        self.ws_max_sessions = max_sessions;
+        let state = RouterState {
+            registry: Arc::clone(&self.registry),
+            identity_provider: Arc::clone(&self.identity_provider),
+            decoy: self.decoy.clone(),
+            openapi_doc: self.openapi_doc.clone(),
+            ws_sessions: Arc::clone(&self.ws_sessions),
+            ws_session_slots: Self::rebuild_session_slots(max_sessions),
+        };
+        self.router = build_router(state, self.extra_routes.clone());
+        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
@@ -171,6 +201,12 @@ impl HttpAdapter {
         Arc::clone(&self.ws_sessions)
     }
 
+    /// A fresh semaphore for the new cap; the retained handles of
+    /// already-open sessions are unaffected (they hold their permits).
+    fn rebuild_session_slots(max_sessions: usize) -> Arc<tokio::sync::Semaphore> {
+        Arc::new(tokio::sync::Semaphore::new(max_sessions))
+    }
+
     pub fn decoy(&self) -> &DecoyConfig {
         &self.decoy
     }
@@ -956,6 +992,9 @@ mod tests {
             decoy: DecoyConfig::default(),
             openapi_doc: CachedOpenAPIDoc::new(&OperationRegistry::new()),
             ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
+            ws_session_slots: Arc::new(tokio::sync::Semaphore::new(
+                crate::websocket::DEFAULT_WS_MAX_SESSIONS,
+            )),
         }
     }
 

src/server/state.rs

@@ -26,8 +26,9 @@ 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, the pre-serialized
-/// `/openapi.json` projection cache (SRV-09), and the shared WS session
-/// registry the WS upgrade retains pump handles in (WS-08).
+/// `/openapi.json` projection cache (SRV-09), and the WS session
+/// lifecycle config the upgrade handler enforces: the shared pump-handle
+/// registry (WS-08) and the concurrent-session cap (WS-09).
 #[derive(Clone)]
 pub(crate) struct RouterState {
     pub(crate) registry: Arc<OperationRegistry>,
@@ -35,6 +36,10 @@ pub(crate) struct RouterState {
     pub(crate) decoy: DecoyConfig,
     pub(crate) openapi_doc: crate::server::adapter::CachedOpenAPIDoc,
     pub(crate) ws_sessions: Arc<crate::websocket::WsSessions>,
+    /// Session cap (WS-09): one `HttpAdapter`-wide semaphore, built
+    /// once at construction — the upgrade handler acquires one permit
+    /// per upgrade and holds it for the session's lifetime.
+    pub(crate) ws_session_slots: Arc<tokio::sync::Semaphore>,
 }
 
 impl axum::extract::FromRef<RouterState> for crate::websocket::SessionState {
@@ -42,6 +47,7 @@ impl axum::extract::FromRef<RouterState> for crate::websocket::SessionState {
         crate::websocket::SessionState::new(
             Arc::clone(&state.registry),
             Arc::clone(&state.ws_sessions),
+            Arc::clone(&state.ws_session_slots),
         )
     }
 }
@@ -89,6 +95,9 @@ mod tests {
             },
             openapi_doc: crate::server::adapter::CachedOpenAPIDoc::new(&OperationRegistry::new()),
             ws_sessions: Arc::new(crate::websocket::WsSessions::new()),
+            ws_session_slots: Arc::new(tokio::sync::Semaphore::new(
+                crate::websocket::DEFAULT_WS_MAX_SESSIONS,
+            )),
         };
         let extracted: DecoyConfig = axum::extract::FromRef::from_ref(&state);
         assert!(matches!(extracted, DecoyConfig::Redirect { .. }));

src/websocket/mod.rs

@@ -24,7 +24,7 @@ pub use byte_adapter::split_tungstenite_to_bytes;
 pub(crate) use upgrade::adapter_install_channel_zero;
 pub use upgrade::{
     run_channels_session, ws_bearer_auth, ws_upgrade_handler, ChannelsPolicy, SessionState,
-    WsSessions,
+    WsSessions, DEFAULT_WS_MAX_SESSIONS,
 };
 
 #[cfg(any(test, feature = "test-support"))]

src/websocket/upgrade.rs

@@ -47,34 +47,54 @@ pub struct WsSessions {
     sessions: Arc<Mutex<HashMap<u64, Arc<WsPumps>>>>,
 }
 
+/// Default bound on concurrent WS sessions (WS-09): the semaphore
+/// initial capacity in [`SessionState`]; a deployment overrides it with
+/// `HttpAdapter::with_ws_max_sessions`.
+pub const DEFAULT_WS_MAX_SESSIONS: usize = 64;
+
 /// 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).
+/// shared [`WsSessions`] instance and the configured session cap) or a
+/// bare `Arc<OperationRegistry>` (custom upgrade routes / integration
+/// tests get a handler-private registry and the default cap; eviction
+/// still works in-crate, just not shared).
 #[derive(Clone)]
 pub struct SessionState {
     registry: Arc<OperationRegistry>,
     sessions: Arc<WsSessions>,
+    /// Session cap (WS-09): acquired post-auth, pre-upgrade; a caller
+    /// over the cap is rejected with 503.
+    session_slots: Arc<tokio::sync::Semaphore>,
 }
 
 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).
+    /// not shared with the assembly layer) and the default session cap.
     pub(crate) fn from_registry(registry: &Arc<OperationRegistry>) -> Self {
         Self {
             registry: Arc::clone(registry),
             sessions: Arc::new(WsSessions::new()),
+            session_slots: Arc::new(tokio::sync::Semaphore::new(DEFAULT_WS_MAX_SESSIONS)),
         }
     }
 
     /// 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 }
+    /// instance the assembly layer can hold for eviction, and the
+    /// pre-built session-cap semaphore (one per `HttpAdapter`, shared
+    /// across requests).
+    pub(crate) fn new(
+        registry: Arc<OperationRegistry>,
+        sessions: Arc<WsSessions>,
+        session_slots: Arc<tokio::sync::Semaphore>,
+    ) -> Self {
+        Self {
+            registry,
+            sessions,
+            session_slots,
+        }
     }
 
     pub(crate) fn registry(&self) -> &Arc<OperationRegistry> {
@@ -266,6 +286,13 @@ pub struct ChannelsPolicy(pub Arc<dyn ChannelLifecyclePolicy>);
 /// 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).
+///
+/// Session cap (WS-09): one semaphore permit is acquired per upgrade,
+/// post-auth and pre-upgrade; when the configured cap
+/// (`HttpAdapter::with_ws_max_sessions`, default
+/// [`DEFAULT_WS_MAX_SESSIONS`]) is exhausted the upgrade is rejected
+/// with **503 Service Unavailable** — holding the permit for the
+/// session's lifetime, so ended sessions free their slot.
 pub async fn ws_upgrade_handler(
     sessions: Option<axum::Extension<WsSessions>>,
     axum::extract::State(state): axum::extract::State<SessionState>,
@@ -273,6 +300,13 @@ pub async fn ws_upgrade_handler(
     policy: Option<axum::Extension<ChannelsPolicy>>,
     ws_upgrade: WebSocketUpgrade,
 ) -> Response {
+    let Ok(permit) = state.session_slots.clone().try_acquire_owned() else {
+        return (
+            StatusCode::SERVICE_UNAVAILABLE,
+            "503 Service Unavailable: WS session cap reached",
+        )
+            .into_response();
+    };
     let sessions = Some(
         sessions
             .map(|axum::Extension(s)| s)
@@ -286,6 +320,7 @@ pub async fn ws_upgrade_handler(
         .max_frame_size(INBOUND_WS_FRAME_CAP)
         .max_message_size(INBOUND_WS_MESSAGE_CAP)
         .on_upgrade(move |socket| async move {
+            let _permit = permit;
             run_channels_session(socket, registry, identity, policy, sessions).await
         })
 }

tests/ws_upgrade_session.rs

@@ -688,3 +688,73 @@ async fn ws_sessions_registry_tracks_and_aborts_live_sessions() {
     assert!(drained);
     ws.close().await;
 }
+
+/// WS-09 acceptance: the session cap is enforced on the upgrade — a
+/// caller over the configured cap is rejected with 503, and an ended
+/// session frees its slot for the next caller.
+#[tokio::test]
+async fn session_cap_rejects_over_limit_with_503_and_frees_slots_on_end() {
+    use alkhttp::server::HttpAdapter;
+
+    let mut registry_val = OperationRegistry::new();
+    registry_val
+        .register(HandlerRegistration::new(
+            OperationSpec::new(
+                "echo/run",
+                OperationType::Query,
+                Visibility::External,
+                serde_json::json!({}),
+                serde_json::json!({}),
+                vec![],
+                AccessControl::default(),
+                None,
+            ),
+            HandlerKind::Once(echo_handler()),
+            OperationProvenance::Local,
+            None,
+            None,
+            alkcall::core::types::Capabilities::new(),
+        ))
+        .unwrap();
+    let registry = Arc::new(registry_val);
+    struct StaticTok;
+    impl IdentityProvider for StaticTok {
+        fn resolve_from_fingerprint(&self, _: &str) -> Option<Identity> {
+            None
+        }
+        fn resolve_from_token(&self, token: &alkcall::core::auth::AuthToken) -> Option<Identity> {
+            let s = String::from_utf8_lossy(&token.raw).to_string();
+            (s == "tok-1").then(|| identity("alice", &[]))
+        }
+    }
+
+    // The adapter's built-in surface: WS route + bearer middleware, the
+    // whole deployment shape the cap knob configures.
+    let adapter =
+        HttpAdapter::new(Arc::new(StaticTok), Arc::clone(&registry)).with_ws_max_sessions(1);
+    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = format!("ws://{}", listener.local_addr().unwrap());
+    let app: axum::Router = adapter.router().clone();
+    tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    let url = format!("{addr}/alk/channels");
+
+    // First session: admitted.
+    let mut first = WsClient::connect_authorized(&url, "tok-1").await.unwrap();
+    let env = call_and_await(&mut first, "req-1", "echo/run", serde_json::json!({})).await;
+    assert_eq!(env.r#type, EVENT_RESPONDED, "first session admitted");
+
+    // Second session over the cap: rejected with 503 at the upgrade.
+    let status = WsClient::connect_status(&url, Some("tok-1")).await;
+    assert_eq!(status, Some(503), "second session rejected over the cap");
+
+    // Ending the first session frees the slot: the next caller is admitted.
+    first.close().await;
+    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+    let mut third = WsClient::connect_authorized(&url, "tok-1").await.unwrap();
+    let env = call_and_await(&mut third, "req-2", "echo/run", serde_json::json!({})).await;
+    assert_eq!(env.r#type, EVENT_RESPONDED, "slot freed after session end");
+    third.close().await;
+}