diff --git a/src/websocket/upgrade.rs b/src/websocket/upgrade.rs index 9754ea5..9263e07 100644 --- a/src/websocket/upgrade.rs +++ b/src/websocket/upgrade.rs @@ -448,15 +448,16 @@ fn install_channel_zero( return; } - if let Some(sessions) = &sessions { + let _conn_guard = sessions.as_ref().map(|sessions| { let id = sessions.insert_connection(Arc::clone(&call_connection)); - let _conn_guard = ConnectionGuard { + ConnectionGuard { sessions: sessions.clone(), id, - }; - // The guard lives for this task — its drop removes the - // handle when the dispatcher loop returns. - } + } + }); + // `conn_guard` lives in the dispatcher task's frame: its drop + // removes the handle when the loop below returns — any + // teardown path (peer close, idle eviction, abort, EOF). let dispatcher = alkcall::protocol::dispatch::Dispatcher::new( fork, diff --git a/tests/ws_upgrade_session.rs b/tests/ws_upgrade_session.rs index 4887ab9..ad19b0e 100644 --- a/tests/ws_upgrade_session.rs +++ b/tests/ws_upgrade_session.rs @@ -696,6 +696,76 @@ async fn ws_sessions_registry_tracks_and_aborts_live_sessions() { ws.close().await; } +/// WS-26/WS-28 acceptance: the channel-0 `CallConnection` handle is +/// visible in the shared `WsSessions` registry while the session is +/// live (`live_connections()` non-empty mid-session, after a completed +/// call proves the dispatcher task is up) and drains when the session +/// ends — the guard lives for the channel-0 task, not the `if let` +/// block that inserts it. This is the only gate exercising the WS-26 +/// surface; its absence is why `030c5ef` landed green (review 007 +/// WS-28, reproduced empirically). +#[tokio::test] +async fn live_connections_visible_mid_session_and_drain_after_teardown() { + use alkhttp::websocket::WsSessions; + + let sessions = WsSessions::new(); + let registry = echo_registry(); + + let app = axum::Router::new() + .route( + "/alk/channels", + axum::routing::get(alkhttp::websocket::ws_upgrade_handler), + ) + .layer(axum::middleware::from_fn_with_state( + provider_with(vec![("tok-1", identity("alice", &[]))]), + 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 handle is retained for the channel-0 task's lifetime: visible + // mid-session (the dispatcher loop is fully up — the call returned). + let saw_handle = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if sessions.live_connection_count() == 1 { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("channel-0 connection handle visible mid-session (WS-28)"); + assert!(saw_handle); + assert_eq!(sessions.live_connection_count(), 1); + let handles = sessions.live_connections(); + assert_eq!(handles.len(), 1, "one handle for the one session"); + + // Teardown: the guard's drop removes the handle on any end path. + ws.close().await; + let drained = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if sessions.live_connection_count() == 0 { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + }) + .await + .expect("channel-0 connection handle removed after the session ends"); + assert!(drained); +} + /// 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.