support long running session

plus UI reattaching to existing session

wakey-agent/src/terminal.rs

@@ -1,10 +1,10 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
 use std::path::Path;
 use std::sync::{Arc, Mutex, Weak};
 use std::time::Duration;
 
 use crate::config::AgentConfig;
-use crate::protocol::{TerminalAgentHandshake, TerminalControl, TerminalId};
+use crate::protocol::{AgentTerminalSession, TerminalAgentHandshake, TerminalControl, TerminalId};
 use anyhow::{Context, Result};
 use futures_util::{SinkExt, StreamExt};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -13,11 +13,18 @@ use tokio_tungstenite::tungstenite::Message;
 use tracing::{info, warn};
 
 const MAX_TERMINAL_FRAME_BYTES: usize = 64 * 1024;
+const TERMINAL_REPLAY_BYTES: usize = 256 * 1024;
 const PROCESS_SIGNAL_GRACE: Duration = Duration::from_secs(1);
 
-/// Owns cancellation handles for terminal workers started by the control socket.
+struct ActiveTerminal {
+    cancel: oneshot::Sender<()>,
+    relay_credentials: mpsc::UnboundedSender<String>,
+    created_at_unix: u64,
+}
+
+/// Owns PTY workers independently of any individual control-plane connection.
 pub struct TerminalManager {
-    active: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
+    active: Arc<Mutex<HashMap<String, ActiveTerminal>>>,
     max_sessions: usize,
     events: mpsc::UnboundedSender<TerminalManagerEvent>,
 }
@@ -54,6 +61,11 @@ impl TerminalManager {
 
         let terminal_key = terminal_id.to_string();
         let (cancel_tx, cancel_rx) = oneshot::channel();
+        let (relay_tx, relay_rx) = mpsc::unbounded_channel();
+        let created_at_unix = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs();
         {
             let mut active = self.active.lock().expect("terminal manager poisoned");
             if active.contains_key(&terminal_key) {
@@ -62,15 +74,23 @@ impl TerminalManager {
             if active.len() >= self.max_sessions {
                 anyhow::bail!("agent terminal session limit reached");
             }
-            active.insert(terminal_key.clone(), cancel_tx);
+            active.insert(
+                terminal_key.clone(),
+                ActiveTerminal {
+                    cancel: cancel_tx,
+                    relay_credentials: relay_tx.clone(),
+                    created_at_unix,
+                },
+            );
         }
+        let _ = relay_tx.send(relay_token);
 
         let config = config.clone();
         let active = Arc::downgrade(&self.active);
         let events = self.events.clone();
         tokio::spawn(async move {
             if let Err(err) =
-                run_terminal(&config, &terminal_id, &relay_token, rows, cols, cancel_rx).await
+                run_terminal(&config, &terminal_id, rows, cols, cancel_rx, relay_rx).await
             {
                 warn!(terminal_id = %terminal_id, error = %err, "terminal worker failed");
                 let _ = events.send(TerminalManagerEvent {
@@ -88,7 +108,34 @@ impl TerminalManager {
             .lock()
             .expect("terminal manager poisoned")
             .remove(terminal_id.as_str())
-            .is_some_and(|cancel| cancel.send(()).is_ok())
+            .is_some_and(|active| active.cancel.send(()).is_ok())
+    }
+
+    pub fn resume(&self, terminal_id: &TerminalId, relay_token: String) -> Result<()> {
+        let active = self.active.lock().expect("terminal manager poisoned");
+        let session = active
+            .get(terminal_id.as_str())
+            .with_context(|| format!("terminal session {terminal_id} is not active"))?;
+        session
+            .relay_credentials
+            .send(relay_token)
+            .map_err(|_| anyhow::anyhow!("terminal session {terminal_id} has stopped"))
+    }
+
+    pub fn sessions(&self) -> Vec<AgentTerminalSession> {
+        self.active
+            .lock()
+            .expect("terminal manager poisoned")
+            .iter()
+            .filter_map(|(terminal_id, active)| {
+                TerminalId::new(terminal_id.clone())
+                    .ok()
+                    .map(|terminal_id| AgentTerminalSession {
+                        terminal_id,
+                        created_at_unix: active.created_at_unix,
+                    })
+            })
+            .collect()
     }
 }
 
@@ -97,14 +144,14 @@ impl Drop for TerminalManager {
         if Arc::strong_count(&self.active) == 1
             && let Ok(mut active) = self.active.lock()
         {
-            for (_, cancel) in active.drain() {
-                let _ = cancel.send(());
+            for (_, active) in active.drain() {
+                let _ = active.cancel.send(());
             }
         }
     }
 }
 
-fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>, terminal_id: &str) {
+fn remove_completed(active: &Weak<Mutex<HashMap<String, ActiveTerminal>>>, terminal_id: &str) {
     if let Some(active) = active.upgrade()
         && let Ok(mut active) = active.lock()
     {
@@ -116,28 +163,11 @@ fn remove_completed(active: &Weak<Mutex<HashMap<String, oneshot::Sender<()>>>>,
 async fn run_terminal(
     config: &AgentConfig,
     terminal_id: &TerminalId,
-    relay_token: &str,
     rows: u16,
     cols: u16,
     mut cancel: oneshot::Receiver<()>,
+    mut relay_credentials: mpsc::UnboundedReceiver<String>,
 ) -> Result<()> {
-    let ws_url = terminal_websocket_url(&config.server_url, terminal_id)?;
-    let (stream, _) = tokio::select! {
-        _ = &mut cancel => return Ok(()),
-        result = tokio_tungstenite::connect_async(ws_url.as_str()) => {
-            result.context("failed to connect terminal relay websocket")?
-        }
-    };
-    let (mut sink, mut source) = stream.split();
-    send_json(
-        &mut sink,
-        &TerminalAgentHandshake::Auth {
-            agent_id: config.agent_id.clone(),
-            relay_token: relay_token.to_string(),
-        },
-    )
-    .await?;
-
     let terminal = match wakey::wakey_linux::terminal::TerminalPty::spawn(
         Path::new(&config.terminal.shell),
         rows,
@@ -145,15 +175,6 @@ async fn run_terminal(
     ) {
         Ok(terminal) => terminal,
         Err(err) => {
-            let _ = send_json(
-                &mut sink,
-                &TerminalControl::Error {
-                    code: "terminal_spawn_failed".into(),
-                    message: err.to_string(),
-                },
-            )
-            .await;
-            let _ = sink.send(Message::Close(None)).await;
             return Err(err);
         }
     };
@@ -163,9 +184,14 @@ async fn run_terminal(
         mut child,
     } = terminal;
     let process_group = child.id();
-    send_json(&mut sink, &TerminalControl::Ready).await?;
     info!(terminal_id = %terminal_id, shell = %config.terminal.shell.display(), "terminal PTY ready");
 
+    let (relay_input_tx, mut relay_input_rx) = mpsc::unbounded_channel();
+    let mut relay_output: Option<mpsc::Sender<Message>> = None;
+    let mut relay_task: Option<tokio::task::JoinHandle<()>> = None;
+    let mut relay_generation = 0_u64;
+    let mut replay = VecDeque::new();
+    let mut replay_bytes = 0_usize;
     let mut output = [0_u8; 16 * 1024];
     let mut requested_close = false;
     let mut observed_status = None;
@@ -182,48 +208,78 @@ async fn run_terminal(
             read = reader.read(&mut output) => {
                 match read {
                     Ok(0) => break,
-                    Ok(count) => sink
-                        .send(Message::Binary(output[..count].to_vec().into()))
-                        .await
-                        .context("failed to send PTY output")?,
+                    Ok(count) => send_terminal_output(
+                        Message::Binary(output[..count].to_vec().into()),
+                        &mut relay_output,
+                        &mut replay,
+                        &mut replay_bytes,
+                    ).await,
                     // Linux PTY masters commonly report EIO after the slave closes.
                     Err(err) if err.raw_os_error() == Some(5) => break,
                     Err(err) => return Err(err).context("failed to read PTY output"),
                 }
             }
-            incoming = source.next() => {
+            credential = relay_credentials.recv() => {
+                let Some(relay_token) = credential else { break; };
+                if let Some(task) = relay_task.take() {
+                    task.abort();
+                }
+                relay_generation = relay_generation.wrapping_add(1);
+                let generation = relay_generation;
+                let (output_tx, output_rx) = mpsc::channel(32);
+                relay_output = None;
+                let initial_replay = replay.drain(..).collect();
+                replay_bytes = 0;
+                let config = config.clone();
+                let terminal_id = terminal_id.clone();
+                let relay_input_tx = relay_input_tx.clone();
+                relay_task = Some(tokio::spawn(async move {
+                    if let Err(err) = run_terminal_relay(RelayConnection {
+                        config,
+                        terminal_id: terminal_id.clone(),
+                        relay_token,
+                        generation,
+                        initial_replay,
+                        output_tx,
+                        output_rx,
+                        input: relay_input_tx.clone(),
+                    }).await {
+                        warn!(terminal_id = %terminal_id, error = %err, "terminal relay disconnected");
+                    }
+                    let _ = relay_input_tx.send(RelayInput::Disconnected { generation });
+                }));
+            }
+            incoming = relay_input_rx.recv() => {
                 let Some(message) = incoming else { break; };
-                match message.context("terminal relay websocket receive failed")? {
-                    Message::Binary(bytes) => {
+                match message {
+                    RelayInput::Binary { generation, bytes } if generation == relay_generation => {
                         if bytes.len() > MAX_TERMINAL_FRAME_BYTES {
                             anyhow::bail!("terminal input frame exceeds size limit");
                         }
                         writer.write_all(&bytes).await.context("failed to write PTY input")?;
                     }
-                    Message::Text(text) => {
-                        match serde_json::from_str::<TerminalControl>(&text)
-                            .context("invalid terminal control frame")?
-                        {
-                            TerminalControl::Resize { rows, cols } => {
-                                validate_size(rows, cols)?;
-                                wakey::wakey_linux::terminal::resize_terminal(
-                                    &writer, rows, cols,
-                                )?;
-                            }
-                            TerminalControl::Close => {
-                                requested_close = true;
-                                break;
-                            }
-                            _ => anyhow::bail!("terminal control frame has invalid direction"),
-                        }
+                    RelayInput::Resize { generation, rows, cols } if generation == relay_generation => {
+                        validate_size(rows, cols)?;
+                        wakey::wakey_linux::terminal::resize_terminal(&writer, rows, cols)?;
                     }
-                    Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
-                    Message::Pong(_) => {}
-                    Message::Close(_) => {
+                    RelayInput::Close { generation } if generation == relay_generation => {
                         requested_close = true;
                         break;
                     }
-                    Message::Frame(_) => {}
+                    RelayInput::Connected { generation, output } if generation == relay_generation => {
+                        relay_output = Some(output.clone());
+                        while let Some(frame) = replay.pop_front() {
+                            replay_bytes = replay_bytes.saturating_sub(message_size(&frame));
+                            if output.send(frame).await.is_err() {
+                                relay_output = None;
+                                break;
+                            }
+                        }
+                    }
+                    RelayInput::Disconnected { generation } if generation == relay_generation => {
+                        relay_output = None;
+                    }
+                    _ => {}
                 }
             }
         }
@@ -233,26 +289,173 @@ async fn run_terminal(
         Some(status) => status,
         None => terminate_process_group(&mut child, process_group).await?,
     };
-    let _ = send_json(
-        &mut sink,
-        &TerminalControl::Exited {
+    if let Some(output) = relay_output {
+        let control = TerminalControl::Exited {
             exit_code: status.code(),
-        },
-    )
-    .await;
-    let _ = sink.send(Message::Close(None)).await;
+        };
+        if let Ok(text) = serde_json::to_string(&control) {
+            let _ = output.send(Message::Text(text.into())).await;
+        }
+    }
     info!(terminal_id = %terminal_id, exit_code = ?status.code(), requested_close, "terminal worker exited");
     Ok(())
 }
 
+enum RelayInput {
+    Connected {
+        generation: u64,
+        output: mpsc::Sender<Message>,
+    },
+    Binary {
+        generation: u64,
+        bytes: Vec<u8>,
+    },
+    Resize {
+        generation: u64,
+        rows: u16,
+        cols: u16,
+    },
+    Close {
+        generation: u64,
+    },
+    Disconnected {
+        generation: u64,
+    },
+}
+
+#[cfg(unix)]
+struct RelayConnection {
+    config: AgentConfig,
+    terminal_id: TerminalId,
+    relay_token: String,
+    generation: u64,
+    initial_replay: Vec<Message>,
+    output_tx: mpsc::Sender<Message>,
+    output_rx: mpsc::Receiver<Message>,
+    input: mpsc::UnboundedSender<RelayInput>,
+}
+
+#[cfg(unix)]
+async fn run_terminal_relay(relay: RelayConnection) -> Result<()> {
+    let ws_url = terminal_websocket_url(&relay.config.server_url, &relay.terminal_id)?;
+    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();
+    send_json(
+        &mut sink,
+        &TerminalAgentHandshake::Auth {
+            agent_id: relay.config.agent_id.clone(),
+            relay_token: relay.relay_token,
+        },
+    )
+    .await?;
+    send_json(&mut sink, &TerminalControl::Ready).await?;
+    for frame in relay.initial_replay {
+        sink.send(frame)
+            .await
+            .context("failed to replay detached terminal output")?;
+    }
+    relay
+        .input
+        .send(RelayInput::Connected {
+            generation: relay.generation,
+            output: relay.output_tx,
+        })
+        .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
+
+    let mut output = relay.output_rx;
+    loop {
+        tokio::select! {
+            outgoing = output.recv() => {
+                let Some(message) = outgoing else { break; };
+                sink.send(message).await.context("failed to send terminal relay output")?;
+            }
+            incoming = source.next() => {
+                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,
+                            bytes: bytes.to_vec(),
+                        }).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,
+                                rows,
+                                cols,
+                            })
+                                .map_err(|_| anyhow::anyhow!("terminal worker stopped"))?;
+                        }
+                        TerminalControl::Close => {
+                            let _ = relay.input.send(RelayInput::Close {
+                                generation: relay.generation,
+                            });
+                            break;
+                        }
+                        _ => anyhow::bail!("terminal control frame has invalid direction"),
+                    },
+                    Message::Ping(payload) => sink.send(Message::Pong(payload)).await?,
+                    Message::Pong(_) => {}
+                    // Transport closure only detaches the relay. The agent-owned
+                    // PTY remains alive and waits for replacement credentials.
+                    Message::Close(_) => break,
+                    Message::Frame(_) => {}
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
+async fn send_terminal_output(
+    frame: Message,
+    relay: &mut Option<mpsc::Sender<Message>>,
+    replay: &mut VecDeque<Message>,
+    replay_bytes: &mut usize,
+) {
+    if let Some(tx) = relay.as_ref() {
+        if let Err(error) = tx.send(frame).await {
+            *relay = None;
+            push_local_replay(error.0, replay, replay_bytes);
+        }
+    } else {
+        push_local_replay(frame, replay, replay_bytes);
+    }
+}
+
+fn push_local_replay(frame: Message, replay: &mut VecDeque<Message>, replay_bytes: &mut usize) {
+    *replay_bytes += message_size(&frame);
+    replay.push_back(frame);
+    while *replay_bytes > TERMINAL_REPLAY_BYTES {
+        if let Some(dropped) = replay.pop_front() {
+            *replay_bytes = replay_bytes.saturating_sub(message_size(&dropped));
+        } else {
+            break;
+        }
+    }
+}
+
+fn message_size(message: &Message) -> usize {
+    match message {
+        Message::Text(text) => text.len(),
+        Message::Binary(bytes) | Message::Ping(bytes) | Message::Pong(bytes) => bytes.len(),
+        Message::Close(_) | Message::Frame(_) => 0,
+    }
+}
+
 #[cfg(not(unix))]
 async fn run_terminal(
     _config: &AgentConfig,
     _terminal_id: &TerminalId,
-    _relay_token: &str,
     _rows: u16,
     _cols: u16,
     _cancel: oneshot::Receiver<()>,
+    _relay_credentials: mpsc::UnboundedReceiver<String>,
 ) -> Result<()> {
     anyhow::bail!("terminal sessions are unsupported on this platform")
 }
@@ -342,4 +545,27 @@ mod tests {
             "wss://example.com/api/v1/agent/terminals/term-1/ws"
         );
     }
+
+    #[test]
+    fn detached_replay_drops_oldest_output_at_bound() {
+        let mut replay = VecDeque::new();
+        let mut replay_bytes = 0;
+        for marker in 0_u8..10 {
+            push_local_replay(
+                Message::Binary(vec![marker; TERMINAL_REPLAY_BYTES / 4].into()),
+                &mut replay,
+                &mut replay_bytes,
+            );
+        }
+
+        assert!(replay_bytes <= TERMINAL_REPLAY_BYTES);
+        assert_eq!(replay.len(), 4);
+        assert_eq!(
+            replay.front().and_then(|frame| match frame {
+                Message::Binary(bytes) => bytes.first().copied(),
+                _ => None,
+            }),
+            Some(6)
+        );
+    }
 }