fix(startup): recover from race condition that causes Zellij to start in the wrong size (#3218)

* fix(startup): recover from Zellij starting up in the wrong size

* style(fmt): rustfmt

* fix tests

5 files changed, 24 insertions, 2 deletions

diff --git a/zellij-client/src/lib.rs b/zellij-client/src/lib.rs
index 06c2477a8..2d74e0dd4 100644
--- a/zellij-client/src/lib.rs
+++ b/zellij-client/src/lib.rs
@@ -51,6 +51,7 @@ pub(crate) enum ClientInstruction {
     SetSynchronizedOutput(Option<SyncOutput>),
     UnblockCliPipeInput(String),   // String -> pipe name
     CliPipeOutput(String, String), // String -> pipe name, String -> output
+    QueryTerminalSize,
 }
 
 impl From<ServerToClientMsg> for ClientInstruction {
@@ -75,6 +76,7 @@ impl From<ServerToClientMsg> for ClientInstruction {
             ServerToClientMsg::CliPipeOutput(pipe_name, output) => {
                 ClientInstruction::CliPipeOutput(pipe_name, output)
             },
+            ServerToClientMsg::QueryTerminalSize => ClientInstruction::QueryTerminalSize,
         }
     }
 }
@@ -97,6 +99,7 @@ impl From<&ClientInstruction> for ClientContext {
             ClientInstruction::SetSynchronizedOutput(..) => ClientContext::SetSynchronisedOutput,
             ClientInstruction::UnblockCliPipeInput(..) => ClientContext::UnblockCliPipeInput,
             ClientInstruction::CliPipeOutput(..) => ClientContext::CliPipeOutput,
+            ClientInstruction::QueryTerminalSize => ClientContext::QueryTerminalSize,
         }
     }
 }
@@ -499,6 +502,11 @@ pub fn start_client(
             ClientInstruction::SetSynchronizedOutput(enabled) => {
                 synchronised_output = enabled;
             },
+            ClientInstruction::QueryTerminalSize => {
+                os_input.send_to_server(ClientToServerMsg::TerminalResize(
+                    os_input.get_terminal_size_using_fd(0),
+                ));
+            },
             _ => {},
         }
     }
diff --git a/zellij-server/src/screen.rs b/zellij-server/src/screen.rs
index bba6baeb9..94b91777d 100644
--- a/zellij-server/src/screen.rs
+++ b/zellij-server/src/screen.rs
@@ -3000,6 +3000,18 @@ pub(crate) fn screen_thread_main(
 
                 screen.unblock_input()?;
                 screen.render(None)?;
+                // we do this here in order to recover from a race condition on app start
+                // that sometimes causes Zellij to think the terminal window is a different size
+                // than it actually is - here, we query the client for its terminal size after
+                // we've finished the setup and handle it as we handle a normal resize,
+                // while this can affect other instances of a layout being applied, the query is
+                // very short and cheap and shouldn't cause any trouble
+                if let Some(os_input) = &mut screen.bus.os_input {
+                    for client_id in screen.connected_clients.borrow().iter() {
+                        let _ = os_input
+                            .send_to_client(*client_id, ServerToClientMsg::QueryTerminalSize);
+                    }
+                }
             },
             ScreenInstruction::GoToTab(tab_index, client_id) => {
                 let client_id_to_switch = if client_id.is_none() {
diff --git a/zellij-server/src/unit/snapshots/zellij_server__screen__screen_tests__send_cli_switch_mode_action.snap b/zellij-server/src/unit/snapshots/zellij_server__screen__screen_tests__send_cli_switch_mode_action.snap
index 49e3d8fb1..f79859e99 100644
--- a/zellij-server/src/unit/snapshots/zellij_server__screen__screen_tests__send_cli_switch_mode_action.snap
+++ b/zellij-server/src/unit/snapshots/zellij_server__screen__screen_tests__send_cli_switch_mode_action.snap
@@ -1,6 +1,6 @@
 ---
 source: zellij-server/src/./unit/screen_tests.rs
-assertion_line: 1711
+assertion_line: 2465
 expression: "format!(\"{:?}\", *\n    mock_screen.os_input.server_to_client_messages.lock().unwrap())"
 ---
-{1: [SwitchToMode(Locked)]}
+{1: [QueryTerminalSize, SwitchToMode(Locked)]}
diff --git a/zellij-utils/src/errors.rs b/zellij-utils/src/errors.rs
index c6ef80e68..d05df3a25 100644
--- a/zellij-utils/src/errors.rs
+++ b/zellij-utils/src/errors.rs
@@ -424,6 +424,7 @@ pub enum ClientContext {
     SetSynchronisedOutput,
     UnblockCliPipeInput,
     CliPipeOutput,
+    QueryTerminalSize,
 }
 
 /// Stack call representations corresponding to the different types of [`ServerInstruction`]s.
diff --git a/zellij-utils/src/ipc.rs b/zellij-utils/src/ipc.rs
index b91ff6973..fc0044576 100644
--- a/zellij-utils/src/ipc.rs
+++ b/zellij-utils/src/ipc.rs
@@ -105,6 +105,7 @@ pub enum ServerToClientMsg {
     SwitchSession(ConnectToSession),
     UnblockCliPipeInput(String),   // String -> pipe name
     CliPipeOutput(String, String), // String -> pipe name, String -> Output
+    QueryTerminalSize,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]