refactor: Simplify AssetBrowserModal and useModelUpload

Author: DrJKL

src/platform/assets/components/AssetBrowserModal.test.ts

@@ -4,22 +4,23 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import AssetBrowserModal from '@/platform/assets/components/AssetBrowserModal.vue'
 import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
-
-const mockAssetsStore = vi.hoisted(() => ({
-  modelAssetsByNodeType: new Map<string, unknown[]>(),
-  modelLoadingByNodeType: new Map<string, boolean>(),
-  updateModelsForNodeType: vi.fn()
-}))
+import { useAssetsStore } from '@/stores/assetsStore'
 
 vi.mock('@/i18n', () => ({
   t: (key: string, params?: Record<string, string>) =>
     params ? `${key}:${JSON.stringify(params)}` : key,
   d: (date: Date) => date.toLocaleDateString()
 }))
 
-vi.mock('@/stores/assetsStore', () => ({
-  useAssetsStore: () => mockAssetsStore
-}))
+vi.mock('@/stores/assetsStore', () => {
+  const store = {
+    modelAssetsByNodeType: new Map<string, unknown[]>(),
+    modelLoadingByNodeType: new Map<string, boolean>(),
+    updateModelsForNodeType: vi.fn(),
+    updateModelsForTag: vi.fn()
+  }
+  return { useAssetsStore: () => store }
+})
 
 vi.mock('@/stores/modelToNodeStore', () => ({
   useModelToNodeStore: () => ({
@@ -182,10 +183,12 @@ describe('AssetBrowserModal', () => {
     })
   }
 
+  const mockStore = useAssetsStore()
+
   beforeEach(() => {
-    mockAssetsStore.modelAssetsByNodeType.clear()
-    mockAssetsStore.modelLoadingByNodeType.clear()
-    mockAssetsStore.updateModelsForNodeType.mockReset()
+    mockStore.modelAssetsByNodeType.clear()
+    mockStore.modelLoadingByNodeType.clear()
+    vi.clearAllMocks()
   })
 
   describe('Integration with useAssetBrowser', () => {
@@ -194,10 +197,7 @@ describe('AssetBrowserModal', () => {
         createTestAsset('asset1', 'Model A', 'checkpoints'),
         createTestAsset('asset2', 'Model B', 'loras')
       ]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
       await flushPromises()
@@ -214,10 +214,7 @@ describe('AssetBrowserModal', () => {
         createTestAsset('c1', 'model.safetensors', 'checkpoints'),
         createTestAsset('l1', 'lora.pt', 'loras')
       ]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({
         nodeType: 'CheckpointLoaderSimple',
@@ -237,17 +234,14 @@ describe('AssetBrowserModal', () => {
       createWrapper({ nodeType: 'CheckpointLoaderSimple' })
       await flushPromises()
 
-      expect(mockAssetsStore.updateModelsForNodeType).toHaveBeenCalledWith(
+      expect(mockStore.updateModelsForNodeType).toHaveBeenCalledWith(
         'CheckpointLoaderSimple'
       )
     })
 
     it('displays cached assets immediately from store', async () => {
       const assets = [createTestAsset('asset1', 'Cached Model', 'checkpoints')]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
 
@@ -257,15 +251,33 @@ describe('AssetBrowserModal', () => {
       expect(gridAssets).toHaveLength(1)
       expect(gridAssets[0].name).toBe('Cached Model')
     })
+
+    it('triggers store refresh for asset type (tag) on mount', async () => {
+      createWrapper({ assetType: 'models' })
+      await flushPromises()
+
+      expect(mockStore.updateModelsForTag).toHaveBeenCalledWith('models')
+    })
+
+    it('uses tag: prefix for cache key when assetType is provided', async () => {
+      const assets = [createTestAsset('asset1', 'Tagged Model', 'models')]
+      mockStore.modelAssetsByNodeType.set('tag:models', assets)
+
+      const wrapper = createWrapper({ assetType: 'models' })
+      await flushPromises()
+
+      const assetGrid = wrapper.findComponent({ name: 'AssetGrid' })
+      const gridAssets = assetGrid.props('assets') as AssetItem[]
+
+      expect(gridAssets).toHaveLength(1)
+      expect(gridAssets[0].name).toBe('Tagged Model')
+    })
   })
 
   describe('Asset Selection', () => {
     it('emits asset-select event when asset is selected', async () => {
       const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
       await flushPromises()
@@ -278,10 +290,7 @@ describe('AssetBrowserModal', () => {
 
     it('executes onSelect callback when provided', async () => {
       const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const onSelect = vi.fn()
       const wrapper = createWrapper({
@@ -324,10 +333,7 @@ describe('AssetBrowserModal', () => {
         createTestAsset('asset1', 'Model A', 'checkpoints'),
         createTestAsset('asset2', 'Model B', 'loras')
       ]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({
         nodeType: 'CheckpointLoaderSimple',
@@ -360,10 +366,7 @@ describe('AssetBrowserModal', () => {
 
     it('passes computed contentTitle to BaseModalLayout when no title prop', async () => {
       const assets = [createTestAsset('asset1', 'Model A', 'checkpoints')]
-      mockAssetsStore.modelAssetsByNodeType.set(
-        'CheckpointLoaderSimple',
-        assets
-      )
+      mockStore.modelAssetsByNodeType.set('CheckpointLoaderSimple', assets)
 
       const wrapper = createWrapper({ nodeType: 'CheckpointLoaderSimple' })
       await flushPromises()

src/platform/assets/components/AssetBrowserModal.vue

@@ -64,7 +64,7 @@
 
 <script setup lang="ts">
 import { breakpointsTailwind, useBreakpoints } from '@vueuse/core'
-import { computed, provide, ref, watchEffect } from 'vue'
+import { computed, provide } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import SearchBox from '@/components/common/SearchBox.vue'
@@ -124,22 +124,6 @@ const isLoading = computed(
   () => isStoreLoading.value && fetchedAssets.value.length === 0
 )
 
-// Track if we've triggered a refresh for this key
-const refreshedKey = ref<string | null>(null)
-
-// Trigger refresh when nodeType/assetType changes (runs on mount and prop change)
-watchEffect(async () => {
-  const key = cacheKey.value
-  if (key && key !== refreshedKey.value) {
-    refreshedKey.value = key
-    if (props.nodeType) {
-      await assetStore.updateModelsForNodeType(props.nodeType)
-    } else if (props.assetType) {
-      await assetStore.updateModelsForTag(props.assetType)
-    }
-  }
-})
-
 async function refreshAssets(): Promise<AssetItem[]> {
   if (props.nodeType) {
     return await assetStore.updateModelsForNodeType(props.nodeType)
@@ -150,6 +134,9 @@ async function refreshAssets(): Promise<AssetItem[]> {
   return []
 }
 
+// Trigger background refresh on mount
+refreshAssets()
+
 const { isUploadButtonEnabled, showUploadDialog } =
   useModelUpload(refreshAssets)
 

src/platform/assets/composables/useModelUpload.ts

@@ -3,13 +3,11 @@ import UploadModelDialog from '@/platform/assets/components/UploadModelDialog.vu
 import UploadModelDialogHeader from '@/platform/assets/components/UploadModelDialogHeader.vue'
 import UploadModelUpgradeModal from '@/platform/assets/components/UploadModelUpgradeModal.vue'
 import UploadModelUpgradeModalHeader from '@/platform/assets/components/UploadModelUpgradeModalHeader.vue'
-import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
 import { useDialogStore } from '@/stores/dialogStore'
-import type { UseAsyncStateReturn } from '@vueuse/core'
 import { computed } from 'vue'
 
 export function useModelUpload(
-  execute?: UseAsyncStateReturn<AssetItem[], [], true>['execute']
+  onUploadSuccess?: () => Promise<unknown> | void
 ) {
   const dialogStore = useDialogStore()
   const { flags } = useFeatureFlags()
@@ -37,7 +35,7 @@ export function useModelUpload(
         component: UploadModelDialog,
         props: {
           onUploadSuccess: async () => {
-            await execute?.()
+            await onUploadSuccess?.()
           }
         },
         dialogComponentProps: {