[backport cloud/1.36] refactor: simplify asset download state and fix deletion UI (#7991)

Backport of #7974 to `cloud/1.36`

Automatically created by backport workflow.

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-7991-backport-cloud-1-36-refactor-simplify-asset-download-state-and-fix-deletion-UI-2e76d73d365081ac8a99c9d4c849325f)
by [Unito](https://www.unito.io)

Co-authored-by: Alexander Brown <drjkl@comfy.org>
Co-authored-by: Amp <amp@ampcode.com>

src/platform/assets/components/AssetBrowserModal.vue

@@ -57,6 +57,7 @@
         :assets="filteredAssets"
         :loading="isLoading"
         @asset-select="handleAssetSelectAndEmit"
+        @asset-deleted="refreshAssets"
       />
     </template>
   </BaseModalLayout>

src/platform/assets/components/AssetCard.vue

@@ -1,6 +1,5 @@
 <template>
   <div
-    v-if="!deletedLocal"
     data-component-id="AssetCard"
     :data-asset-id="asset.id"
     :aria-labelledby="titleId"
@@ -139,8 +138,9 @@ const { asset, interactive } = defineProps<{
   interactive?: boolean
 }>()
 
-defineEmits<{
+const emit = defineEmits<{
   select: [asset: AssetDisplayItem]
+  deleted: [asset: AssetDisplayItem]
 }>()
 
 const { t } = useI18n()
@@ -158,7 +158,6 @@ const descId = useId()
 
 const isEditing = ref(false)
 const newNameRef = ref<string>()
-const deletedLocal = ref(false)
 
 const displayName = computed(() => newNameRef.value ?? asset.name)
 
@@ -211,7 +210,7 @@ function confirmDeletion() {
           })
           // Give a second for the completion message
           await new Promise((resolve) => setTimeout(resolve, 1_000))
-          deletedLocal.value = true
+          emit('deleted', asset)
         } catch (err: unknown) {
           console.error(err)
           promptText.value = t('assetBrowser.deletion.failed', {

src/platform/assets/components/AssetGrid.vue

@@ -35,6 +35,7 @@
           :asset="item"
           :interactive="true"
           @select="$emit('assetSelect', $event)"
+          @deleted="$emit('assetDeleted', $event)"
         />
       </template>
     </VirtualGrid>
@@ -56,6 +57,7 @@ const { assets } = defineProps<{
 
 defineEmits<{
   assetSelect: [asset: AssetDisplayItem]
+  assetDeleted: [asset: AssetDisplayItem]
 }>()
 
 const assetsWithKey = computed(() =>

src/platform/assets/composables/useUploadModelWizard.ts

@@ -245,7 +245,8 @@ export function useUploadModelWizard(modelTypes: Ref<ModelTypeOption[]>) {
           if (selectedModelType.value) {
             assetDownloadStore.trackDownload(
               result.task.task_id,
-              selectedModelType.value
+              selectedModelType.value,
+              filename
             )
           }
           uploadStatus.value = 'processing'

src/stores/assetDownloadStore.test.ts

@@ -117,15 +117,29 @@ describe('useAssetDownloadStore', () => {
     it('associates task with model type for completion tracking', () => {
       const store = useAssetDownloadStore()
 
-      store.trackDownload('task-123', 'checkpoints')
+      store.trackDownload('task-123', 'checkpoints', 'model.safetensors')
       dispatch(createDownloadMessage({ status: 'completed', progress: 100 }))
 
-      expect(store.completedDownloads).toHaveLength(1)
+      expect(store.lastCompletedDownload).toMatchObject({
-      expect(store.completedDownloads[0]).toMatchObject({
         taskId: 'task-123',
         modelType: 'checkpoints'
       })
     })
+
+    it('handles out-of-order messages where completed arrives before progress', () => {
+      const store = useAssetDownloadStore()
+
+      store.trackDownload('task-123', 'checkpoints', 'model.safetensors')
+
+      dispatch(createDownloadMessage({ status: 'completed', progress: 100 }))
+
+      dispatch(createDownloadMessage({ status: 'running', progress: 50 }))
+
+      expect(store.activeDownloads).toHaveLength(0)
+      expect(store.finishedDownloads).toHaveLength(1)
+      expect(store.finishedDownloads[0].status).toBe('completed')
+      expect(store.lastCompletedDownload?.modelType).toBe('checkpoints')
+    })
   })
 
   describe('stale download polling', () => {

src/stores/assetDownloadStore.ts

@@ -16,6 +16,7 @@ export interface AssetDownload {
   lastUpdate: number
   assetId?: string
   error?: string
+  modelType?: string
 }
 
 interface CompletedDownload {
@@ -23,15 +24,29 @@ interface CompletedDownload {
   modelType: string
   timestamp: number
 }
-
-const MAX_COMPLETED_DOWNLOADS = 10
 const STALE_THRESHOLD_MS = 10_000
 const POLL_INTERVAL_MS = 10_000
 
+function generateDownloadTrackingPlaceholder(
+  taskId: string,
+  modelType: string,
+  assetName: string
+): AssetDownload {
+  return {
+    taskId,
+    modelType,
+    assetName,
+    bytesTotal: 0,
+    bytesDownloaded: 0,
+    progress: 0,
+    status: 'created',
+    lastUpdate: Date.now()
+  }
+}
+
 export const useAssetDownloadStore = defineStore('assetDownload', () => {
   const downloads = ref<Map<string, AssetDownload>>(new Map())
-  const pendingModelTypes = new Map<string, string>()
+  const lastCompletedDownload = ref<CompletedDownload | null>(null)
-  const completedDownloads = ref<CompletedDownload[]>([])
 
   const downloadList = computed(() => Array.from(downloads.value.values()))
   const activeDownloads = computed(() =>
@@ -47,8 +62,13 @@ export const useAssetDownloadStore = defineStore('assetDownload', () => {
   const hasActiveDownloads = computed(() => activeDownloads.value.length > 0)
   const hasDownloads = computed(() => downloads.value.size > 0)
 
-  function trackDownload(taskId: string, modelType: string) {
+  function trackDownload(taskId: string, modelType: string, assetName: string) {
-    pendingModelTypes.set(taskId, modelType)
+    if (downloads.value.has(taskId)) return
+
+    downloads.value.set(
+      taskId,
+      generateDownloadTrackingPlaceholder(taskId, modelType, assetName)
+    )
   }
 
   function handleAssetDownload(e: CustomEvent<AssetDownloadWsMessage>) {
@@ -69,24 +89,18 @@ export const useAssetDownloadStore = defineStore('assetDownload', () => {
       progress: data.progress,
       status: data.status,
       error: data.error,
-      lastUpdate: Date.now()
+      lastUpdate: Date.now(),
+      modelType: existing?.modelType
     }
 
     downloads.value.set(data.task_id, download)
 
-    if (data.status === 'completed') {
+    if (data.status === 'completed' && download.modelType) {
-      const modelType = pendingModelTypes.get(data.task_id)
+      lastCompletedDownload.value = {
-      if (modelType) {
+        taskId: data.task_id,
-        const updated = [
+        modelType: download.modelType,
-          ...completedDownloads.value,
+        timestamp: Date.now()
-          { taskId: data.task_id, modelType, timestamp: Date.now() }
-        ]
-        if (updated.length > MAX_COMPLETED_DOWNLOADS) updated.shift()
-        completedDownloads.value = updated
-        pendingModelTypes.delete(data.task_id)
       }
-    } else if (data.status === 'failed') {
-      pendingModelTypes.delete(data.task_id)
     }
   }
 
@@ -157,7 +171,7 @@ export const useAssetDownloadStore = defineStore('assetDownload', () => {
     hasActiveDownloads,
     hasDownloads,
     downloadList,
-    completedDownloads,
+    lastCompletedDownload,
     trackDownload,
     clearFinishedDownloads
   }

src/stores/assetsStore.ts

@@ -1,7 +1,7 @@
-import { useAsyncState } from '@vueuse/core'
+import { useAsyncState, whenever } from '@vueuse/core'
 import { isEqual } from 'es-toolkit'
 import { defineStore } from 'pinia'
-import { computed, shallowReactive, ref, watch } from 'vue'
+import { computed, shallowReactive, ref } from 'vue'
 import {
   mapInputFileToAssetItem,
   mapTaskOutputToAssetItem
@@ -376,24 +376,32 @@ export const useAssetsStore = defineStore('assets', () => {
   } = getModelState()
 
   // Watch for completed downloads and refresh model caches
-  watch(
+  whenever(
-    () => assetDownloadStore.completedDownloads.at(-1),
+    () => assetDownloadStore.lastCompletedDownload,
     async (latestDownload) => {
-      if (!latestDownload) return
-
       const { modelType } = latestDownload
 
       const providers = modelToNodeStore
         .getAllNodeProviders(modelType)
         .filter((provider) => provider.nodeDef?.name)
-      const results = await Promise.allSettled(
+
-        providers.map((provider) =>
+      const nodeTypeUpdates = providers.map((provider) =>
-          updateModelsForNodeType(provider.nodeDef.name).then(
+        updateModelsForNodeType(provider.nodeDef.name).then(
-            () => provider.nodeDef.name
+          () => provider.nodeDef.name
-          )
         )
       )
 
+      // Also update by tag in case modal was opened with assetType
+      const tagUpdates = [
+        updateModelsForTag(modelType),
+        updateModelsForTag('models')
+      ]
+
+      const results = await Promise.allSettled([
+        ...nodeTypeUpdates,
+        ...tagUpdates
+      ])
+
       for (const result of results) {
         if (result.status === 'rejected') {
           console.error(