Allow saving current session for resurrection on demand (#4654)

* initial implementation

* session manager ui and plugin api

* remove event

* cleanups and formatting

* add pr

7 files changed, 212 insertions, 5 deletions

diff --git a/zellij-server/src/background_jobs.rs b/zellij-server/src/background_jobs.rs
index 22828e15..371f99b9 100644
--- a/zellij-server/src/background_jobs.rs
+++ b/zellij-server/src/background_jobs.rs
@@ -190,6 +190,15 @@ pub(crate) fn background_jobs_main(
             },
             BackgroundJob::ReportLayoutInfo(session_layout) => {
                 *current_session_layout.lock().unwrap() = session_layout;
+
+                // Update session save time for plugin query
+                let timestamp_millis = std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .unwrap_or_default()
+                    .as_millis() as u64;
+                let _ = bus
+                    .senders
+                    .send_to_plugin(PluginInstruction::UpdateSessionSaveTime(timestamp_millis));
             },
             BackgroundJob::ReadAllSessionInfosOnMachine => {
                 // this job should only be run once and it keeps track of other sessions (as well
@@ -220,6 +229,13 @@ pub(crate) fn background_jobs_main(
                                     current_session_info,
                                     current_session_layout,
                                 );
+
+                                // Send SavedCurrentSession instruction to plugin thread
+                                let timestamp_millis = std::time::SystemTime::now()
+                                    .duration_since(std::time::UNIX_EPOCH)
+                                    .unwrap_or_default()
+                                    .as_millis()
+                                    as u64;
                             }
                             let mut session_infos_on_machine =
                                 read_other_live_session_states(&current_session_name);
@@ -511,7 +527,7 @@ fn job_already_running(
     }
 }
 
-fn write_session_state_to_disk(
+pub fn write_session_state_to_disk(
     current_session_name: String,
     current_session_info: SessionInfo,
     current_session_layout: (String, BTreeMap<String, String>),
diff --git a/zellij-server/src/plugins/mod.rs b/zellij-server/src/plugins/mod.rs
index c6c931c8..a72c1ffa 100644
--- a/zellij-server/src/plugins/mod.rs
+++ b/zellij-server/src/plugins/mod.rs
@@ -207,6 +207,10 @@ pub enum PluginInstruction {
     },
     LayoutListUpdate(Vec<LayoutInfo>, Vec<LayoutWithError>),
     RequestStateUpdateForPlugin(PluginId),
+    UpdateSessionSaveTime(u64), // u64 = milliseconds since UNIX epoch
+    GetLastSessionSaveTime {
+        response_channel: crossbeam::channel::Sender<Option<u64>>,
+    },
     Exit,
 }
 
@@ -263,6 +267,10 @@ impl From<&PluginInstruction> for PluginContext {
             PluginInstruction::RequestStateUpdateForPlugin(..) => {
                 PluginContext::RequestStateUpdateForPlugin
             },
+            PluginInstruction::UpdateSessionSaveTime(..) => PluginContext::UpdateSessionSaveTime,
+            PluginInstruction::GetLastSessionSaveTime { .. } => {
+                PluginContext::GetLastSessionSaveTime
+            },
         }
     }
 }
@@ -1203,6 +1211,14 @@ pub(crate) fn plugin_thread_main(
             PluginInstruction::RequestStateUpdateForPlugin(plugin_id) => {
                 wasm_bridge.state_update_for_plugin(plugin_id);
             },
+            PluginInstruction::UpdateSessionSaveTime(timestamp_millis) => {
+                // Store timestamp in WasmBridge (as Unix epoch for internal use)
+                *wasm_bridge.last_session_save_time.lock().unwrap() = Some(timestamp_millis);
+            },
+            PluginInstruction::GetLastSessionSaveTime { response_channel } => {
+                let timestamp = *wasm_bridge.last_session_save_time.lock().unwrap();
+                let _ = response_channel.send(timestamp);
+            },
             PluginInstruction::Exit => {
                 break;
             },
diff --git a/zellij-server/src/plugins/wasm_bridge.rs b/zellij-server/src/plugins/wasm_bridge.rs
index bdbe25f3..7244a2d3 100644
--- a/zellij-server/src/plugins/wasm_bridge.rs
+++ b/zellij-server/src/plugins/wasm_bridge.rs
@@ -187,6 +187,7 @@ pub struct WasmBridge {
     base_modes: HashMap<ClientId, InputMode>,
     downloader: Downloader,
     previous_pane_render_report: Option<PaneRenderReport>,
+    pub last_session_save_time: Arc<Mutex<Option<u64>>>, // milliseconds since UNIX epoch
 }
 
 impl WasmBridge {
@@ -252,6 +253,7 @@ impl WasmBridge {
             base_modes: HashMap::new(),
             downloader,
             previous_pane_render_report: None,
+            last_session_save_time: Arc::new(Mutex::new(None)),
         }
     }
     pub fn load_plugin(
diff --git a/zellij-server/src/plugins/zellij_exports.rs b/zellij-server/src/plugins/zellij_exports.rs
index 5b03c0e6..184936bd 100644
--- a/zellij-server/src/plugins/zellij_exports.rs
+++ b/zellij-server/src/plugins/zellij_exports.rs
@@ -62,11 +62,12 @@ use zellij_utils::{
         },
         plugin_command::{
             dump_layout_response, dump_session_layout_response, parse_layout_response,
-            ProtobufDeleteLayoutResponse, ProtobufDumpLayoutResponse,
+            save_session_response, ProtobufDeleteLayoutResponse, ProtobufDumpLayoutResponse,
             ProtobufDumpSessionLayoutResponse, ProtobufEditLayoutResponse,
             ProtobufGenerateRandomNameResponse, ProtobufGetFocusedPaneInfoResponse,
             ProtobufGetLayoutDirResponse, ProtobufGetPanePidResponse, ProtobufParseLayoutResponse,
             ProtobufPluginCommand, ProtobufRenameLayoutResponse, ProtobufSaveLayoutResponse,
+            ProtobufSaveSessionResponse,
         },
         plugin_ids::{ProtobufPluginIds, ProtobufZellijVersion},
     },
@@ -129,6 +130,10 @@ fn host_run_plugin_command(mut caller: Caller<'_, PluginEnv>) {
                     PluginCommand::ParseLayout(layout_string) => parse_layout(env, layout_string),
                     PluginCommand::GetLayoutDir => get_layout_dir(env),
                     PluginCommand::GetFocusedPaneInfo => get_focused_pane_info(env),
+                    PluginCommand::SaveSession => save_session(env),
+                    PluginCommand::CurrentSessionLastSavedTime => {
+                        current_session_last_saved_time(env)
+                    },
                     PluginCommand::OpenFile(file_to_open, context) => {
                         open_file(env, file_to_open, context)
                     },
@@ -2564,6 +2569,64 @@ fn dump_session_layout(env: &PluginEnv, tab_index: Option<usize>) {
         .with_context(|| format!("failed to send session layout to plugin {}", env.name()));
 }
 
+fn save_session(env: &PluginEnv) {
+    use save_session_response::Result as SaveSessionResult;
+    let (completion_tx, completion_rx) = oneshot::channel();
+
+    let send_result = env.senders.send_to_screen(ScreenInstruction::SaveSession(
+        env.client_id,
+        Some(NotificationEnd::new(completion_tx)),
+    ));
+
+    let response = if let Err(e) = send_result {
+        ProtobufSaveSessionResponse {
+            result: Some(SaveSessionResult::Error(format!(
+                "Failed to send save request: {}",
+                e
+            ))),
+        }
+    } else {
+        let wait_forever = false;
+        let _result = wait_for_action_completion(completion_rx, "save_session", wait_forever);
+        ProtobufSaveSessionResponse {
+            result: Some(SaveSessionResult::Success(true)),
+        }
+    };
+
+    let _ = wasi_write_object(env, &response.encode_to_vec());
+}
+
+fn current_session_last_saved_time(env: &PluginEnv) {
+    use zellij_utils::plugin_api::plugin_command::ProtobufCurrentSessionLastSavedTimeResponse;
+
+    let (response_tx, response_rx) = crossbeam::channel::bounded(1);
+    let send_result = env
+        .senders
+        .send_to_plugin(PluginInstruction::GetLastSessionSaveTime {
+            response_channel: response_tx,
+        });
+
+    let saved_timestamp_millis = if send_result.is_ok() {
+        response_rx
+            .recv_timeout(std::time::Duration::from_secs(5))
+            .ok()
+            .flatten()
+    } else {
+        None
+    };
+
+    let timestamp_millis = saved_timestamp_millis.map(|saved_time| {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_millis() as u64;
+        now.saturating_sub(saved_time)
+    });
+
+    let response = ProtobufCurrentSessionLastSavedTimeResponse { timestamp_millis };
+    let _ = wasi_write_object(env, &response.encode_to_vec());
+}
+
 fn list_clients(env: &PluginEnv) {
     let _ = env.senders.to_screen.as_ref().map(|sender| {
         sender.send(ScreenInstruction::ListClientsToPlugin(
@@ -3842,7 +3905,9 @@ fn check_command_permission(
         | PluginCommand::GetLayoutDir
         | PluginCommand::GetFocusedPaneInfo
         | PluginCommand::DumpLayout(..)
-        | PluginCommand::ParseLayout(..) => PermissionType::ReadApplicationState,
+        | PluginCommand::ParseLayout(..)
+        | PluginCommand::SaveSession
+        | PluginCommand::CurrentSessionLastSavedTime => PermissionType::ReadApplicationState,
         PluginCommand::RebindKeys { .. } | PluginCommand::Reconfigure(..) => {
             PermissionType::Reconfigure
         },
diff --git a/zellij-server/src/pty.rs b/zellij-server/src/pty.rs
index 82a24b55..587fd18d 100644
--- a/zellij-server/src/pty.rs
+++ b/zellij-server/src/pty.rs
@@ -1,3 +1,4 @@
+use crate::background_jobs::write_session_state_to_disk;
 use crate::background_jobs::BackgroundJob;
 use crate::route::NotificationEnd;
 use crate::terminal_bytes::TerminalBytes;
@@ -19,7 +20,7 @@ use std::{collections::HashMap, os::unix::io::RawFd, path::PathBuf};
 use zellij_utils::{
     data::{
         CommandOrPlugin, Event, FloatingPaneCoordinates, GetPanePidResponse, NewPanePlacement,
-        OriginatingPlugin,
+        OriginatingPlugin, SessionInfo,
     },
     errors::prelude::*,
     errors::{ContextType, PtyContext},
@@ -110,6 +111,12 @@ pub enum PtyInstruction {
         response_channel: crossbeam::channel::Sender<DumpSessionLayoutResponse>,
     },
     LogLayoutToHd(SessionLayoutMetadata),
+    SaveSessionToDisk {
+        session_name: String,
+        session_info: SessionInfo,
+        session_layout_metadata: SessionLayoutMetadata,
+        completion_tx: Option<NotificationEnd>,
+    },
     FillPluginCwd(
         Option<bool>,   // should float
         bool,           // should be opened in place
@@ -160,6 +167,7 @@ impl From<&PtyInstruction> for PtyContext {
             PtyInstruction::DumpLayout(..) => PtyContext::DumpLayout,
             PtyInstruction::DumpLayoutToPlugin { .. } => PtyContext::DumpLayoutToPlugin,
             PtyInstruction::LogLayoutToHd(..) => PtyContext::LogLayoutToHd,
+            PtyInstruction::SaveSessionToDisk { .. } => PtyContext::SaveSessionToDisk,
             PtyInstruction::FillPluginCwd(..) => PtyContext::FillPluginCwd,
             PtyInstruction::ListClientsMetadata(..) => PtyContext::ListClientsMetadata,
             PtyInstruction::Reconfigure { .. } => PtyContext::Reconfigure,
@@ -760,6 +768,41 @@ pub(crate) fn pty_thread_main(mut pty: Pty, layout: Box<Layout>) -> Result<()> {
                     }
                 }
             },
+            PtyInstruction::SaveSessionToDisk {
+                session_name,
+                session_info,
+                mut session_layout_metadata,
+                completion_tx: _completion_tx, // Dropped at end to signal completion
+            } => {
+                pty.populate_session_layout_metadata(&mut session_layout_metadata);
+                match session_serialization::serialize_session_layout(
+                    session_layout_metadata.into(),
+                ) {
+                    Ok(kdl_and_files) => {
+                        write_session_state_to_disk(
+                            session_name,
+                            session_info,
+                            kdl_and_files.clone(),
+                        );
+
+                        // Update session save time for plugin query
+                        let timestamp_millis = std::time::SystemTime::now()
+                            .duration_since(std::time::UNIX_EPOCH)
+                            .unwrap_or_default()
+                            .as_millis() as u64;
+                        let _ = pty.bus.senders.send_to_plugin(
+                            PluginInstruction::UpdateSessionSaveTime(timestamp_millis),
+                        );
+
+                        let _ = pty.bus.senders.send_to_background_jobs(
+                            BackgroundJob::ReportLayoutInfo(kdl_and_files),
+                        );
+                    },
+                    Err(e) => {
+                        log::error!("Failed to serialize layout: {}", e);
+                    },
+                };
+            },
             PtyInstruction::FillPluginCwd(
                 should_float,
                 should_be_open_in_place,
diff --git a/zellij-server/src/route.rs b/zellij-server/src/route.rs
index f643f024..7ab11dc0 100644
--- a/zellij-server/src/route.rs
+++ b/zellij-server/src/route.rs
@@ -389,6 +389,14 @@ pub(crate) fn route_action(
                 ))
                 .with_context(err_context)?;
         },
+        Action::SaveSession => {
+            senders
+                .send_to_screen(ScreenInstruction::SaveSession(
+                    client_id,
+                    Some(NotificationEnd::new(completion_tx)),
+                ))
+                .with_context(err_context)?;
+        },
         Action::EditScrollback => {
             senders
                 .send_to_screen(ScreenInstruction::EditScrollback(
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index 3bbfc670..17e83230 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -262,6 +262,7 @@ pub enum ScreenInstruction {
     DumpScreen(String, ClientId, bool, Option<NotificationEnd>),
     DumpLayout(Option<PathBuf>, ClientId, Option<NotificationEnd>), // PathBuf is the default configured
     // shell
+    SaveSession(ClientId, Option<NotificationEnd>),
     DumpLayoutToPlugin {
         plugin_id: PluginId,
         tab_index: Option<usize>,
@@ -634,6 +635,7 @@ impl From<&ScreenInstruction> for ScreenContext {
             ScreenInstruction::ClearScreen(..) => ScreenContext::ClearScreen,
             ScreenInstruction::DumpScreen(..) => ScreenContext::DumpScreen,
             ScreenInstruction::DumpLayout(..) => ScreenContext::DumpLayout,
+            ScreenInstruction::SaveSession(..) => ScreenContext::SaveSession,
             ScreenInstruction::DumpLayoutToPlugin { .. } => ScreenContext::DumpLayoutToPlugin,
             ScreenInstruction::GetFocusedPaneInfo { .. } => ScreenContext::GetFocusedPaneInfo,
             ScreenInstruction::EditScrollback(..) => ScreenContext::EditScrollback,
@@ -3914,7 +3916,6 @@ pub(crate) fn screen_thread_main(
 
                 let blocking_notification = if set_blocking { completion_tx } else { None };
 
-                log::info!("client_or_tab_index: {:?}", client_or_tab_index);
                 match client_or_tab_index {
                     ClientTabIndexOrPaneId::ClientId(client_id) => {
                         active_tab_and_connected_client_id_with_first_tab_fallback!(screen, client_id, |tab: &mut Tab, client_id: Option<ClientId>| {
@@ -6168,6 +6169,62 @@ pub(crate) fn screen_thread_main(
                     screen.dump_layout_to_hd()?;
                 }
             },
+            ScreenInstruction::SaveSession(_client_id, completion_tx) => {
+                let err_context = || "Failed to save session";
+
+                screen.update_active_pane_ids();
+                let pane_manifest = screen.generate_and_report_pane_state()?;
+                let tab_infos = screen.generate_and_report_tab_state()?;
+
+                #[cfg(not(test))]
+                let (available_layouts, _layout_errors) = Layout::list_available_layouts(
+                    screen.layout_dir.clone(),
+                    &screen.default_layout_name,
+                );
+                #[cfg(test)]
+                let available_layouts = vec![];
+
+                let session_info = SessionInfo {
+                    name: screen.session_name.clone(),
+                    tabs: tab_infos,
+                    panes: pane_manifest,
+                    connected_clients: screen.active_tab_indices.keys().len(),
+                    is_current_session: true,
+                    available_layouts,
+                    web_clients_allowed: screen.web_sharing.web_clients_allowed(),
+                    web_client_count: screen
+                        .connected_clients
+                        .borrow()
+                        .iter()
+                        .filter(|(_client_id, is_web_client)| **is_web_client)
+                        .count(),
+                    plugins: Default::default(),
+                    tab_history: screen.tab_history.clone(),
+                    pane_history: screen
+                        .pane_history
+                        .iter()
+                        .map(|(k, v)| (*k, v.iter().map(|v| (*v).into()).collect()))
+                        .collect(),
+                };
+
+                let session_layout_metadata = if screen.session_serialization {
+                    screen.get_layout_metadata(Some(screen.default_shell.clone()), None)
+                } else {
+                    // Create empty metadata if serialization is disabled
+                    SessionLayoutMetadata::new(screen.default_layout.clone())
+                };
+
+                screen
+                    .bus
+                    .senders
+                    .send_to_pty(PtyInstruction::SaveSessionToDisk {
+                        session_name: screen.session_name.clone(),
+                        session_info,
+                        session_layout_metadata,
+                        completion_tx,
+                    })
+                    .with_context(err_context)?;
+            },
             ScreenInstruction::RenameSession(
                 name,
                 client_id,