screen relay: close the websocket when the vnc side hangs up and stop the surviving task

Cargo.lock

@@ -1697,6 +1697,7 @@ dependencies = [
  "tempfile",
  "tokio",
  "tokio-stream",
+ "tokio-tungstenite",
  "tower",
  "tower-http 0.7.0",
  "tracing",

hive-agent/Cargo.toml

@@ -70,6 +70,9 @@ tracing.workspace = true
 # refresh test waits out its one-minute interval on paused time.
 [dev-dependencies]
 tokio = { workspace = true, features = ["test-util"] }
+# A WebSocket client for the screen relay teardown test; the same version
+# axum's `ws` feature already pulls in.
+tokio-tungstenite = "0.29"
 tower = { version = "0.5", features = ["util"] }
 # A layer that captures `swarm_queue`'s exhausted-connect log event in a test.
 tracing-subscriber.workspace = true

hive-agent/src/web_ui/screen.rs

@@ -55,7 +55,7 @@ async fn relay_ws_vnc(socket: axum::extract::ws::WebSocket, vnc_port: u16) {
     let (mut ws_tx, mut ws_rx) = socket.split();
 
     // WS → TCP
-    let ws_to_tcp = tokio::spawn(async move {
+    let mut ws_to_tcp = tokio::spawn(async move {
         while let Some(Ok(msg)) = futures_util::StreamExt::next(&mut ws_rx).await {
             match msg {
                 Message::Binary(data) if tcp_tx.write_all(&data).await.is_err() => {
@@ -68,7 +68,7 @@ async fn relay_ws_vnc(socket: axum::extract::ws::WebSocket, vnc_port: u16) {
     });
 
     // TCP → WS
-    let tcp_to_ws = tokio::spawn(async move {
+    let mut tcp_to_ws = tokio::spawn(async move {
         let mut buf = vec![0u8; 8192];
         loop {
             match tcp_rx.read(&mut buf).await {
@@ -84,11 +84,64 @@ async fn relay_ws_vnc(socket: axum::extract::ws::WebSocket, vnc_port: u16) {
                 }
             }
         }
+        // The VNC side is gone: tell the browser instead of leaving it
+        // holding a connection that will never deliver anything.
+        let _ = ws_tx.send(Message::Close(None)).await;
     });
 
-    // Wait for either direction to close, then let both tasks drop.
+    // When either direction ends, stop the other one too. Dropping a
+    // `JoinHandle` only detaches its task, so it is aborted explicitly.
     tokio::select! {
-        _ = ws_to_tcp => {}
-        _ = tcp_to_ws => {}
+        _ = &mut ws_to_tcp => tcp_to_ws.abort(),
+        _ = &mut tcp_to_ws => ws_to_tcp.abort(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::relay_ws_vnc;
+    use futures_util::StreamExt as _;
+    use std::time::Duration;
+
+    /// When the VNC server hangs up, the browser's WebSocket must be closed
+    /// rather than left open and silent.
+    #[tokio::test]
+    async fn the_websocket_closes_when_the_vnc_side_hangs_up() {
+        let vnc = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let vnc_port = vnc.local_addr().unwrap().port();
+        // Accept the relay's connection, then drop it at once.
+        tokio::spawn(async move {
+            let _ = vnc.accept().await;
+        });
+
+        let app = axum::Router::new().route(
+            "/ws",
+            axum::routing::get(move |ws: axum::extract::ws::WebSocketUpgrade| async move {
+                ws.on_upgrade(move |socket| relay_ws_vnc(socket, vnc_port))
+            }),
+        );
+        let http = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let http_port = http.local_addr().unwrap().port();
+        tokio::spawn(async move { axum::serve(http, app).await });
+
+        let (mut client, _) =
+            tokio_tungstenite::connect_async(format!("ws://127.0.0.1:{http_port}/ws"))
+                .await
+                .unwrap();
+        let ended = tokio::time::timeout(Duration::from_secs(5), async {
+            loop {
+                match client.next().await {
+                    None | Some(Err(_) | Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => {
+                        break;
+                    }
+                    Some(Ok(_)) => {}
+                }
+            }
+        })
+        .await;
+        assert!(
+            ended.is_ok(),
+            "websocket still open after the VNC side hung up"
+        );
     }
 }