diff --git a/src/data/service.ts b/src/data/service.ts
index f93312ecf2..70b05a2f14 100644
--- a/src/data/service.ts
+++ b/src/data/service.ts
@@ -10,5 +10,15 @@ export const callExecuteScript = (
     sequence,
   });
 
-export const serviceCallWillDisconnect = (domain: string, service: string) =>
-  domain === "homeassistant" && ["restart", "stop"].includes(service);
+export const serviceCallWillDisconnect = (
+  domain: string,
+  service: string,
+  serviceData?: Record<string, any>
+) =>
+  (domain === "homeassistant" && ["restart", "stop"].includes(service)) ||
+  (domain === "update" &&
+    service === "install" &&
+    [
+      "update.home_assistant_core_update",
+      "update.home_assistant_operating_system_update",
+    ].includes(serviceData?.entity_id));
diff --git a/src/state/connection-mixin.ts b/src/state/connection-mixin.ts
index c2c81cc87f..8ba5091a02 100644
--- a/src/state/connection-mixin.ts
+++ b/src/state/connection-mixin.ts
@@ -100,7 +100,7 @@ export const connectionMixin = <T extends Constructor<HassBaseEl>>(
           } catch (err: any) {
             if (
               err.error?.code === ERR_CONNECTION_LOST &&
-              serviceCallWillDisconnect(domain, service)
+              serviceCallWillDisconnect(domain, service, serviceData)
             ) {
               return { context: { id: "" } };
             }
@@ -121,7 +121,12 @@ export const connectionMixin = <T extends Constructor<HassBaseEl>>(
                 "ui.notification_toast.service_call_failed",
                 "service",
                 `${domain}/${service}`
-              ) + ` ${err.message}`;
+              ) +
+              ` ${
+                err.message || err.error?.code === ERR_CONNECTION_LOST
+                  ? "connection lost"
+                  : "unknown error"
+              }`;
             fireEvent(this as any, "hass-notification", { message });
             throw err;
           }