From 2a648da557afe4b53ffd26d8464ad0a2a724446e Mon Sep 17 00:00:00 2001
From: Jennifer Player <chainchompa@gmail.com>
Date: Mon, 11 Mar 2024 16:58:51 -0400
Subject: [PATCH] updated model manager to display when import item is
 cancelled

---
 invokeai/app/services/events/events_base.py    |  4 ++--
 .../model_install/model_install_default.py     |  2 +-
 .../listeners/socketio/socketModelInstall.ts   | 18 ++++++++++++++++++
 .../web/src/services/events/actions.ts         |  5 +++++
 .../frontend/web/src/services/events/types.ts  |  7 +++++++
 5 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py
index b52b919edd..98a62fb323 100644
--- a/invokeai/app/services/events/events_base.py
+++ b/invokeai/app/services/events/events_base.py
@@ -405,7 +405,7 @@ class EventServiceBase:
             payload={"source": source, "total_bytes": total_bytes, "key": key, "id": id},
         )
 
-    def emit_model_install_cancelled(self, source: str) -> None:
+    def emit_model_install_cancelled(self, source: str, id: int) -> None:
         """
         Emit when an install job is cancelled.
 
@@ -413,7 +413,7 @@ class EventServiceBase:
         """
         self.__emit_model_event(
             event_name="model_install_cancelled",
-            payload={"source": source},
+            payload={"source": source, "id": id},
         )
 
     def emit_model_install_error(self, source: str, error_type: str, error: str, id: int) -> None:
diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py
index 8c8f74012b..138bde8bbf 100644
--- a/invokeai/app/services/model_install/model_install_default.py
+++ b/invokeai/app/services/model_install/model_install_default.py
@@ -892,7 +892,7 @@ class ModelInstallService(ModelInstallServiceBase):
     def _signal_job_cancelled(self, job: ModelInstallJob) -> None:
         self._logger.info(f"{job.source}: model installation was cancelled")
         if self._event_bus:
-            self._event_bus.emit_model_install_cancelled(str(job.source))
+            self._event_bus.emit_model_install_cancelled(str(job.source), id=job.id)
 
     @staticmethod
     def get_fetcher_from_url(url: str):
diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts
index 59c88bf2a8..33940b9c55 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketModelInstall.ts
@@ -2,6 +2,7 @@ import type { AppStartListening } from 'app/store/middleware/listenerMiddleware'
 import { api } from 'services/api';
 import { modelsApi } from 'services/api/endpoints/models';
 import {
+  socketModelInstallCancelled,
   socketModelInstallCompleted,
   socketModelInstallDownloading,
   socketModelInstallError,
@@ -63,4 +64,21 @@ export const addModelInstallEventListener = (startAppListening: AppStartListenin
       );
     },
   });
+
+  startAppListening({
+    actionCreator: socketModelInstallCancelled,
+    effect: (action, { dispatch }) => {
+      const { id } = action.payload.data;
+
+      dispatch(
+        modelsApi.util.updateQueryData('listModelInstalls', undefined, (draft) => {
+          const modelImport = draft.find((m) => m.id === id);
+          if (modelImport) {
+            modelImport.status = 'cancelled';
+          }
+          return draft;
+        })
+      );
+    },
+  });
 };
diff --git a/invokeai/frontend/web/src/services/events/actions.ts b/invokeai/frontend/web/src/services/events/actions.ts
index 31f9e4d11d..8dd1cfd4fa 100644
--- a/invokeai/frontend/web/src/services/events/actions.ts
+++ b/invokeai/frontend/web/src/services/events/actions.ts
@@ -9,6 +9,7 @@ import type {
   InvocationErrorEvent,
   InvocationRetrievalErrorEvent,
   InvocationStartedEvent,
+  ModelInstallCancelledEvent,
   ModelInstallCompletedEvent,
   ModelInstallDownloadingEvent,
   ModelInstallErrorEvent,
@@ -71,6 +72,10 @@ export const socketModelInstallError = createAction<{
   data: ModelInstallErrorEvent;
 }>('socket/socketModelInstallError');
 
+export const socketModelInstallCancelled = createAction<{
+  data: ModelInstallCancelledEvent;
+}>('socket/socketModelInstallCancelled');
+
 export const socketSessionRetrievalError = createAction<{
   data: SessionRetrievalErrorEvent;
 }>('socket/socketSessionRetrievalError');
diff --git a/invokeai/frontend/web/src/services/events/types.ts b/invokeai/frontend/web/src/services/events/types.ts
index 6fda5a16b1..28e8c5238c 100644
--- a/invokeai/frontend/web/src/services/events/types.ts
+++ b/invokeai/frontend/web/src/services/events/types.ts
@@ -69,6 +69,12 @@ export type ModelInstallErrorEvent = {
   id: number;
 };
 
+export type ModelInstallCancelledEvent = {
+  source: string;
+  timestamp: number;
+  id: number;
+};
+
 /**
  * A `generator_progress` socket.io event.
  *
@@ -264,6 +270,7 @@ export type ServerToClientEvents = {
   model_install_downloading: (payload: ModelInstallDownloadingEvent) => void;
   model_install_completed: (payload: ModelInstallCompletedEvent) => void;
   model_install_error: (payload: ModelInstallErrorEvent) => void;
+  model_install_canceled: (payload: ModelInstallCancelledEvent) => void;
   session_retrieval_error: (payload: SessionRetrievalErrorEvent) => void;
   invocation_retrieval_error: (payload: InvocationRetrievalErrorEvent) => void;
   queue_item_status_changed: (payload: QueueItemStatusChangedEvent) => void;