From 2d6725a41a4f05e65c788ad200881a5dbce996bf Mon Sep 17 00:00:00 2001
From: Conrad Irwin <conrad.irwin@gmail.com>
Date: Wed, 11 Oct 2023 09:45:57 -0600
Subject: [PATCH] Make collaboration warning more useful

---
 crates/collab_ui/src/collab_titlebar_item.rs | 49 +++++++++++++-------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/crates/collab_ui/src/collab_titlebar_item.rs b/crates/collab_ui/src/collab_titlebar_item.rs
index d85aca164a..211ee863e8 100644
--- a/crates/collab_ui/src/collab_titlebar_item.rs
+++ b/crates/collab_ui/src/collab_titlebar_item.rs
@@ -2,6 +2,7 @@ use crate::{
     contact_notification::ContactNotification, face_pile::FacePile, toggle_deafen, toggle_mute,
     toggle_screen_sharing, LeaveCall, ToggleDeafen, ToggleMute, ToggleScreenSharing,
 };
+use auto_update::AutoUpdateStatus;
 use call::{ActiveCall, ParticipantLocation, Room};
 use client::{proto::PeerId, Client, ContactEventKind, SignIn, SignOut, User, UserStore};
 use clock::ReplicaId;
@@ -1177,22 +1178,38 @@ impl CollabTitlebarItem {
                     .with_style(theme.titlebar.offline_icon.container)
                     .into_any(),
             ),
-            client::Status::UpgradeRequired => Some(
-                MouseEventHandler::new::<ConnectionStatusButton, _>(0, cx, |_, _| {
-                    Label::new(
-                        "Please update Zed to collaborate",
-                        theme.titlebar.outdated_warning.text.clone(),
-                    )
-                    .contained()
-                    .with_style(theme.titlebar.outdated_warning.container)
-                    .aligned()
-                })
-                .with_cursor_style(CursorStyle::PointingHand)
-                .on_click(MouseButton::Left, |_, _, cx| {
-                    auto_update::check(&Default::default(), cx);
-                })
-                .into_any(),
-            ),
+            client::Status::UpgradeRequired => {
+                let auto_updater = auto_update::AutoUpdater::get(cx);
+                let label = match auto_updater.map(|auto_update| auto_update.read(cx).status()) {
+                    Some(AutoUpdateStatus::Updated) => "Please restart Zed to Collaborate",
+                    Some(AutoUpdateStatus::Installing)
+                    | Some(AutoUpdateStatus::Downloading)
+                    | Some(AutoUpdateStatus::Checking) => "Updating...",
+                    Some(AutoUpdateStatus::Idle) | Some(AutoUpdateStatus::Errored) | None => {
+                        "Please update Zed to Collaborate"
+                    }
+                };
+
+                Some(
+                    MouseEventHandler::new::<ConnectionStatusButton, _>(0, cx, |_, _| {
+                        Label::new(label, theme.titlebar.outdated_warning.text.clone())
+                            .contained()
+                            .with_style(theme.titlebar.outdated_warning.container)
+                            .aligned()
+                    })
+                    .with_cursor_style(CursorStyle::PointingHand)
+                    .on_click(MouseButton::Left, |_, _, cx| {
+                        if let Some(auto_updater) = auto_update::AutoUpdater::get(cx) {
+                            if auto_updater.read(cx).status() == AutoUpdateStatus::Updated {
+                                workspace::restart(&Default::default(), cx);
+                                return;
+                            }
+                        }
+                        auto_update::check(&Default::default(), cx);
+                    })
+                    .into_any(),
+                )
+            }
             _ => None,
         }
     }