From a166ec91a6e95b6dc4d2d139b929b1d24df5f608 Mon Sep 17 00:00:00 2001 From: Alexander Brown Date: Mon, 12 Jan 2026 17:57:25 -0800 Subject: [PATCH] refactor: simplify asset download state and fix deletion UI (#7974) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Refactors asset download state management and fixes asset deletion UI issues. ## Changes ### assetDownloadStore simplification - Replace `pendingModelTypes` Map with `modelType` stored directly on `AssetDownload` - Replace `completedDownloads` array with single `lastCompletedDownload` ref - `trackDownload()` now creates a placeholder entry immediately - Use VueUse `whenever` instead of `watch` for cleaner null handling ### Asset refresh on download completion - Refresh all relevant caches when a download completes: - Node type caches (e.g., "CheckpointLoaderSimple") - Tag caches (e.g., "tag:checkpoints") - "All Models" cache ("tag:models") ### Asset deletion fix - Remove local `deletedLocal` state that caused blank grid cells - Emit `deleted` event from AssetCard → AssetGrid → AssetBrowserModal - Trigger store refresh on deletion to properly remove the asset from the grid ## Testing - Added test for out-of-order websocket message handling - All existing tests pass ┆Issue is synchronized with this [Notion page](https://www.notion.so/PR-7974-refactor-simplify-asset-download-state-and-fix-deletion-UI-2e76d73d365081c69bcde9150a0d460c) by [Unito](https://www.unito.io) --------- Co-authored-by: Amp --- .../assets/components/AssetBrowserModal.vue | 1 + src/platform/assets/components/AssetCard.vue | 7 ++- src/platform/assets/components/AssetGrid.vue | 2 + .../composables/useUploadModelWizard.ts | 3 +- src/stores/assetDownloadStore.test.ts | 20 +++++-- src/stores/assetDownloadStore.ts | 54 ++++++++++++------- src/stores/assetsStore.ts | 30 +++++++---- 7 files changed, 78 insertions(+), 39 deletions(-) diff --git a/src/platform/assets/components/AssetBrowserModal.vue b/src/platform/assets/components/AssetBrowserModal.vue index 001387d05..b7db7494c 100644 --- a/src/platform/assets/components/AssetBrowserModal.vue +++ b/src/platform/assets/components/AssetBrowserModal.vue @@ -57,6 +57,7 @@ :assets="filteredAssets" :loading="isLoading" @asset-select="handleAssetSelectAndEmit" + @asset-deleted="refreshAssets" /> diff --git a/src/platform/assets/components/AssetCard.vue b/src/platform/assets/components/AssetCard.vue index a24610022..11e404d41 100644 --- a/src/platform/assets/components/AssetCard.vue +++ b/src/platform/assets/components/AssetCard.vue @@ -1,6 +1,5 @@ @@ -56,6 +57,7 @@ const { assets } = defineProps<{ defineEmits<{ assetSelect: [asset: AssetDisplayItem] + assetDeleted: [asset: AssetDisplayItem] }>() const assetsWithKey = computed(() => diff --git a/src/platform/assets/composables/useUploadModelWizard.ts b/src/platform/assets/composables/useUploadModelWizard.ts index 207fed562..9a945bad9 100644 --- a/src/platform/assets/composables/useUploadModelWizard.ts +++ b/src/platform/assets/composables/useUploadModelWizard.ts @@ -245,7 +245,8 @@ export function useUploadModelWizard(modelTypes: Ref) { if (selectedModelType.value) { assetDownloadStore.trackDownload( result.task.task_id, - selectedModelType.value + selectedModelType.value, + filename ) } uploadStatus.value = 'processing' diff --git a/src/stores/assetDownloadStore.test.ts b/src/stores/assetDownloadStore.test.ts index 28d58b9c1..5882c4b9a 100644 --- a/src/stores/assetDownloadStore.test.ts +++ b/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.completedDownloads[0]).toMatchObject({ + expect(store.lastCompletedDownload).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', () => { diff --git a/src/stores/assetDownloadStore.ts b/src/stores/assetDownloadStore.ts index 8aa5788bd..88818c31c 100644 --- a/src/stores/assetDownloadStore.ts +++ b/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>(new Map()) - const pendingModelTypes = new Map() - const completedDownloads = ref([]) + const lastCompletedDownload = ref(null) 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) { - pendingModelTypes.set(taskId, modelType) + function trackDownload(taskId: string, modelType: string, assetName: string) { + if (downloads.value.has(taskId)) return + + downloads.value.set( + taskId, + generateDownloadTrackingPlaceholder(taskId, modelType, assetName) + ) } function handleAssetDownload(e: CustomEvent) { @@ -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') { - const modelType = pendingModelTypes.get(data.task_id) - if (modelType) { - const updated = [ - ...completedDownloads.value, - { taskId: data.task_id, modelType, timestamp: Date.now() } - ] - if (updated.length > MAX_COMPLETED_DOWNLOADS) updated.shift() - completedDownloads.value = updated - pendingModelTypes.delete(data.task_id) + if (data.status === 'completed' && download.modelType) { + lastCompletedDownload.value = { + taskId: data.task_id, + modelType: download.modelType, + timestamp: Date.now() } - } 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 } diff --git a/src/stores/assetsStore.ts b/src/stores/assetsStore.ts index 5fe64e911..56efe3daa 100644 --- a/src/stores/assetsStore.ts +++ b/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( - () => assetDownloadStore.completedDownloads.at(-1), + whenever( + () => 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) => - updateModelsForNodeType(provider.nodeDef.name).then( - () => provider.nodeDef.name - ) + + const nodeTypeUpdates = providers.map((provider) => + updateModelsForNodeType(provider.nodeDef.name).then( + () => 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(