fix terminal exit delivery and stale control-plane views

wakey-agent/src/terminal.rs

@@ -522,7 +522,7 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
     let (stream, _) = tokio_tungstenite::connect_async(ws_url.as_str())
         .await
         .context("failed to connect terminal relay websocket")?;
-    let (mut sink, mut source) = stream.split();
+    let (mut sink, source) = stream.split();
     send_json(
         &mut sink,
         &TerminalAgentHandshake::Auth {
@@ -541,40 +541,61 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
         .await
         .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
 
-    let mut output = relay.output_rx;
+    relay_loop(source, sink, relay.input, relay.generation, relay.output_rx).await
+}
+
+async fn relay_loop<S, St>(
+    mut source: St,
+    mut sink: S,
+    input: mpsc::Sender<RelayInput>,
+    generation: u64,
+    mut output: mpsc::Receiver<Message>,
+) -> Result<()>
+where
+    S: futures_util::Sink<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin,
+    St: futures_util::Stream<
+            Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>,
+        > + Unpin,
+{
+    let mut worker_stopped = false;
+    let mut worker_stop_deadline: Option<tokio::time::Instant> = None;
+    // Reserve at most one control so a full worker queue cannot block output.
+    let mut pending_input = None;
+    let mut close_after_input = false;
     loop {
         tokio::select! {
             biased;
-            incoming = source.next() => {
+            _ = async {
+                match worker_stop_deadline {
+                    Some(deadline) => tokio::time::sleep_until(deadline).await,
+                    None => std::future::pending().await,
+                }
+            } => break,
+            incoming = source.next(), if !worker_stopped && pending_input.is_none() => {
                 let Some(message) = incoming else { break; };
                 match message.context("terminal relay websocket receive failed")? {
                     Message::Binary(bytes) => {
-                        relay.input.send(RelayInput::Binary {
-                            generation: relay.generation,
+                        pending_input = Some(RelayInput::Binary {
+                            generation,
                             bytes: bytes.to_vec(),
-                        }).await.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
+                        });
                     }
                     Message::Text(text) => match serde_json::from_str::<TerminalControl>(&text)
                         .context("invalid terminal control frame")?
                     {
                         TerminalControl::Resize { rows, cols } => {
-                            relay.input.send(RelayInput::Resize {
-                                generation: relay.generation,
+                            pending_input = Some(RelayInput::Resize {
+                                generation,
                                 rows,
                                 cols,
-                            }).await
-                                .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
+                            });
                         }
                         TerminalControl::Snapshot => {
-                            relay.input.send(RelayInput::Snapshot {
-                                generation: relay.generation,
-                            }).await.map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
+                            pending_input = Some(RelayInput::Snapshot { generation });
                         }
                         TerminalControl::Close => {
-                            let _ = relay.input.send(RelayInput::Close {
-                                generation: relay.generation,
-                            }).await;
-                            break;
+                            pending_input = Some(RelayInput::Close { generation });
+                            close_after_input = true;
                         }
                         _ => anyhow::bail!("terminal control frame has invalid direction"),
                     },
@@ -586,9 +607,32 @@ async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
                     Message::Frame(_) => {}
                 }
             }
+            permit = input.reserve(), if !worker_stopped && pending_input.is_some() => {
+                match permit {
+                    Ok(permit) => {
+                        permit.send(pending_input.take().expect("pending terminal input"));
+                        if close_after_input {
+                            break;
+                        }
+                    }
+                    Err(_) => {
+                        worker_stopped = true;
+                        pending_input = None;
+                        worker_stop_deadline =
+                            Some(tokio::time::Instant::now() + PTY_EXIT_DRAIN_TIMEOUT);
+                    }
+                }
+            }
             outgoing = output.recv() => {
                 let Some(message) = outgoing else { break; };
-                sink.send(message).await.context("failed to send terminal relay output")?;
+                if let Some(deadline) = worker_stop_deadline {
+                    tokio::time::timeout_at(deadline, sink.send(message))
+                        .await
+                        .context("terminal relay output drain timed out")?
+                        .context("failed to send terminal relay output")?;
+                } else {
+                    sink.send(message).await.context("failed to send terminal relay output")?;
+                }
             }
         }
     }
@@ -871,6 +915,58 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn relay_drains_final_output_after_worker_stops_during_concurrent_controls() {
+        let (input, input_rx) = mpsc::channel(RELAY_INPUT_QUEUE);
+        drop(input_rx);
+        let (output_tx, output_rx) = mpsc::channel(RELAY_OUTPUT_QUEUE);
+        output_tx
+            .send(Message::Binary(vec![1, 2, 3].into()))
+            .await
+            .expect("live output queues");
+        let exit_frame = serde_json::to_string(&TerminalControl::Exited { exit_code: Some(0) })
+            .expect("exit frame encodes");
+        output_tx
+            .send(Message::Text(exit_frame.clone().into()))
+            .await
+            .expect("exit output queues");
+        drop(output_tx);
+
+        let source = futures_util::stream::iter([
+            Ok(Message::Binary(vec![4].into())),
+            Ok(Message::Text(
+                serde_json::to_string(&TerminalControl::Resize {
+                    rows: 30,
+                    cols: 100,
+                })
+                .expect("resize encodes")
+                .into(),
+            )),
+            Ok(Message::Text(
+                serde_json::to_string(&TerminalControl::Snapshot)
+                    .expect("snapshot encodes")
+                    .into(),
+            )),
+        ]);
+        let delivered = Arc::new(Mutex::new(Vec::new()));
+        let sink = futures_util::sink::unfold(delivered.clone(), |delivered, frame| async move {
+            delivered.lock().expect("delivered frames").push(frame);
+            Ok::<_, tokio_tungstenite::tungstenite::Error>(delivered)
+        });
+
+        relay_loop(source, Box::pin(sink), input, 7, output_rx)
+            .await
+            .expect("relay drains after worker completion");
+
+        let delivered = delivered.lock().expect("delivered frames");
+        assert_eq!(delivered.len(), 2);
+        assert_eq!(delivered[0].clone().into_data(), vec![1, 2, 3]);
+        let Message::Text(exit) = &delivered[1] else {
+            panic!("exit control should follow final bytes");
+        };
+        assert_eq!(exit.as_str(), exit_frame);
+    }
+
     #[test]
     fn open_rejects_invalid_initial_size_before_spawning() {
         let mut config = crate::config::DEFAULT_CONFIG.clone();