mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-01-26 19:09:52 +00:00
refactor: simplify asset download state and fix deletion UI (#7974)
## 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 <amp@ampcode.com>
This commit is contained in:
@@ -57,6 +57,7 @@
|
||||
:assets="filteredAssets"
|
||||
:loading="isLoading"
|
||||
@asset-select="handleAssetSelectAndEmit"
|
||||
@asset-deleted="refreshAssets"
|
||||
/>
|
||||
</template>
|
||||
</BaseModalLayout>
|
||||
|
||||
@@ -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', {
|
||||
|
||||
@@ -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(() =>
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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', () => {
|
||||
|
||||
@@ -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 completedDownloads = ref<CompletedDownload[]>([])
|
||||
const lastCompletedDownload = ref<CompletedDownload | null>(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<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') {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user