Add optional support for raw scrollback with ansi (#4617)

* scrollback with ansi support

* add command to kdl parser

* removed an unused esc (actually could cause some visualization issue)

* correctly reset style after new lines

* remove redundant reset

* preserve styled spaces

* more readable code

9 files changed, 136 insertions, 1 deletions

diff --git a/zellij-server/src/panes/grid.rs b/zellij-server/src/panes/grid.rs
index 58b10cde..262c5040 100644
--- a/zellij-server/src/panes/grid.rs
+++ b/zellij-server/src/panes/grid.rs
@@ -288,6 +288,45 @@ macro_rules! dump_screen {
     }};
 }
 
+macro_rules! dump_screen_with_ansi {
+    ($lines:expr) => {{
+        use std::fmt::Write;
+        let mut is_first = true;
+        let mut buf = String::new();
+        let mut last_styles: Option<RcCharacterStyles> = None;
+
+        for line in &$lines {
+            if line.is_canonical && !is_first {
+                buf.push_str("\n");
+                last_styles = None;
+            }
+
+            let last_non_space = line
+                .columns
+                .iter()
+                .rposition(|tc| {
+                    let space = tc.character == ' ';
+                    let styled = !matches!(tc.styles.background, Some(AnsiCode::Reset) | None);
+                    !space || styled // it's, something drawable
+                })
+                .map(|i| i + 1)
+                .unwrap_or(0);
+
+            for tc in line.columns.iter().take(last_non_space) {
+                // Only output style codes if style changed
+                if last_styles.as_ref() != Some(&tc.styles) {
+                    write!(buf, "{}", tc.styles).unwrap();
+                    last_styles = Some(tc.styles.clone());
+                }
+                buf.push(tc.character);
+            }
+            is_first = false;
+        }
+        buf.push_str("\u{1b}[m");
+        buf
+    }};
+}
+
 fn utf8_mouse_coordinates(column: usize, line: isize) -> Vec<u8> {
     let mut coordinates = vec![];
     let mouse_pos_encode = |pos: usize| -> Vec<u8> {
@@ -1244,6 +1283,19 @@ impl Grid {
         scrollback.push_str(&viewport);
         scrollback
     }
+    /// Dumps all lines (with ansi) above terminal viewport and the viewport itself to a string
+    pub fn dump_screen_with_ansi(&self, full: bool) -> String {
+        let viewport: String = dump_screen_with_ansi!(self.viewport);
+        if !full {
+            return viewport;
+        }
+        let mut scrollback: String = dump_screen_with_ansi!(self.lines_above);
+        if !scrollback.is_empty() {
+            scrollback.push('\n');
+        }
+        scrollback.push_str(&viewport);
+        scrollback
+    }
     pub fn move_viewport_up(&mut self, count: usize) {
         for _ in 0..count {
             self.scroll_up_one_line();
diff --git a/zellij-server/src/panes/terminal_pane.rs b/zellij-server/src/panes/terminal_pane.rs
index 37a01678..87fac0ae 100644
--- a/zellij-server/src/panes/terminal_pane.rs
+++ b/zellij-server/src/panes/terminal_pane.rs
@@ -537,6 +537,9 @@ impl Pane for TerminalPane {
     fn dump_screen(&self, full: bool, _client_id: Option<ClientId>) -> String {
         self.grid.dump_screen(full)
     }
+    fn dump_screen_with_ansi(&self, full: bool, _client_id: Option<ClientId>) -> String {
+        self.grid.dump_screen_with_ansi(full)
+    }
     fn clear_screen(&mut self) {
         self.grid.clear_screen()
     }
diff --git a/zellij-server/src/route.rs b/zellij-server/src/route.rs
index be3cf492..4d03bd0b 100644
--- a/zellij-server/src/route.rs
+++ b/zellij-server/src/route.rs
@@ -443,6 +443,14 @@ pub(crate) fn route_action(
                 ))
                 .with_context(err_context)?;
         },
+        Action::EditScrollbackRaw => {
+            senders
+                .send_to_screen(ScreenInstruction::EditScrollbackRaw(
+                    client_id,
+                    Some(NotificationEnd::new(completion_tx)),
+                ))
+                .with_context(err_context)?;
+        },
         Action::ScrollUp => {
             senders
                 .send_to_screen(ScreenInstruction::ScrollUp(
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 7f4a4030..19cf0548 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -310,6 +310,7 @@ pub enum ScreenInstruction {
         response_channel: crossbeam::channel::Sender<Option<TabInfo>>,
     },
     EditScrollback(ClientId, Option<NotificationEnd>),
+    EditScrollbackRaw(ClientId, Option<NotificationEnd>),
     GetPaneScrollback {
         pane_id: PaneId,
         client_id: ClientId,
@@ -716,6 +717,7 @@ impl From<&ScreenInstruction> for ScreenContext {
             ScreenInstruction::GetPaneInfo { .. } => ScreenContext::GetPaneInfo,
             ScreenInstruction::GetTabInfo { .. } => ScreenContext::GetTabInfo,
             ScreenInstruction::EditScrollback(..) => ScreenContext::EditScrollback,
+            ScreenInstruction::EditScrollbackRaw(..) => ScreenContext::EditScrollback, // fallback
             ScreenInstruction::GetPaneScrollback { .. } => ScreenContext::GetPaneScrollback,
             ScreenInstruction::ScrollUp(..) => ScreenContext::ScrollUp,
             ScreenInstruction::ScrollDown(..) => ScreenContext::ScrollDown,
@@ -4741,6 +4743,16 @@ pub(crate) fn screen_thread_main(
                 screen.render(None)?;
                 screen.log_and_report_session_state()?;
             },
+            ScreenInstruction::EditScrollbackRaw(client_id, completion_tx) => {
+                active_tab_and_connected_client_id!(
+                    screen,
+                    client_id,
+                    |tab: &mut Tab, client_id: ClientId| tab.edit_scrollback_raw(client_id, completion_tx),
+                    ?
+                );
+                screen.render(None)?;
+                screen.log_and_report_session_state()?;
+            },
             ScreenInstruction::GetPaneScrollback {
                 pane_id,
                 client_id,
diff --git a/zellij-server/src/tab/mod.rs b/zellij-server/src/tab/mod.rs
index 2ec01383..e5f14f32 100644
--- a/zellij-server/src/tab/mod.rs
+++ b/zellij-server/src/tab/mod.rs
@@ -286,6 +286,9 @@ pub trait Pane {
     fn dump_screen(&self, _full: bool, _client_id: Option<ClientId>) -> String {
         "".to_owned()
     }
+    fn dump_screen_with_ansi(&self, _full: bool, _client_id: Option<ClientId>) -> String {
+        "".to_owned()
+    }
     fn scroll_up(&mut self, count: usize, client_id: ClientId);
     fn scroll_down(&mut self, count: usize, client_id: ClientId);
     fn clear_scroll(&mut self);
@@ -3884,6 +3887,23 @@ impl Tab {
         }
         Ok(())
     }
+    pub fn dump_with_ansi_active_terminal_screen(
+        &mut self,
+        file: Option<String>,
+        client_id: ClientId,
+        full: bool,
+    ) -> Result<()> {
+        let err_context =
+            || format!("failed to dump active terminal screen for client {client_id}");
+
+        if let Some(active_pane) = self.get_active_pane_or_floating_pane_mut(client_id) {
+            let dump = active_pane.dump_screen_with_ansi(full, Some(client_id));
+            self.os_api
+                .write_to_file(dump, file)
+                .with_context(err_context)?;
+        }
+        Ok(())
+    }
     pub fn dump_terminal_screen(
         &mut self,
         file: Option<String>,
@@ -3923,6 +3943,33 @@ impl Tab {
             ))
             .with_context(err_context)
     }
+    pub fn edit_scrollback_raw(
+        &mut self,
+        client_id: ClientId,
+        completion_tx: Option<NotificationEnd>,
+    ) -> Result<()> {
+        let err_context = || format!("failed to edit scrollback for client {client_id}");
+
+        let mut file = temp_dir();
+        file.push(format!("{}.dump", Uuid::new_v4()));
+        self.dump_with_ansi_active_terminal_screen(
+            Some(String::from(file.to_string_lossy())),
+            client_id,
+            true,
+        )
+        .with_context(err_context)?;
+        let line_number = self
+            .get_active_pane(client_id)
+            .and_then(|a_t| a_t.get_line_number());
+        self.senders
+            .send_to_pty(PtyInstruction::OpenInPlaceEditor(
+                file,
+                line_number,
+                ClientTabIndexOrPaneId::ClientId(client_id),
+                completion_tx,
+            ))
+            .with_context(err_context)
+    }
     pub fn edit_scrollback_for_pane_with_id(
         &mut self,
         pane_id: PaneId,
diff --git a/zellij-utils/src/input/actions.rs b/zellij-utils/src/input/actions.rs
index f53a5c69..cb49c433 100644
--- a/zellij-utils/src/input/actions.rs
+++ b/zellij-utils/src/input/actions.rs
@@ -178,8 +178,9 @@ pub enum Action {
     DumpLayout,
     /// Save the current session state to disk
     SaveSession,
-    /// Scroll up in focus pane.
     EditScrollback,
+    EditScrollbackRaw,
+    /// Scroll up in focus pane.
     ScrollUp,
     /// Scroll up at point
     ScrollUpAt {
diff --git a/zellij-utils/src/ipc/protobuf_conversion.rs b/zellij-utils/src/ipc/protobuf_conversion.rs
index 023e0503..7ff73969 100644
--- a/zellij-utils/src/ipc/protobuf_conversion.rs
+++ b/zellij-utils/src/ipc/protobuf_conversion.rs
@@ -831,6 +831,9 @@ impl From<crate::input::actions::Action>
             crate::input::actions::Action::EditScrollback => {
                 ActionType::EditScrollback(EditScrollbackAction {})
             },
+            crate::input::actions::Action::EditScrollbackRaw => {
+                ActionType::EditScrollback(EditScrollbackAction {}) // fallback to default edit scrollback
+            },
             crate::input::actions::Action::ScrollUp => ActionType::ScrollUp(ScrollUpAction {}),
             crate::input::actions::Action::ScrollUpAt { position } => {
                 ActionType::ScrollUpAt(ScrollUpAtAction {
diff --git a/zellij-utils/src/kdl/mod.rs b/zellij-utils/src/kdl/mod.rs
index 6915c79f..99e539cc 100644
--- a/zellij-utils/src/kdl/mod.rs
+++ b/zellij-utils/src/kdl/mod.rs
@@ -49,6 +49,7 @@ macro_rules! parse_kdl_action_arguments {
                 "FocusPreviousPane" => Ok(Action::FocusPreviousPane),
                 "SwitchFocus" => Ok(Action::SwitchFocus),
                 "EditScrollback" => Ok(Action::EditScrollback),
+                "EditScrollbackRaw" => Ok(Action::EditScrollbackRaw),
                 "ScrollUp" => Ok(Action::ScrollUp),
                 "ScrollDown" => Ok(Action::ScrollDown),
                 "ScrollToBottom" => Ok(Action::ScrollToBottom),
@@ -697,6 +698,7 @@ impl Action {
             },
             Action::DumpLayout => Some(KdlNode::new("DumpLayout")),
             Action::EditScrollback => Some(KdlNode::new("EditScrollback")),
+            Action::EditScrollbackRaw => Some(KdlNode::new("EditScrollbackRaw")),
             Action::ScrollUp => Some(KdlNode::new("ScrollUp")),
             Action::ScrollDown => Some(KdlNode::new("ScrollDown")),
             Action::ScrollToBottom => Some(KdlNode::new("ScrollToBottom")),
@@ -1453,6 +1455,9 @@ impl TryFrom<(&KdlNode, &Options)> for Action {
             "EditScrollback" => {
                 parse_kdl_action_arguments!(action_name, action_arguments, kdl_action)
             },
+            "EditScrollbackRaw" => {
+                parse_kdl_action_arguments!(action_name, action_arguments, kdl_action)
+            },
             "ScrollUp" => parse_kdl_action_arguments!(action_name, action_arguments, kdl_action),
             "ScrollDown" => parse_kdl_action_arguments!(action_name, action_arguments, kdl_action),
             "ScrollToBottom" => {
diff --git a/zellij-utils/src/plugin_api/action.rs b/zellij-utils/src/plugin_api/action.rs
index 3e472a30..0463a40d 100644
--- a/zellij-utils/src/plugin_api/action.rs
+++ b/zellij-utils/src/plugin_api/action.rs
@@ -1132,6 +1132,10 @@ impl TryFrom<Action> for ProtobufAction {
                 name: ProtobufActionName::EditScrollback as i32,
                 optional_payload: None,
             }),
+            Action::EditScrollbackRaw => Ok(ProtobufAction {
+                name: ProtobufActionName::EditScrollback as i32, // fallback to default edit scrollback
+                optional_payload: None,
+            }),
             Action::ScrollUp => Ok(ProtobufAction {
                 name: ProtobufActionName::ScrollUp as i32,
                 optional_payload: None,