Comfy-Org/ComfyUI_frontend
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI_frontend.git synced 2026-02-23 08:14:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add session download tracking to assetDownloadStore (#8213)

Browse Source 
## Summary

Add session download tracking to track which assets were downloaded
during the current session. This enables UI features like:
- Badge count on "Imported" nav showing newly downloaded assets
- Visual indicator on asset cards for recently downloaded items

## Changes

- Add `acknowledged` flag to `AssetDownload` interface
- Add `unacknowledgedDownloads` computed for filtering
- Add `sessionDownloadCount` computed for badge display
- Add `isDownloadedThisSession(identifier)` to check individual assets
- Add `acknowledgeDownload(identifier)` to mark assets as seen

## Testing

- 6 new unit tests covering all session tracking functionality
- Run: `pnpm test:unit -- src/stores/assetDownloadStore.test.ts`

## Related

- Part of Asset Browser improvements (#8090)

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-8213-feat-add-session-download-tracking-to-assetDownloadStore-2ef6d73d365081538045e8544d26bafa)
by [Unito](https://www.unito.io)

---------

Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Alexander Brown
2026-01-21 16:32:30 -08:00
committed by GitHub
parent d12c6d7814
commit f1d1747582
6 changed files with 272 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/locales/en/main.json
View File

@@ -2315,6 +2315,12 @@
"assetBrowser": {
"allCategory": "All {category}",
"allModels": "All Models",
"byType": "By type",
"emptyImported": {
"canImport": "No imported models yet. Click \"Import Model\" to add your own.",
"restricted": "Personal models are only available at Creator tier and above."
},
"imported": "Imported",
"assetCollection": "Asset collection",
"assets": "Assets",
"baseModels": "Base models",

9
src/platform/assets/components/AssetBrowserModal.vue
View File

@@ -7,12 +7,12 @@
>
<template v-if="shouldShowLeftPanel" #leftPanel>
<LeftSidePanel
v-model="selectedCategory"
v-model="selectedNavItem"
data-component-id="AssetBrowserModal-LeftSidePanel"
:nav-items="availableCategories"
:nav-items
>
<template #header-icon>
<div class="icon-[lucide--folder] size-4" />
<div class="icon-[comfy--ai-model] size-4" />
</template>
<template #header-title>
<span class="capitalize">{{ displayTitle }}</span>
@@ -143,8 +143,9 @@ const { isUploadButtonEnabled, showUploadDialog } =
const {
searchQuery,
selectedNavItem,
selectedCategory,
availableCategories,
navItems,
categoryFilteredAssets,
filteredAssets,
updateFilters

158
src/platform/assets/composables/useAssetBrowser.test.ts
View File

@@ -1,3 +1,4 @@
import { createPinia, setActivePinia } from 'pinia'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { nextTick, ref } from 'vue'
@@ -8,6 +9,8 @@ vi.mock('@/i18n', () => ({
t: (key: string) => {
const translations: Record<string, string> = {
'assetBrowser.allModels': 'All Models',
'assetBrowser.imported': 'Imported',
'assetBrowser.byType': 'By type',
'assetBrowser.assets': 'Assets',
'assetBrowser.unknown': 'unknown'
}
@@ -18,6 +21,7 @@ vi.mock('@/i18n', () => ({
describe('useAssetBrowser', () => {
beforeEach(() => {
setActivePinia(createPinia())
vi.restoreAllMocks()
})
@@ -48,7 +52,7 @@ describe('useAssetBrowser', () => {
tags: ['models', 'loras']
})
const { selectedCategory, categoryFilteredAssets } = useAssetBrowser(
const { selectedNavItem, categoryFilteredAssets } = useAssetBrowser(
ref([checkpointAsset, loraAsset])
)
@@ -56,11 +60,11 @@ describe('useAssetBrowser', () => {
expect(categoryFilteredAssets.value).toHaveLength(2)
// When category selected, should only show that category
selectedCategory.value = 'checkpoints'
selectedNavItem.value = 'checkpoints'
expect(categoryFilteredAssets.value).toHaveLength(1)
expect(categoryFilteredAssets.value[0].id).toBe('checkpoint-1')
selectedCategory.value = 'loras'
selectedNavItem.value = 'loras'
expect(categoryFilteredAssets.value).toHaveLength(1)
expect(categoryFilteredAssets.value[0].id).toBe('lora-1')
})
@@ -150,9 +154,9 @@ describe('useAssetBrowser', () => {
createApiAsset({ id: '3', tags: ['models', 'checkpoints'] })
]
const { selectedCategory, filteredAssets } = useAssetBrowser(ref(assets))
const { selectedNavItem, filteredAssets } = useAssetBrowser(ref(assets))
selectedCategory.value = 'checkpoints'
selectedNavItem.value = 'checkpoints'
await nextTick()
expect(filteredAssets.value).toHaveLength(2)
@@ -169,9 +173,9 @@ describe('useAssetBrowser', () => {
createApiAsset({ id: '2', tags: ['models', 'loras'] })
]
const { selectedCategory, filteredAssets } = useAssetBrowser(ref(assets))
const { selectedNavItem, filteredAssets } = useAssetBrowser(ref(assets))
selectedCategory.value = 'all'
selectedNavItem.value = 'all'
await nextTick()
expect(filteredAssets.value).toHaveLength(2)
@@ -354,12 +358,12 @@ describe('useAssetBrowser', () => {
})
]
const { searchQuery, selectedCategory, filteredAssets } = useAssetBrowser(
const { searchQuery, selectedNavItem, filteredAssets } = useAssetBrowser(
ref(assets)
)
searchQuery.value = 'realistic'
selectedCategory.value = 'checkpoints'
selectedNavItem.value = 'checkpoints'
await nextTick()
expect(filteredAssets.value).toHaveLength(1)
@@ -448,7 +452,7 @@ describe('useAssetBrowser', () => {
expect(filteredAssets.value).toHaveLength(3)
})
it('filters by ownership - my models only', async () => {
it('filters by ownership - imported models only via nav selection', async () => {
const assets = [
createApiAsset({ name: 'my-model.safetensors', is_immutable: false }),
createApiAsset({
@@ -461,14 +465,10 @@ describe('useAssetBrowser', () => {
})
]
const { updateFilters, filteredAssets } = useAssetBrowser(ref(assets))
const { selectedNavItem, filteredAssets } = useAssetBrowser(ref(assets))
updateFilters({
sortBy: 'name-asc',
fileFormats: [],
baseModels: [],
ownership: 'my-models'
})
// Selecting 'imported' nav item filters to my-models (non-immutable)
selectedNavItem.value = 'imported'
await nextTick()
expect(filteredAssets.value).toHaveLength(2)
@@ -477,7 +477,7 @@ describe('useAssetBrowser', () => {
)
})
it('filters by ownership - public models only', async () => {
it('shows all models when nav is "all"', async () => {
const assets = [
createApiAsset({ name: 'my-model.safetensors', is_immutable: false }),
createApiAsset({
@@ -490,41 +490,47 @@ describe('useAssetBrowser', () => {
})
]
const { updateFilters, filteredAssets } = useAssetBrowser(ref(assets))
const { selectedNavItem, filteredAssets } = useAssetBrowser(ref(assets))
updateFilters({
sortBy: 'name-asc',
fileFormats: [],
baseModels: [],
ownership: 'public-models'
})
// Selecting 'all' nav item shows all models
selectedNavItem.value = 'all'
await nextTick()
expect(filteredAssets.value).toHaveLength(2)
expect(filteredAssets.value.every((asset) => asset.is_immutable)).toBe(
true
)
expect(filteredAssets.value).toHaveLength(3)
})
})
describe('Dynamic Category Extraction', () => {
it('extracts categories from asset tags', () => {
it('extracts categories from asset tags into navItems', () => {
const assets = [
createApiAsset({ tags: ['models', 'checkpoints'] }),
createApiAsset({ tags: ['models', 'loras'] }),
createApiAsset({ tags: ['models', 'checkpoints'] }) // duplicate
]
const { availableCategories } = useAssetBrowser(ref(assets))
const { navItems } = useAssetBrowser(ref(assets))
expect(availableCategories.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--folder]' },
// navItems includes quick filters plus a "By type" group
expect(navItems.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--list]' },
{
id: 'checkpoints',
label: 'Checkpoints',
icon: 'icon-[lucide--package]'
id: 'imported',
label: 'Imported',
icon: 'icon-[lucide--folder-input]',
badge: undefined
},
{ id: 'loras', label: 'Loras', icon: 'icon-[lucide--package]' }
{
title: 'By type',
collapsible: false,
items: [
{
id: 'checkpoints',
label: 'Checkpoints',
icon: 'icon-[lucide--folder]'
},
{ id: 'loras', label: 'Loras', icon: 'icon-[lucide--folder]' }
]
}
])
})
@@ -534,11 +540,21 @@ describe('useAssetBrowser', () => {
createApiAsset({ tags: ['models', 'vae'] })
]
const { availableCategories } = useAssetBrowser(ref(assets))
const { navItems } = useAssetBrowser(ref(assets))
expect(availableCategories.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--folder]' },
{ id: 'vae', label: 'Vae', icon: 'icon-[lucide--package]' }
expect(navItems.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--list]' },
{
id: 'imported',
label: 'Imported',
icon: 'icon-[lucide--folder-input]',
badge: undefined
},
{
title: 'By type',
collapsible: false,
items: [{ id: 'vae', label: 'Vae', icon: 'icon-[lucide--folder]' }]
}
])
})
@@ -548,31 +564,47 @@ describe('useAssetBrowser', () => {
createApiAsset({ tags: ['models', 'checkpoints'] })
]
const { availableCategories } = useAssetBrowser(ref(assets))
const { navItems } = useAssetBrowser(ref(assets))
expect(availableCategories.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--folder]' },
expect(navItems.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--list]' },
{
id: 'checkpoints',
label: 'Checkpoints',
icon: 'icon-[lucide--package]'
id: 'imported',
label: 'Imported',
icon: 'icon-[lucide--folder-input]',
badge: undefined
},
{
title: 'By type',
collapsible: false,
items: [
{
id: 'checkpoints',
label: 'Checkpoints',
icon: 'icon-[lucide--folder]'
}
]
}
])
})
it('computes content title from selected category', () => {
it('computes content title from selected nav item', () => {
const assets = [createApiAsset({ tags: ['models', 'checkpoints'] })]
const { selectedCategory, contentTitle } = useAssetBrowser(ref(assets))
const { selectedNavItem, contentTitle } = useAssetBrowser(ref(assets))
// Default
expect(contentTitle.value).toBe('All Models')
// Set specific category
selectedCategory.value = 'checkpoints'
selectedNavItem.value = 'checkpoints'
expect(contentTitle.value).toBe('Checkpoints')
// Set imported
selectedNavItem.value = 'imported'
expect(contentTitle.value).toBe('Imported')
// Unknown category
selectedCategory.value = 'unknown'
selectedNavItem.value = 'unknown'
expect(contentTitle.value).toBe('Assets')
})
@@ -596,26 +628,18 @@ describe('useAssetBrowser', () => {
})
]
const { availableCategories, selectedCategory, categoryFilteredAssets } =
const { navItems, selectedNavItem, categoryFilteredAssets } =
useAssetBrowser(ref(assets))
// Should group all Chatterbox subfolders under single category
expect(availableCategories.value).toEqual([
{ id: 'all', label: 'All Models', icon: 'icon-[lucide--folder]' },
{
id: 'Chatterbox',
label: 'Chatterbox',
icon: 'icon-[lucide--package]'
},
{
id: 'OtherFolder',
label: 'OtherFolder',
icon: 'icon-[lucide--package]'
}
// Should group all Chatterbox subfolders under single category in the type group
const typeGroup = navItems.value[2] as { items: { id: string }[] }
expect(typeGroup.items.map((i) => i.id)).toEqual([
'Chatterbox',
'OtherFolder'
])
// When selecting Chatterbox category, should include all models from its subfolders
selectedCategory.value = 'Chatterbox'
selectedNavItem.value = 'Chatterbox'
expect(categoryFilteredAssets.value).toHaveLength(3)
expect(categoryFilteredAssets.value.map((a) => a.id)).toEqual([
'asset-1',
@@ -624,7 +648,7 @@ describe('useAssetBrowser', () => {
])
// When selecting OtherFolder category, should include only its models
selectedCategory.value = 'OtherFolder'
selectedNavItem.value = 'OtherFolder'
expect(categoryFilteredAssets.value).toHaveLength(1)
expect(categoryFilteredAssets.value[0].id).toBe('asset-4')
})

83
src/platform/assets/composables/useAssetBrowser.ts
View File

@@ -2,6 +2,7 @@ import { computed, ref } from 'vue'
import type { Ref } from 'vue'
import { useFuse } from '@vueuse/integrations/useFuse'
import type { UseFuseOptions } from '@vueuse/integrations/useFuse'
import { storeToRefs } from 'pinia'
import { d, t } from '@/i18n'
import type { FilterState } from '@/platform/assets/components/AssetFilterBar.vue'
@@ -10,9 +11,13 @@ import {
getAssetBaseModel,
getAssetDescription
} from '@/platform/assets/utils/assetMetadataUtils'
import { useAssetDownloadStore } from '@/stores/assetDownloadStore'
import type { NavGroupData, NavItemData } from '@/types/navTypes'
export type OwnershipOption = 'all' | 'my-models' | 'public-models'
type NavId = 'all' | 'imported' | (string & {})
function filterByCategory(category: string) {
return (asset: AssetItem) => {
if (category === 'all') return true
@@ -81,9 +86,12 @@ export function useAssetBrowser(
assetsSource: Ref<AssetItem[] | undefined> = ref<AssetItem[] | undefined>([])
) {
const assets = computed<AssetItem[]>(() => assetsSource.value ?? [])
const assetDownloadStore = useAssetDownloadStore()
const { sessionDownloadCount } = storeToRefs(assetDownloadStore)
// State
const searchQuery = ref('')
const selectedCategory = ref('all')
const selectedNavItem = ref<NavId>('all')
const filters = ref<FilterState>({
sortBy: 'recent',
fileFormats: [],
@@ -91,6 +99,21 @@ export function useAssetBrowser(
ownership: 'all'
})
const selectedOwnership = computed<OwnershipOption>(() => {
if (selectedNavItem.value === 'imported') return 'my-models'
return 'all'
})
const selectedCategory = computed(() => {
if (
selectedNavItem.value === 'all' ||
selectedNavItem.value === 'imported'
) {
return 'all'
}
return selectedNavItem.value
})
// Transform API asset to display asset
function transformAssetForDisplay(asset: AssetItem): AssetDisplayItem {
// Extract description from metadata or create from tags
@@ -136,39 +159,69 @@ export function useAssetBrowser(
}
}
const availableCategories = computed(() => {
const typeCategories = computed<NavItemData[]>(() => {
const categories = assets.value
.filter((asset) => asset.tags[0] === 'models')
.map((asset) => asset.tags[1])
.filter((tag): tag is string => typeof tag === 'string' && tag.length > 0)
.map((tag) => tag.split('/')[0]) // Extract top-level folder name
.map((tag) => tag.split('/')[0])
const uniqueCategories = Array.from(new Set(categories))
return Array.from(new Set(categories))
.sort()
.map((category) => ({
id: category,
label: category.charAt(0).toUpperCase() + category.slice(1),
icon: 'icon-[lucide--package]'
icon: 'icon-[lucide--folder]'
}))
})
return [
const navItems = computed<(NavItemData | NavGroupData)[]>(() => {
const quickFilters: NavItemData[] = [
{
id: 'all',
label: t('assetBrowser.allModels'),
icon: 'icon-[lucide--folder]'
icon: 'icon-[lucide--list]'
},
...uniqueCategories
{
id: 'imported',
label: t('assetBrowser.imported'),
icon: 'icon-[lucide--folder-input]',
badge:
sessionDownloadCount.value > 0
? sessionDownloadCount.value
: undefined
}
]
if (typeCategories.value.length === 0) {
return quickFilters
}
return [
...quickFilters,
{
title: t('assetBrowser.byType'),
items: typeCategories.value,
collapsible: false
}
]
})
// Compute content title from selected category
const isImportedSelected = computed(
() => selectedNavItem.value === 'imported'
)
// Compute content title from selected nav item
const contentTitle = computed(() => {
if (selectedCategory.value === 'all') {
if (selectedNavItem.value === 'all') {
return t('assetBrowser.allModels')
}
if (selectedNavItem.value === 'imported') {
return t('assetBrowser.imported')
}
const category = availableCategories.value.find(
(cat) => cat.id === selectedCategory.value
const category = typeCategories.value.find(
(cat) => cat.id === selectedNavItem.value
)
return category?.label || t('assetBrowser.assets')
})
@@ -205,7 +258,7 @@ export function useAssetBrowser(
const filtered = searchFiltered.value
.filter(filterByFileFormats(filters.value.fileFormats))
.filter(filterByBaseModels(filters.value.baseModels))
.filter(filterByOwnership(filters.value.ownership))
.filter(filterByOwnership(selectedOwnership.value))
const sortedAssets = [...filtered]
sortedAssets.sort((a, b) => {
@@ -234,11 +287,13 @@ export function useAssetBrowser(
return {
searchQuery,
selectedNavItem,
selectedCategory,
availableCategories,
navItems,
contentTitle,
categoryFilteredAssets,
filteredAssets,
isImportedSelected,
updateFilters
}
}

76
src/stores/assetDownloadStore.test.ts
View File

@@ -236,4 +236,80 @@ describe('useAssetDownloadStore', () => {
expect(store.finishedDownloads).toHaveLength(0)
})
})
describe('session download tracking', () => {
it('counts unacknowledged completed downloads with asset IDs', () => {
const store = useAssetDownloadStore()
dispatch(
createDownloadMessage({
status: 'completed',
progress: 100,
asset_id: 'asset-456'
})
)
expect(store.sessionDownloadCount).toBe(1)
})
it('does not count completed downloads without asset IDs', () => {
const store = useAssetDownloadStore()
dispatch(
createDownloadMessage({
status: 'completed',
progress: 100,
asset_id: undefined
})
)
expect(store.sessionDownloadCount).toBe(0)
})
it('does not count failed downloads', () => {
const store = useAssetDownloadStore()
dispatch(
createDownloadMessage({
status: 'failed',
asset_id: 'asset-456'
})
)
expect(store.sessionDownloadCount).toBe(0)
})
it('isDownloadedThisSession returns true for unacknowledged downloads', () => {
const store = useAssetDownloadStore()
dispatch(
createDownloadMessage({
status: 'completed',
progress: 100,
asset_id: 'asset-456'
})
)
expect(store.isDownloadedThisSession('asset-456')).toBe(true)
expect(store.isDownloadedThisSession('other-asset')).toBe(false)
})
it('acknowledgeAsset decrements session count', () => {
const store = useAssetDownloadStore()
dispatch(
createDownloadMessage({
status: 'completed',
progress: 100,
asset_id: 'asset-456'
})
)
expect(store.sessionDownloadCount).toBe(1)
store.acknowledgeAsset('asset-456')
expect(store.sessionDownloadCount).toBe(0)
expect(store.isDownloadedThisSession('asset-456')).toBe(false)
})
})
})

26
src/stores/assetDownloadStore.ts
View File

@@ -17,6 +17,7 @@ export interface AssetDownload {
assetId?: string
error?: string
modelType?: string
acknowledged?: boolean
}
interface CompletedDownload {
@@ -59,9 +60,29 @@ export const useAssetDownloadStore = defineStore('assetDownload', () => {
(d) => d.status === 'completed' || d.status === 'failed'
)
)
const unacknowledgedDownloads = computed(() =>
finishedDownloads.value.filter(
(d) => d.status === 'completed' && d.assetId && !d.acknowledged
)
)
const sessionDownloadCount = computed(
() => unacknowledgedDownloads.value.length
)
const hasActiveDownloads = computed(() => activeDownloads.value.length > 0)
const hasDownloads = computed(() => downloads.value.size > 0)
function isDownloadedThisSession(assetId: string): boolean {
return unacknowledgedDownloads.value.some((d) => d.assetId === assetId)
}
function acknowledgeAsset(assetId: string) {
for (const download of downloads.value.values()) {
if (download.assetId === assetId) {
download.acknowledged = true
}
}
}
function trackDownload(taskId: string, modelType: string, assetName: string) {
if (downloads.value.has(taskId)) return
@@ -172,7 +193,10 @@ export const useAssetDownloadStore = defineStore('assetDownload', () => {
hasDownloads,
downloadList,
lastCompletedDownload,
sessionDownloadCount,
trackDownload,
clearFinishedDownloads
clearFinishedDownloads,
isDownloadedThisSession,
acknowledgeAsset
}
})