feat: add direct terminal attach stream

1 files changed, 377 insertions, 35 deletions

diff --git a/src/server/headless.rs b/src/server/headless.rs
index fbfef251..d1c8d144 100644
--- a/src/server/headless.rs
+++ b/src/server/headless.rs
@@ -28,6 +28,7 @@ use tokio::sync::mpsc;
 use tracing::{debug, error, info, warn};
 
 use base64::Engine;
+use bytes::Bytes;
 
 use crate::api;
 use crate::app;
@@ -200,8 +201,24 @@ fn derive_client_socket_from_api_socket(api_socket_path: &Path) -> PathBuf {
 // Connected client state
 // ---------------------------------------------------------------------------
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum ClientConnectionMode {
+    App,
+    TerminalAttach { terminal_id: String },
+}
+
+type RenderTarget = (
+    u64,
+    (u16, u16),
+    crate::kitty_graphics::HostCellSize,
+    bool,
+    ClientConnectionMode,
+);
+
 /// A connected client tracked by the server.
 struct ClientConnection {
+    /// Whether this connection is the full app client or a direct terminal attach.
+    mode: ClientConnectionMode,
     /// The client's terminal size (after clamping).
     terminal_size: (u16, u16),
     /// Pixel size of one client terminal cell.
@@ -225,6 +242,7 @@ struct ClientConnection {
 }
 
 impl ClientConnection {
+    #[cfg(test)]
     fn new(
         terminal_size: (u16, u16),
         cell_size: crate::kitty_graphics::HostCellSize,
@@ -234,7 +252,30 @@ impl ClientConnection {
         render_encoding: RenderEncoding,
         writer: Option<ClientWriter>,
     ) -> Self {
+        Self::new_with_mode(
+            ClientConnectionMode::App,
+            terminal_size,
+            cell_size,
+            host_terminal_theme,
+            outer_terminal_focus,
+            last_activity,
+            render_encoding,
+            writer,
+        )
+    }
+
+    fn new_with_mode(
+        mode: ClientConnectionMode,
+        terminal_size: (u16, u16),
+        cell_size: crate::kitty_graphics::HostCellSize,
+        host_terminal_theme: crate::terminal_theme::TerminalTheme,
+        outer_terminal_focus: Option<bool>,
+        last_activity: u64,
+        render_encoding: RenderEncoding,
+        writer: Option<ClientWriter>,
+    ) -> Self {
         Self {
+            mode,
             terminal_size,
             cell_size,
             host_terminal_theme,
@@ -291,6 +332,8 @@ pub struct HeadlessServer {
     next_client_id: u64,
     /// The client currently driving the shared pane runtime size and theme.
     foreground_client_id: Option<u64>,
+    /// Writable direct attach owner per terminal id string.
+    terminal_attach_owners: HashMap<String, u64>,
     /// Monotonic activity counter used to pick the most recently active client.
     next_activity_stamp: u64,
     /// Shared pane runtime size derived from the foreground client,
@@ -336,6 +379,7 @@ impl HeadlessServer {
             clients: HashMap::new(),
             next_client_id: 1,
             foreground_client_id: None,
+            terminal_attach_owners: HashMap::new(),
             next_activity_stamp: 1,
             effective_size: (MIN_COLS, MIN_ROWS),
             shutting_down: false,
@@ -596,6 +640,7 @@ impl HeadlessServer {
         let next_foreground = self
             .clients
             .iter()
+            .filter(|(_, client)| matches!(client.mode, ClientConnectionMode::App))
             .max_by_key(|(_, client)| client.last_activity)
             .map(|(&client_id, _)| client_id);
         let changed = next_foreground != self.foreground_client_id;
@@ -607,7 +652,20 @@ impl HeadlessServer {
     fn remove_client(&mut self, client_id: u64) -> bool {
         let was_foreground = self.foreground_client_id == Some(client_id);
         self.send_client_graphics_cleanup(client_id);
-        self.clients.remove(&client_id);
+        let removed = self.clients.remove(&client_id);
+        if let Some(ClientConnection {
+            mode: ClientConnectionMode::TerminalAttach { terminal_id },
+            ..
+        }) = removed
+        {
+            self.terminal_attach_owners.remove(&terminal_id);
+            if let Some(terminal_id) = self.terminal_id_by_string(&terminal_id) {
+                self.app
+                    .state
+                    .direct_attach_resize_locks
+                    .remove(&terminal_id);
+            }
+        }
         if was_foreground {
             self.promote_latest_remaining_client()
         } else {
@@ -759,6 +817,23 @@ impl HeadlessServer {
         changed
     }
 
+    fn terminal_id_by_string(&self, terminal_id: &str) -> Option<crate::terminal::TerminalId> {
+        self.app
+            .state
+            .terminals
+            .keys()
+            .find(|id| id.to_string() == terminal_id)
+            .cloned()
+    }
+
+    fn runtime_for_terminal_id_string(
+        &self,
+        terminal_id: &str,
+    ) -> Option<&crate::terminal::TerminalRuntime> {
+        let terminal_id = self.terminal_id_by_string(terminal_id)?;
+        self.app.state.terminal_runtimes.get(&terminal_id)
+    }
+
     fn pane_effective_state(&self, pane_id: crate::layout::PaneId) -> crate::detect::AgentState {
         self.app
             .state
@@ -1000,6 +1075,26 @@ impl HeadlessServer {
 
                 true
             }
+            AppEvent::PaneDied { pane_id } => {
+                let terminal_id = self.app.state.workspaces.iter().find_map(|ws| {
+                    ws.tabs.iter().find_map(|tab| {
+                        tab.panes
+                            .get(pane_id)
+                            .map(|pane| pane.attached_terminal_id.to_string())
+                    })
+                });
+
+                self.app.handle_internal_event(ev);
+
+                if let Some(terminal_id) = terminal_id {
+                    self.shutdown_terminal_attach_clients(
+                        &terminal_id,
+                        format!("terminal {terminal_id} exited"),
+                    );
+                }
+
+                true
+            }
             _ => {
                 self.app.handle_internal_event(ev);
                 true
@@ -1130,6 +1225,104 @@ impl HeadlessServer {
         }
     }
 
+    fn shutdown_terminal_attach_clients(&mut self, terminal_id: &str, reason: String) {
+        let client_ids: Vec<u64> = self
+            .clients
+            .iter()
+            .filter_map(|(&client_id, client)| match &client.mode {
+                ClientConnectionMode::TerminalAttach {
+                    terminal_id: attached,
+                } if attached == terminal_id => Some(client_id),
+                _ => None,
+            })
+            .collect();
+
+        for client_id in client_ids {
+            self.send_to_client(
+                client_id,
+                ServerMessage::ServerShutdown {
+                    reason: Some(reason.clone()),
+                },
+            );
+            let foreground_changed = self.remove_client(client_id);
+            if foreground_changed {
+                self.resize_shared_runtime_to_effective_size();
+            }
+        }
+    }
+
+    fn attach_terminal_client(
+        &mut self,
+        client_id: u64,
+        terminal_id: String,
+        takeover: bool,
+    ) -> bool {
+        let Some(real_terminal_id) = self.terminal_id_by_string(&terminal_id) else {
+            self.send_to_client(
+                client_id,
+                ServerMessage::ServerShutdown {
+                    reason: Some(format!(
+                        "terminal attach failed: terminal {terminal_id} not found"
+                    )),
+                },
+            );
+            self.remove_client(client_id);
+            return false;
+        };
+
+        if let Some(existing_owner) = self.terminal_attach_owners.get(&terminal_id).copied() {
+            if existing_owner != client_id && !takeover {
+                self.send_to_client(
+                    client_id,
+                    ServerMessage::ServerShutdown {
+                        reason: Some(format!(
+                            "terminal attach failed: terminal {terminal_id} already has an attached client; retry with --takeover"
+                        )),
+                    },
+                );
+                self.remove_client(client_id);
+                return false;
+            }
+            if existing_owner != client_id {
+                self.send_to_client(
+                    existing_owner,
+                    ServerMessage::ServerShutdown {
+                        reason: Some("terminal attach taken over".to_owned()),
+                    },
+                );
+                self.remove_client(existing_owner);
+            }
+        }
+
+        let stamp = self.allocate_activity_stamp();
+        let Some(client) = self.clients.get_mut(&client_id) else {
+            return false;
+        };
+        let (cols, rows) = client.terminal_size;
+        let cell_size = client.cell_size;
+        client.mode = ClientConnectionMode::TerminalAttach {
+            terminal_id: terminal_id.clone(),
+        };
+        client.render_state.reset_baseline();
+        client.last_activity = stamp;
+        let was_foreground = self.foreground_client_id == Some(client_id);
+        if was_foreground {
+            self.promote_latest_remaining_client();
+        }
+
+        info!(client_id, cols, rows, terminal_id = %terminal_id, "terminal attach client connected");
+        self.terminal_attach_owners
+            .insert(terminal_id.clone(), client_id);
+        self.app
+            .state
+            .direct_attach_resize_locks
+            .insert(real_terminal_id.clone());
+        if let Some(runtime) = self.app.state.terminal_runtimes.get(&real_terminal_id) {
+            runtime.resize(rows, cols, cell_size.width_px, cell_size.height_px);
+        }
+        true
+    }
+
     /// Handles a server event. Returns true if the event requires a re-render.
     fn handle_server_event(&mut self, ev: ServerEvent) -> bool {
         match ev {
@@ -1154,7 +1347,8 @@ impl HeadlessServer {
                 let last_activity = self.allocate_activity_stamp();
                 self.clients.insert(
                     client_id,
-                    ClientConnection::new(
+                    ClientConnection::new_with_mode(
+                        ClientConnectionMode::App,
                         (cols, rows),
                         crate::kitty_graphics::HostCellSize {
                             width_px: cell_width_px,
@@ -1172,8 +1366,25 @@ impl HeadlessServer {
                 self.resize_shared_runtime_to_effective_size();
                 true
             }
+            ServerEvent::ClientAttachTerminal {
+                client_id,
+                terminal_id,
+                takeover,
+            } => self.attach_terminal_client(client_id, terminal_id, takeover),
             ServerEvent::ClientInput { client_id, data } => {
                 debug!(client_id, len = data.len(), "client input received");
+                if let Some(ClientConnection {
+                    mode: ClientConnectionMode::TerminalAttach { terminal_id },
+                    ..
+                }) = self.clients.get(&client_id)
+                {
+                    if let Some(runtime) = self.runtime_for_terminal_id_string(terminal_id) {
+                        if let Err(err) = runtime.try_send_bytes(Bytes::from(data)) {
+                            warn!(client_id, terminal_id = %terminal_id, err = %err, "terminal attach input failed");
+                        }
+                    }
+                    return true;
+                }
                 let events = crate::raw_input::parse_raw_input_bytes_sync(&data);
                 let host_surface_redraw =
                     crate::raw_input::events_require_host_surface_redraw(&events);
@@ -1251,6 +1462,30 @@ impl HeadlessServer {
                     client_id,
                     cols, rows, cell_width_px, cell_height_px, "client resize"
                 );
+                let direct_terminal_id = if let Some(ClientConnection {
+                    mode: ClientConnectionMode::TerminalAttach { terminal_id },
+                    terminal_size,
+                    cell_size,
+                    render_state,
+                    ..
+                }) = self.clients.get_mut(&client_id)
+                {
+                    *terminal_size = (cols, rows);
+                    *cell_size = crate::kitty_graphics::HostCellSize {
+                        width_px: cell_width_px,
+                        height_px: cell_height_px,
+                    };
+                    render_state.reset_baseline();
+                    Some(terminal_id.clone())
+                } else {
+                    None
+                };
+                if let Some(terminal_id) = direct_terminal_id {
+                    if let Some(runtime) = self.runtime_for_terminal_id_string(&terminal_id) {
+                        runtime.resize(rows, cols, cell_width_px, cell_height_px);
+                    }
+                    return true;
+                }
                 if let Some(client) = self.clients.get_mut(&client_id) {
                     client.terminal_size = (cols, rows);
                     client.cell_size = crate::kitty_graphics::HostCellSize {
@@ -1508,6 +1743,9 @@ impl HeadlessServer {
 
         let mut broken_clients: Vec<u64> = Vec::new();
         for (&client_id, client) in &mut self.clients {
+            if !matches!(client.mode, ClientConnectionMode::App) {
+                continue;
+            }
             if client.host_mouse_capture_active == Some(enabled) {
                 continue;
             }
@@ -1537,21 +1775,23 @@ impl HeadlessServer {
     /// frames to all connected clients.
     fn render_and_stream(&mut self) {
         let foreground_client_id = self.foreground_client_id;
-        let mut render_targets: Vec<(u64, (u16, u16), crate::kitty_graphics::HostCellSize, bool)> =
-            self.clients
-                .iter()
-                .filter(|(_, client)| client.writer.is_some())
-                .map(|(&client_id, client)| {
-                    (
-                        client_id,
-                        client.terminal_size,
-                        client.cell_size,
-                        foreground_client_id == Some(client_id),
-                    )
-                })
-                .collect();
+        let mut render_targets: Vec<RenderTarget> = self
+            .clients
+            .iter()
+            .filter(|(_, client)| client.writer.is_some())
+            .map(|(&client_id, client)| {
+                (
+                    client_id,
+                    client.terminal_size,
+                    client.cell_size,
+                    foreground_client_id == Some(client_id),
+                    client.mode.clone(),
+                )
+            })
+            .collect();
 
-        render_targets.sort_by_key(|(client_id, _, _, is_foreground)| (*is_foreground, *client_id));
+        render_targets
+            .sort_by_key(|(client_id, _, _, is_foreground, _)| (*is_foreground, *client_id));
 
         if render_targets.is_empty() {
             let (cols, rows) = self.effective_size;
@@ -1570,32 +1810,55 @@ impl HeadlessServer {
         }
 
         let mut broken_clients: Vec<u64> = Vec::new();
-        for (client_id, (cols, rows), cell_size, is_foreground) in render_targets {
+        for (client_id, (cols, rows), cell_size, is_foreground, mode) in render_targets {
             let area = Rect::new(0, 0, cols, rows);
-            let (buffer, cursor) = if self.app.state.kitty_graphics_enabled && cell_size.is_known()
-            {
-                crate::server::render_stream::render_virtual_with_cell_size(
-                    &mut self.app.state,
-                    area,
-                    is_foreground,
-                    cell_size,
-                )
-            } else {
-                crate::server::render_stream::render_virtual(
-                    &mut self.app.state,
-                    area,
-                    is_foreground,
-                )
+            let is_app_client = matches!(mode, ClientConnectionMode::App);
+            let mut frame = match mode {
+                ClientConnectionMode::App => {
+                    let (buffer, cursor) =
+                        if self.app.state.kitty_graphics_enabled && cell_size.is_known() {
+                            crate::server::render_stream::render_virtual_with_cell_size(
+                                &mut self.app.state,
+                                area,
+                                is_foreground,
+                                cell_size,
+                            )
+                        } else {
+                            crate::server::render_stream::render_virtual(
+                                &mut self.app.state,
+                                area,
+                                is_foreground,
+                            )
+                        };
+                    let hyperlinks =
+                        crate::server::render_stream::visible_hyperlinks(&self.app.state);
+                    FrameData::from_ratatui_buffer_with_hyperlinks(&buffer, cursor, &hyperlinks)
+                }
+                ClientConnectionMode::TerminalAttach { terminal_id } => {
+                    let Some(runtime) = self.runtime_for_terminal_id_string(&terminal_id) else {
+                        self.send_to_client(
+                            client_id,
+                            ServerMessage::ServerShutdown {
+                                reason: Some(format!(
+                                    "terminal attach ended: terminal {terminal_id} not found"
+                                )),
+                            },
+                        );
+                        broken_clients.push(client_id);
+                        continue;
+                    };
+                    let (buffer, cursor) =
+                        crate::server::render_stream::render_terminal_virtual(runtime, area);
+                    let hyperlinks = runtime.visible_hyperlinks(area);
+                    FrameData::from_ratatui_buffer_with_hyperlinks(&buffer, cursor, &hyperlinks)
+                }
             };
-            let hyperlinks = crate::server::render_stream::visible_hyperlinks(&self.app.state);
-            let mut frame =
-                FrameData::from_ratatui_buffer_with_hyperlinks(&buffer, cursor, &hyperlinks);
 
             let Some(client) = self.clients.get_mut(&client_id) else {
                 continue;
             };
             let mut next_graphics_cache = client.graphics_cache.clone();
-            if self.app.state.kitty_graphics_enabled && cell_size.is_known() {
+            if is_app_client && self.app.state.kitty_graphics_enabled && cell_size.is_known() {
                 frame.graphics = crate::kitty_graphics::encode_local_pane_graphics(
                     &self.app.state,
                     cell_size,
@@ -1999,6 +2262,7 @@ mod tests {
             clients: HashMap::new(),
             next_client_id: 1,
             foreground_client_id: None,
+            terminal_attach_owners: HashMap::new(),
             next_activity_stamp: 1,
             effective_size: (MIN_COLS, MIN_ROWS),
             shutting_down: false,
@@ -2020,6 +2284,13 @@ mod tests {
         }
     }
 
+    fn read_server_shutdown_reason(bytes: Vec<u8>) -> Option<String> {
+        match read_server_message(bytes) {
+            ServerMessage::ServerShutdown { reason } => reason,
+            other => panic!("expected shutdown, got {other:?}"),
+        }
+    }
+
     fn test_client_writer() -> (
         ClientWriter,
         std::sync::mpsc::Receiver<Vec<u8>>,
@@ -2038,6 +2309,77 @@ mod tests {
     }
 
     #[test]
+    fn terminal_attach_rejects_missing_terminal_and_removes_client() {
+        let mut server = test_headless_server();
+        let (writer, control_rx, _render_rx) = test_client_writer();
+
+        assert!(server.handle_server_event(ServerEvent::ClientConnected {
+            client_id: 7,
+            cols: 80,
+            rows: 24,
+            cell_width_px: 0,
+            cell_height_px: 0,
+            render_encoding: RenderEncoding::TerminalAnsi,
+            writer,
+        }));
+        assert!(server.clients.contains_key(&7));
+
+        assert!(
+            !server.handle_server_event(ServerEvent::ClientAttachTerminal {
+                client_id: 7,
+                terminal_id: "term_missing".to_owned(),
+                takeover: false,
+            })
+        );
+        assert!(!server.clients.contains_key(&7));
+        let reason = read_server_shutdown_reason(control_rx.recv().expect("shutdown message"));
+        assert_eq!(
+            reason,
+            Some("terminal attach failed: terminal term_missing not found".to_owned())
+        );
+    }
+
+    #[test]
+    fn terminal_attach_client_exits_when_attached_pane_dies() {
+        let mut server = test_headless_server();
+        let workspace = crate::workspace::Workspace::test_new("attached");
+        let pane_id = workspace.tabs[0].root_pane;
+        server.app.state.workspaces = vec![workspace];
+        server.app.state.ensure_test_terminals();
+        let terminal_id = server.app.state.workspaces[0]
+            .pane_state(pane_id)
+            .expect("pane")
+            .attached_terminal_id
+            .to_string();
+        let (writer, control_rx, _render_rx) = test_client_writer();
+
+        assert!(server.handle_server_event(ServerEvent::ClientConnected {
+            client_id: 7,
+            cols: 80,
+            rows: 24,
+            cell_width_px: 0,
+            cell_height_px: 0,
+            render_encoding: RenderEncoding::TerminalAnsi,
+            writer,
+        }));
+        assert!(
+            server.handle_server_event(ServerEvent::ClientAttachTerminal {
+                client_id: 7,
+                terminal_id: terminal_id.clone(),
+                takeover: false,
+            })
+        );
+        assert_eq!(server.terminal_attach_owners.get(&terminal_id), Some(&7));
+
+        assert!(server.handle_internal_event_with_forwarding(AppEvent::PaneDied { pane_id }));
+
+        assert!(!server.clients.contains_key(&7));
+        assert!(!server.terminal_attach_owners.contains_key(&terminal_id));
+        let reason = read_server_shutdown_reason(control_rx.recv().expect("shutdown message"));
+        assert_eq!(reason, Some(format!("terminal {terminal_id} exited")));
+    }
+
+    #[test]
     fn client_socket_path_derived_from_api_socket_override() {
         let path = client_socket_path_from_overrides(Some("/tmp/test-herdr.sock"), None);
         assert_eq!(path, PathBuf::from("/tmp/test-herdr-client.sock"));