fix: restore saved machine screens after reconnect (#4126)

refs #3842

1 files changed, 10 insertions, 4 deletions

diff --git a/src/client/shell/endpoints.rs b/src/client/shell/endpoints.rs
index d3823d54..c27358e5 100644
--- a/src/client/shell/endpoints.rs
+++ b/src/client/shell/endpoints.rs
@@ -208,6 +208,7 @@ impl ClientShellState {
         let Some(snapshot) = endpoint.snapshot.clone() else {
             return false;
         };
+        let generation = endpoint.snapshot_generation;
         let switching_endpoint = endpoint_id != &self.active_endpoint_id;
         let agent_scroll = self.agent_scroll;
         if switching_endpoint {
@@ -215,7 +216,7 @@ impl ClientShellState {
             self.pane_surface = None;
             self.pending_pane_surface = None;
         }
-        self.apply_active_snapshot(snapshot);
+        self.apply_active_snapshot(snapshot, generation);
         if switching_endpoint {
             // The aggregate agent list belongs to the client, not one endpoint.
             self.agent_scroll = agent_scroll;
@@ -653,16 +654,21 @@ impl ClientShellState {
     }
 
     fn apply_cached_endpoint_snapshot(&mut self, endpoint_id: &ClientEndpointId) {
-        let Some(snapshot) = self
+        let Some((snapshot, generation)) = self
             .endpoints
             .iter()
             .find(|endpoint| &endpoint.endpoint_id == endpoint_id)
-            .and_then(|endpoint| endpoint.snapshot.clone())
+            .and_then(|endpoint| {
+                endpoint
+                    .snapshot
+                    .clone()
+                    .map(|snapshot| (snapshot, endpoint.snapshot_generation))
+            })
         else {
             return;
         };
         if endpoint_id == &self.active_endpoint_id {
-            self.apply_active_snapshot(snapshot);
+            self.apply_active_snapshot(snapshot, generation);
         }
     }
 }