Compare commits

..

49 Commits

Author SHA1 Message Date
Simon Pinfold
84491c3c60 Merge branch 'synap5e/disable-change-model-tag-on-core' into synap5e/assets-mode-sidebar 2026-07-12 23:03:10 +12:00
Simon Pinfold
e8f5617da7 Merge branch 'synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag' into synap5e/disable-change-model-tag-on-core 2026-07-12 23:03:10 +12:00
Simon Pinfold
65f80ecee7 Merge branch 'synap5e/fe-1076-fe-group-the-model-library-by-model_type-tags' into synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag 2026-07-12 23:03:10 +12:00
Simon Pinfold
f6bbaf4b9e docs(assets): correct the multi-root limitation to preview-only reads
/view_metadata takes no path index and searches roots in order, so
metadata only misresolves on relative-path collisions; the root-0
targeting applies to preview URLs and row identity.
2026-07-12 23:02:39 +12:00
Simon Pinfold
1e7d48623a fix(assets): represent every category of a shared multi-type asset
A shared-root asset tagged with several model_type values groups under
each of them, but its badge and a node-typed picker's title keyed off
only the first value, so viewing it under its second category showed the
first category's label. Cards now badge every authoritative model_type
value, and a node-typed picker titles off the requested node's category
instead of guessing from the first asset.
2026-07-12 23:01:25 +12:00
Simon Pinfold
c8afd4cf8a docs(assets): mark the multi-root pathIndex limitation at its source
Asset records carry no root identity, so the asset-backed sidebar
reports every model as root 0; secondary-root metadata/preview reads
misresolve on multi-root categories. Accepted for this surface, loader
values are unaffected; lifting it needs the backend to carry root
identity on assets.
2026-07-12 22:52:02 +12:00
Simon Pinfold
abfb89b990 Merge branch 'synap5e/disable-change-model-tag-on-core' into synap5e/assets-mode-sidebar 2026-07-12 21:43:14 +12:00
Simon Pinfold
259fd9a62f Merge branch 'synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag' into synap5e/disable-change-model-tag-on-core 2026-07-12 21:43:13 +12:00
Simon Pinfold
724e31d235 Merge branch 'synap5e/fe-1076-fe-group-the-model-library-by-model_type-tags' into synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag 2026-07-12 21:43:13 +12:00
Simon Pinfold
9a7793ccc8 fix(models): close reload races that strand or misorder folder contents
A scan-complete reload landing while the eager load is still in flight
replaced the Loading folder objects but only re-loaded Loaded ones,
stranding the sidebar on uninitialized folders whose original requests
finish into detached objects; Loading folders now count as previously
loaded. Same-folder concurrent refreshes were last-completion-wins, so
an older slow response could overwrite a newer commit; refreshes now
carry a per-folder generation and stale ones discard.
2026-07-12 21:40:05 +12:00
Simon Pinfold
317cc81196 fix(assets): close pagination-termination and title-preservation gaps
Adversarial review follow-ups: an offset-ignoring backend that reorders
its responses could alternate two page orderings and defeat the
repeated-page check forever, so the walk now also remembers every served
page signature and treats a run of no-new-ID full pages as exhausted.
Modal titles for uncovered assets keep their legacy verbatim tag (only
covered assets retitle to their model_type value), restoring the
pre-alignment behavior the batch accidentally changed, and namespace
residue like a malformed empty model_type: can no longer leak into
titles, nav items, or categories.
2026-07-12 21:34:25 +12:00
Simon Pinfold
64740d9d4a Merge branch 'synap5e/disable-change-model-tag-on-core' into synap5e/assets-mode-sidebar 2026-07-12 21:05:25 +12:00
Simon Pinfold
9788ed2439 Merge branch 'synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag' into synap5e/disable-change-model-tag-on-core 2026-07-12 21:05:00 +12:00
Simon Pinfold
7b841548ce refactor(assets): require modelTypeMode on the write-path helpers
Matches the read-path helpers: new call sites must choose the tag mode
explicitly instead of silently inheriting legacy behavior.
2026-07-12 21:04:01 +12:00
Simon Pinfold
fcb7d838ef Merge branch 'synap5e/fe-1076-fe-group-the-model-library-by-model_type-tags' into synap5e/fe-1084-fe-write-model-type-from-discovery-flip-the-cutover-flag
# Conflicts:
#	src/platform/assets/utils/assetMetadataUtils.test.ts
#	src/platform/assets/utils/assetMetadataUtils.ts
2026-07-12 21:01:01 +12:00
Simon Pinfold
873a85e59e fix(assets): surface refresh errors and harden pagination races
A refresh that fails before its first batch commits now mirrors the
error onto the committed state consumers actually read, instead of
dying silently on the discarded pending state. Pagination stops only
when a full page repeats the previous page verbatim (ignored-offset
signature); an all-duplicate page whose content merely shifted keeps
walking so later pages still load. Teardown of a stale request can no
longer evict a newer request's single-flight entries after an
interleaved invalidation.
2026-07-12 20:57:50 +12:00
Simon Pinfold
ca33752569 fix(assets): align badge and loader resolution with model_type tags
In modelTypeMode the card badge and modal title now key off the same
model_type: value the asset groups under, so they can no longer diverge
from the grouping (review consensus on FE-1076). Loader-node resolution
demotes bare tags unrelated to any model_type: value to trailing
fallbacks: they can no longer pre-empt an authoritative resolvable
model_type: however deep they are, while related hierarchical twins
(LLM/Qwen-VL over flat model_type:LLM) keep winning.

Also: getAssetModelType can no longer return a raw model_type: literal
(kills the prefixed round-trip through edit widgets), modelTypeMode is
now a required parameter so new call sites must choose explicitly,
malformed empty model_type: tags produce no badge, per-asset flag reads
hoisted, and a test pins the local reserved-tag mirrors to
assetService's canonical values.
2026-07-12 20:57:13 +12:00
Simon Pinfold
e6916bb665 chore: correct an overclaiming comment on the model-library routing 2026-07-10 16:34:25 +12:00
Simon Pinfold
fdc4651934 fix(assets): close eager-load and folder-refresh races in the model store
Second-pass review: loadModels now loads the folder structure first when it
has not arrived yet (eager-loading could race app boot and silently load
nothing), and refreshModelFolder no longer commits a stale folder object
over a structure rebuilt by a concurrent reload. The store also adopts the
assetService-owned scan subscription (with scope-disposal cleanup) and
requestModelScan handles its own failure via async try/catch.
2026-07-10 16:33:29 +12:00
Simon Pinfold
62123c4c0d refactor(assets): own the scan-complete subscription in assetService
Second-pass review: the wire-level assets.seed.fast_complete event is now
owned by assetService, which exposes an onModelsScanned callback returning
an unsubscribe, and the bucket memoizer uses async/try-finally over a
then/finally chain.
2026-07-10 16:32:22 +12:00
Simon Pinfold
1cb04bef92 chore(assets): drop the scan-discovery console log
Review feedback: the sidebar visibly updating is signal enough, and this
removes the only no-console suppression in src.
2026-07-10 14:45:47 +12:00
Simon Pinfold
f5e221b955 chore(assets): drop a comment restating behavior pinned by a test 2026-07-10 14:42:56 +12:00
Simon Pinfold
fd1f2726a5 fix(settings): open the asset browser only when the asset API is enabled
Review feedback: rather than prompt-correcting the invalid combination, the
browser routing now derives from UseAssetBrowser AND UseAssetAPI — the asset
browser cannot function without the asset API, so with the API disabled the
model library tab always opens the sidebar tree.
2026-07-10 14:42:04 +12:00
Simon Pinfold
8d51d933fc refactor(assets): resolve model preview urls in the model store
Review feedback: url construction is domain logic that belongs below the
component layer; ModelTreeLeaf now consumes a pure getModelPreviewUrl.
2026-07-10 14:41:47 +12:00
Simon Pinfold
be564b232b Merge remote-tracking branch 'origin/main' into synap5e/fe-1076-fe-group-the-model-library-by-model_type-tags
# Conflicts:
#	src/composables/useFeatureFlags.ts
2026-07-10 14:07:58 +12:00
Simon Pinfold
b62405a6ca fix(assets): group slashed bare tags by their top-level segment
Bare-tag backends can carry hierarchical tags (e.g. Chatterbox/sub/model);
the asset browser has always grouped those under the top-level segment, and
the sidebar's folder names come from /experiment/models, so bucketing under
the full slashed tag made such models unreachable. Match the browser's
grouping. On main's server-side per-folder query these assets were equally
invisible, so this is a fix over both.
2026-07-07 11:12:02 +12:00
Simon Pinfold
33c1806673 fix(assets): drop assets with traversal-shaped loader paths
The loader value flows into viewMetadata URLs and widget values; validate
it against the existing filename schema at bucket time so a '../' path
cannot pass through even if backend validation regresses.
2026-07-07 10:51:42 +12:00
Simon Pinfold
9342caad0a fix(assets): harden concurrent folder reloads
Guard loadModelFolders with a request id so a slow stale response cannot
overwrite a fresher folder structure when a manual refresh races the
scan-complete reload, and log instead of leaking an unhandled rejection
when the post-scan background reload fails.
2026-07-07 10:51:20 +12:00
Simon Pinfold
48aed9e0d9 feat(assets): filter model library folders by their registered extensions
/experiment/models returns each folder's raw registered allowlist; an
empty array means match-all. The asset walk filters folder contents by the
registered set verbatim, substituting an FE-owned copy of core's default
supported_pt_extensions when the allowlist is empty or the field is absent
(older backends). This reproduces the legacy sidebar's global-set display
behavior for match-all folders and extensions-less backends — nothing that
was hidden starts showing — while registered allowlists let exotic folders
(e.g. .gguf) surface their actual models. The legacy listing keeps its
historical server-side filter untouched.
2026-07-07 09:14:02 +12:00
Simon Pinfold
dd3dceeaca fix(assets): keep active search results live across model reloads
Search results were a snapshot filled inside handleSearch, so a reload
(e.g. a completed scan discovering new files) updated the store without the
filtered tree ever re-running. Filter via a computed over the store instead,
and expand results from a watch so late-arriving matches unfold too.
2026-07-07 09:14:02 +12:00
Simon Pinfold
3605fc1d75 feat(assets): eager-load the model library in asset mode
In asset mode every folder resolves from one cached walk, so load all
folders on sidebar mount: search and badge counts are complete from the
start and the load-all button becomes redundant (hidden; it remains for the
request-per-folder legacy path, where AutoLoadAll stays the opt-in). Log
when a scan's fast phase discovered new files, since the sidebar now
updates itself and the console note is the only trace of why.
2026-07-07 09:14:02 +12:00
Simon Pinfold
f97195e392 feat(assets): rescan model roots on refresh and live-update the sidebar
Manual refresh now kicks off POST /assets/seed (fire-and-forget, 409
tolerated, skipped on Cloud and on the legacy listing path) alongside the
immediate re-load of known server state. The scan's fast phase already
persists the category tags and loader filenames the sidebar needs, and the
backend announces it with the assets.seed.fast_complete websocket event, so
the store re-loads previously loaded folders on that event instead of
polling /assets/seed/status (which cannot distinguish the fast phase from
enrichment).
2026-07-07 09:14:02 +12:00
Simon Pinfold
0488d0d8a4 feat(assets): back the model library sidebar with the asset API
Derive folders from model_type: tags and use loader_path for the tree,
widget value, and viewMetadata. Gate local-disk metadata/preview reads
on !isCloud.

The models tag is walked once and bucketed per category (multi-category
assets land in each), memoized behind a request id so an in-flight walk
cannot repopulate an invalidated cache. Assets with a null loader_path
are unloadable (orphans) and are dropped with a warning rather than
minting a widget value from the deprecated name field; bare-tag backends
that predate loader_path keep the name fallback.
2026-07-07 09:14:02 +12:00
Simon Pinfold
aea28a06de feat(assets): gate model library asset view behind a runtime setting (FE-1142)
Add a hidden, experimental Comfy.Assets.ModelLibraryAssetView setting that
controls whether the model library opens the asset browser or the sidebar,
defaulting to the asset browser on Cloud and the sidebar on Core. The
model-library toggle now reads this setting instead of Comfy.Assets.UseAssetAPI,
decoupling the sidebar-vs-asset-browser view from Asset API enablement so the
asset-backed view stays reachable on Core for testing without forcing the
cutover.
2026-07-07 09:14:02 +12:00
Simon Pinfold
5847b3d148 fix(dialog): keep non-modal dialogs open when focus lands on a toast
Toasts teleport to body but were missing from the reka outside-layer
allowlist, so a non-modal dialog (e.g. Settings) dismissed itself when
focus recovery landed on a toast's close button. This happens whenever an
action button disables itself while its request is in flight and the
request fails: the browser drops focus from the disabled button and it can
resurface on the error toast. Latent on main; deterministic once the
sidebar rework shifted chunk timing enough for the loading state to render
before the failure response.
2026-07-07 09:14:02 +12:00
Simon Pinfold
187239712c fix(assets): make model type read-only on core
Retagging a model rewrites its asset tags, but core is filesystem-backed
and does not yet move the file to match the new type, so an edit would
leave the tag and the on-disk location diverged. Gate the model type
dropdown to cloud builds; on core the field renders the existing
read-only label. Revert once core implements move-on-retag (or the asset
system no longer requires the file move).
2026-07-02 14:08:09 +12:00
Simon Pinfold
e53b7b4ba8 fix(assets): recognize model_type: prefix in import type-mismatch guard
The required-type guard and imported-type label compared the bare folder_name
against raw asset tags, so a namespaced model_type:<folder> tag was treated as a
type mismatch (and leaked the raw prefix into the mismatch label). Strip the
model_type: prefix before comparing/labelling so a namespaced tag satisfies the
required type. Not reachable in production (the dual-write keeps the bare twin),
but hardens the import path against a namespaced-only asset.
2026-07-01 14:20:37 +12:00
Simon Pinfold
ab56cf9e82 test(assets): pin import write keeps user_metadata.model_type bare
Assert the import-by-URL path namespaces only the tag (model_type:<folder_name>)
while user_metadata.model_type stays the bare folder_name, since that value
still feeds download tracking, category refresh, and the progress label.
2026-07-01 14:20:37 +12:00
Simon Pinfold
f772fc0123 feat(assets): write model_type: tag from discovery
Prefix the discovery folder_name with model_type: on model import and on
model-type edits, gated on the supports_model_type_tags backend capability.
Default off writes the legacy bare directory-path tag unchanged, so this is a
no-op until the backend declares the cutover; the flag is the same authoritative
signal FE-1076 grouping reads.

The FE writes only the single selected folder_name's model_type: form; the
backend keeps the bare directory-path twin in sync (and forward-expands shared
paths), so the edit path manages just the model_type: tags and the resolver read
path is untouched. Mirrors the read side: an uncovered asset with no model_type:
tag still routes by its bare tag.
2026-07-01 14:20:37 +12:00
Simon Pinfold
a5f4559df8 fix(assets): halt model batch loop when backend ignores offset
The batch loop only stopped on a short page, so a backend that ignores
offset and repeats a full page would refetch forever. Now runs on non-cloud
builds too. Stop when a full page contributes no new IDs.
2026-07-01 14:07:54 +12:00
Simon Pinfold
6103dd164e fix(assets): prefer resolvable model_type over incidental slash tags
An incidental hierarchical tag (e.g. a user-added foo/bar) could outrank
an authoritative model_type:vae in node-category resolution, targeting the
wrong provider or failing with NO_PROVIDER. Resolution now tries category
candidates most-specific-first and takes the first that maps to a provider,
so the deepest resolvable path still wins but unresolvable tags fall through.
2026-07-01 14:07:37 +12:00
Simon Pinfold
f7d672e3a5 test(assets): restore mocks after each resolveModelNodeFromAsset test
Add afterEach(vi.restoreAllMocks) so the console.error spy doesn't persist
across tests (CodeRabbit nit).
2026-07-01 09:35:05 +12:00
Simon Pinfold
6f0eaefe1b fix(assets): keep reserved-tag constants local in assetMetadataUtils
Importing MODELS_TAG/MISSING_TAG from assetService made this leaf util pull
the assetService -> i18n (createI18n) chain. Test suites that import
assetMetadataUtils (e.g. useAssetFilterOptions, AssetsSidebarListView) with a
minimal vue-i18n mock then failed to load. Define the two reserved-tag
literals locally so the util has no heavy transitive dependency.
2026-07-01 09:35:05 +12:00
Simon Pinfold
55ceec0a16 fix(assets): port model_type node-category resolution to resolveModelNodeFromAsset
main (#12765) renamed createModelNodeFromAsset -> resolveModelNodeFromAsset.
Re-apply FE-1076's flag-gated node-category resolution there: use
getAssetNodeCategory(asset, flags.supportsModelTypeTags) instead of the inline
first-non-reserved-tag lookup, so a model_type:* tag resolves the loader when
the supports_model_type_tags flag is on.
2026-07-01 09:35:05 +12:00
Simon Pinfold
5ac1b15266 fix(assets): strip model_type: prefix from the asset browser title
Also realigns the card type badge to the same rule: keep the existing
"first non-models tag" selection and only strip the model_type: namespace
prefix, rather than reordering to prefer the model_type: value. So an asset
like [models, foo, model_type:bar] still shows foo, while
[models, model_type:checkpoints] shows checkpoints. Title (primaryCategoryTag)
strips the same prefix when the flag is on.
2026-07-01 09:35:05 +12:00
Simon Pinfold
df826415ca fix(assets): strip model_type: prefix from asset card type badge
The card type badge derived from the first raw non-models tag, so in
model_type mode it rendered the literal `model_type:checkpoints`. Route it
through a model_type-aware getAssetTypeBadge and key the display cache on the
flag so a late-arriving feature flag re-derives the badge. Also adds
docstrings to the tag helpers to satisfy coverage.
2026-07-01 09:35:05 +12:00
Simon Pinfold
ec25e874a2 refactor(assets): run model-asset store on non-cloud builds
Unwrap the if (isCloud) gate in assetsStore getModelState so the asset-mode
Model Library works on non-cloud distributions (it previously returned an
empty no-op store off-cloud, leaving the library blank even with the asset
API enabled). Scoped to the model-asset store only; the rest of the isCloud
forks are out of scope here.
2026-07-01 09:35:05 +12:00
Simon Pinfold
c8e4029f96 refactor(assets): gate model_type grouping on supports_model_type_tags
Drive the model_type: grouping/resolution off the backend-reported
supports_model_type_tags capability (getServerFeature) instead of sniffing
tags per asset. Default false preserves the legacy bare-tag grouping, so a
backend that does not declare the capability is unchanged. When the flag is
true, assets group by their model_type: values and an uncovered asset (no
model_type: tag) still routes by its bare tags.

The thread-through is a modelTypeMode flag on getAssetCategories /
getAssetNodeCategory / filterByCategory (keeping them pure); call sites read
the capability via useFeatureFlags.
2026-07-01 09:35:05 +12:00
Simon Pinfold
5ebeb580ca feat(assets): group model library by model_type tags
Group the asset-mode Model Library by namespaced model_type:* tags. When an
asset carries model_type:* tags, the stripped values are the authoritative
categories and all other tags are disregarded; otherwise fall back to the
legacy bare-tag top-level grouping. A shared getAssetCategories helper keeps
the "By type" nav and category filtering in agreement.

Resolve node-provider lookup from the same tags via getAssetNodeCategory:
prefer the most specific (deepest parent/child) candidate so a flat
model_type:LLM does not shadow a resolvable hierarchical tag, and strip the
model_type: prefix before the modelNodeMappings lookup.
2026-07-01 09:34:29 +12:00
34 changed files with 2917 additions and 640 deletions

View File

@@ -418,28 +418,26 @@ export class AssetsSidebarTab extends SidebarTab {
async openSettingsMenu() {
await this.dismissToasts()
await this.settingsButton.click()
// Wait for the popover content to render. Use the default timeout so slower
// (e.g. cloud) app inits don't burst-fail a tight explicit timeout.
await expect(
this.listViewOption.or(this.gridViewOption).first()
).toBeVisible()
// Wait for popover content to render
await this.listViewOption
.or(this.gridViewOption)
.first()
.waitFor({ state: 'visible', timeout: 3000 })
}
async openFilterMenu() {
await this.dismissToasts()
await this.filterButton.click()
// Wait for the filter popover to open. Use the default timeout so slower
// (e.g. cloud) app inits don't burst-fail a tight explicit timeout.
await expect(this.filterCheckbox('Image')).toBeVisible()
await this.filterCheckbox('Image').waitFor({
state: 'visible',
timeout: 3000
})
}
async toggleMediaTypeFilter(
filter: MediaFilterKind | MediaFilterLabel
): Promise<void> {
const checkbox = this.filterCheckbox(filter)
// Ensure the popover has finished opening before reading its state; a stale
// read here races the reka-ui slide/fade animation.
await expect(checkbox).toBeVisible()
const before = await checkbox.getAttribute('aria-checked')
await checkbox.click()
const expected = before === 'true' ? 'false' : 'true'

View File

@@ -422,17 +422,9 @@ test.describe('Vue Node Moving', { tag: '@vue-nodes' }, () => {
loadCheckpointHeaderPos
)
// Poll the header position so the assertion retries until the touch pan
// has settled, instead of reading a single mid-animation bounding box.
// A screen bounding box read is pixel-quantized, so assert to the nearest
// pixel (precision 0 == within 0.5px) rather than the default precision 2
// (within 0.005px), which sub-pixel canvas rounding cannot satisfy.
await expect
.poll(() => getHeaderPos(comfyPage, 'Load Checkpoint').then((p) => p.x))
.toBeCloseTo(loadCheckpointHeaderPos.x + 64, 0)
await expect
.poll(() => getHeaderPos(comfyPage, 'Load Checkpoint').then((p) => p.y))
.toBeCloseTo(loadCheckpointHeaderPos.y + 64, 0)
const newHeaderPos = await getLoadCheckpointHeaderPos(comfyPage)
expect(newHeaderPos.x).toBeCloseTo(loadCheckpointHeaderPos.x + 64)
expect(newHeaderPos.y).toBeCloseTo(loadCheckpointHeaderPos.y + 64)
}
)
})

View File

@@ -426,7 +426,7 @@ describe('shouldPreventRekaDismiss', () => {
expect(event.defaultPrevented).toBe(true)
})
it.for(['p-dialog', 'p-select-overlay'])(
it.for(['p-dialog', 'p-select-overlay', 'p-toast'])(
'focus-outside on a sibling %s portal does not dismiss the parent',
(className) => {
const overlay = document.createElement('div')

View File

@@ -1,10 +1,13 @@
// PrimeVue overlays (Select, ColorPicker, Popover, Autocomplete, stacked
// PrimeVue Dialogs) teleport to body. Reka treats clicks on body-portaled
// elements as outside its dialog and would auto-dismiss on the first
// interaction, tearing the overlay down mid-interaction. Treat any
// PrimeVue overlay click as inside.
// PrimeVue Dialogs, Toasts) teleport to body. Reka treats clicks on
// body-portaled elements as outside its dialog and would auto-dismiss on the
// first interaction, tearing the overlay down mid-interaction. Treat any
// PrimeVue overlay click as inside. Toasts matter for focus-outside: when a
// button disables itself mid-action (e.g. a confirm entering its loading
// state), the browser drops focus and recovery can land on the toast's close
// button, which must not dismiss the dialog underneath.
const PRIMEVUE_OVERLAY_SELECTORS =
'.p-select-overlay, .p-colorpicker-panel, .p-popover, .p-autocomplete-overlay, .p-overlay, .p-overlay-mask, .p-dialog'
'.p-select-overlay, .p-colorpicker-panel, .p-popover, .p-autocomplete-overlay, .p-overlay, .p-overlay-mask, .p-dialog, .p-toast'
// Reka portals its own dialogs / popovers / menus into the body too. When a
// nested Reka layer opens on top of a non-modal parent, the parent's

View File

@@ -1,5 +1,6 @@
import { createTestingPinia } from '@pinia/testing'
import { fromPartial } from '@total-typescript/shoehorn'
import userEvent from '@testing-library/user-event'
import { render, screen } from '@testing-library/vue'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { nextTick } from 'vue'
@@ -18,7 +19,10 @@ const {
mockGetNodeProvider,
mockToggleNodeOnEvent,
mockRefreshModelFolder,
downloadStoreState
mockLoadModels,
downloadStoreState,
settingState,
modelsState
} = vi.hoisted(() => {
let capturedRoot: TreeExplorerNode<unknown> | null = null
return {
@@ -33,7 +37,13 @@ const {
mockGetNodeProvider: vi.fn(),
mockToggleNodeOnEvent: vi.fn(),
mockRefreshModelFolder: vi.fn().mockResolvedValue(undefined),
downloadStoreState: { setLastCompleted: (_: unknown) => {} }
mockLoadModels: vi.fn().mockResolvedValue([]),
downloadStoreState: { setLastCompleted: (_: unknown) => {} },
settingState: { useAssetAPI: false, autoLoadAll: false },
modelsState: {
push: (_: unknown) => {},
reset: () => {}
}
}
})
@@ -54,20 +64,30 @@ const mockModel = fromPartial<ComfyModelDef>({
searchable: 'checkpoints/model.safetensors'
})
vi.mock('@/stores/modelStore', () => ({
ResourceState: {
Loading: 'loading',
Loaded: 'loaded'
},
useModelStore: () => ({
modelFolders: [],
models: [mockModel],
loadModels: vi.fn().mockResolvedValue([]),
loadModelFolders: vi.fn().mockResolvedValue([]),
refresh: vi.fn().mockResolvedValue(undefined),
refreshModelFolder: mockRefreshModelFolder
})
}))
vi.mock('@/stores/modelStore', async () => {
const { reactive } = await import('vue')
const models = reactive<ComfyModelDef[]>([])
modelsState.push = (model: unknown) => {
models.push(model as ComfyModelDef)
}
modelsState.reset = () => {
models.splice(0, models.length, mockModel)
}
return {
ResourceState: {
Loading: 'loading',
Loaded: 'loaded'
},
useModelStore: () => ({
modelFolders: [],
models,
loadModels: mockLoadModels,
loadModelFolders: vi.fn().mockResolvedValue([]),
refresh: vi.fn().mockResolvedValue(undefined),
refreshModelFolder: mockRefreshModelFolder
})
}
})
vi.mock('@/stores/assetDownloadStore', async () => {
const { ref } = await import('vue')
@@ -92,6 +112,10 @@ vi.mock('@/platform/settings/settingStore', () => ({
useSettingStore: () => ({
get: vi.fn((key: string) => {
if (key === 'Comfy.ModelLibrary.NameFormat') return 'filename'
if (key === 'Comfy.Assets.UseAssetAPI') return settingState.useAssetAPI
if (key === 'Comfy.ModelLibrary.AutoLoadAll') {
return settingState.autoLoadAll
}
return false
})
})
@@ -104,26 +128,45 @@ vi.mock('@/composables/useTreeExpansion', () => ({
})
}))
vi.mock('@/components/common/TreeExplorer.vue', () => ({
default: {
name: 'TreeExplorer',
template: '<div data-testid="tree-explorer" />',
props: ['root', 'expandedKeys'],
setup(props: { root: TreeExplorerNode<unknown> }) {
captureRoot(props.root)
vi.mock('@/components/common/TreeExplorer.vue', async () => {
const { watchEffect } = await import('vue')
return {
default: {
name: 'TreeExplorer',
template: '<div data-testid="tree-explorer" />',
props: ['root', 'expandedKeys'],
setup(props: { root: TreeExplorerNode<unknown> }) {
watchEffect(() => captureRoot(props.root))
}
}
}
}))
})
vi.mock('@/components/ui/search-input/SearchInput.vue', () => ({
default: {
name: 'SearchInput',
template: '<input data-testid="search-input" />',
template: '<input data-testid="search-input" @input="onInput" />',
props: ['modelValue', 'placeholder'],
setup() {
return { focus: vi.fn() }
},
expose: ['focus']
emits: ['update:modelValue', 'search'],
setup(
_props: unknown,
{
emit,
expose
}: {
emit: (event: 'update:modelValue' | 'search', value: string) => void
expose: (exposed: Record<string, unknown>) => void
}
) {
expose({ focus: vi.fn() })
return {
onInput: (event: Event) => {
const value = (event.target as HTMLInputElement).value
emit('update:modelValue', value)
emit('search', value)
}
}
}
}
}))
@@ -134,7 +177,8 @@ vi.mock('./SidebarTopArea.vue', () => ({
vi.mock('./SidebarTabTemplate.vue', () => ({
default: {
name: 'SidebarTabTemplate',
template: '<div><slot name="header" /><slot name="body" /></div>'
template:
'<div><slot name="tool-buttons" /><slot name="header" /><slot name="body" /></div>'
}
}))
@@ -157,13 +201,17 @@ describe('ModelLibrarySidebarTab', () => {
vi.clearAllMocks()
resetRoot()
downloadStoreState.setLastCompleted(null)
settingState.useAssetAPI = false
settingState.autoLoadAll = false
modelsState.reset()
})
function renderComponent() {
return render(ModelLibrarySidebarTab, {
global: {
plugins: [createTestingPinia({ stubActions: false }), i18n],
stubs: { teleport: true }
stubs: { teleport: true },
directives: { tooltip: {} }
}
})
}
@@ -236,4 +284,67 @@ describe('ModelLibrarySidebarTab', () => {
expect(mockRefreshModelFolder).not.toHaveBeenCalled()
})
describe('search', () => {
it('updates active search results when a reload adds a matching model', async () => {
const user = userEvent.setup()
renderComponent()
await nextTick()
await user.type(screen.getByTestId('search-input'), 'model')
await nextTick()
expect(mockLoadModels).toHaveBeenCalled()
const leafLabels = () => {
const { children: folders = [] } = getRoot()
return folders.flatMap(({ children: leaves = [] }) =>
leaves.map((leaf) => leaf.label)
)
}
expect(leafLabels()).toEqual(['model'])
// A completed scan reloads the store while the search is still active.
modelsState.push(
fromPartial<ComfyModelDef>({
key: 'checkpoints/model-new.safetensors',
file_name: 'model-new.safetensors',
simplified_file_name: 'model-new',
title: 'Model New',
directory: 'checkpoints',
searchable: 'checkpoints/model-new.safetensors'
})
)
await nextTick()
expect(leafLabels()).toEqual(['model', 'model-new'])
})
})
describe('asset mode', () => {
it('hides the load-all button and eager-loads models on mount', async () => {
settingState.useAssetAPI = true
renderComponent()
await nextTick()
expect(screen.queryByLabelText('g.loadAllFolders')).toBeNull()
expect(screen.getByLabelText('g.refresh')).toBeInTheDocument()
expect(mockLoadModels).toHaveBeenCalledTimes(1)
})
it('legacy mode keeps the load-all button and stays lazy by default', async () => {
renderComponent()
await nextTick()
expect(screen.getByLabelText('g.loadAllFolders')).toBeInTheDocument()
expect(mockLoadModels).not.toHaveBeenCalled()
})
it('legacy mode still honors AutoLoadAll', async () => {
settingState.autoLoadAll = true
renderComponent()
await nextTick()
expect(mockLoadModels).toHaveBeenCalledTimes(1)
})
})
})

View File

@@ -11,6 +11,7 @@
<i class="icon-[lucide--refresh-cw] size-4" />
</Button>
<Button
v-if="!usesAssetAPI"
v-tooltip.bottom="$t('g.loadAllFolders')"
variant="muted-textonly"
size="icon"
@@ -77,28 +78,28 @@ import { buildTree } from '@/utils/treeUtil'
const modelStore = useModelStore()
const modelToNodeStore = useModelToNodeStore()
const settingStore = useSettingStore()
const usesAssetAPI = computed(() =>
settingStore.get('Comfy.Assets.UseAssetAPI')
)
const assetDownloadStore = useAssetDownloadStore()
const searchBoxRef = ref()
const searchQuery = ref<string>('')
const expandedKeys = ref<Record<string, boolean>>({})
const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
const filteredModels = ref<ComfyModelDef[]>([])
const filteredModels = computed<ComfyModelDef[]>(() => {
const search = searchQuery.value.toLocaleLowerCase()
if (!search) return []
return modelStore.models.filter((model) => model.searchable.includes(search))
})
const handleSearch = async (query: string) => {
if (!query) {
filteredModels.value = []
expandedKeys.value = {}
return
}
// Load all models to ensure we have the latest data
// Load all models to ensure results cover folders not yet opened
await modelStore.loadModels()
const search = query.toLocaleLowerCase()
filteredModels.value = modelStore.models.filter((model: ComfyModelDef) => {
return model.searchable.includes(search)
})
await nextTick()
expandNode(root.value)
}
type ModelOrFolder = ComfyModelDef | ModelFolder
@@ -112,6 +113,12 @@ const root = computed<TreeNode>(() => {
)
})
watch(root, async (newRoot) => {
if (!searchQuery.value) return
await nextTick()
expandNode(newRoot)
})
const renderedRoot = computed<TreeExplorerNode<ModelOrFolder>>(() => {
const nameFormat = settingStore.get('Comfy.ModelLibrary.NameFormat')
const fillNodeInfo = (node: TreeNode): TreeExplorerNode<ModelOrFolder> => {
@@ -193,7 +200,13 @@ watch(
onMounted(async () => {
searchBoxRef.value?.focus()
if (settingStore.get('Comfy.ModelLibrary.AutoLoadAll')) {
// In asset mode the whole library resolves from one cached walk, so eager
// loading is cheap and keeps search and folder badges complete from the
// start; AutoLoadAll remains the opt-in for the request-per-folder legacy path.
if (
usesAssetAPI.value ||
settingStore.get('Comfy.ModelLibrary.AutoLoadAll')
) {
await modelStore.loadModels()
}
})

View File

@@ -26,6 +26,7 @@ import { computed, nextTick, onMounted, onUnmounted, ref } from 'vue'
import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
import { useSettingStore } from '@/platform/settings/settingStore'
import type { ComfyModelDef } from '@/stores/modelStore'
import { getModelPreviewUrl } from '@/stores/modelStore'
import type { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
import ModelPreview from './ModelPreview.vue'
@@ -37,17 +38,7 @@ const props = defineProps<{
// Note: The leaf node should always have a model definition on node.data.
const modelDef = computed<ComfyModelDef>(() => props.node.data!)
const modelPreviewUrl = computed(() => {
if (modelDef.value.image) {
return modelDef.value.image
}
const folder = modelDef.value.directory
const path_index = modelDef.value.path_index
const extension = modelDef.value.file_name.split('.').pop()
const filename = modelDef.value.file_name.replace(`.${extension}`, '.webp')
const encodedFilename = encodeURIComponent(filename).replace(/%2F/g, '/')
return `/api/experiment/models/preview/${folder}/${path_index}/${encodedFilename}`
})
const modelPreviewUrl = computed(() => getModelPreviewUrl(modelDef.value))
const previewRef = ref<InstanceType<typeof ModelPreview> | null>(null)
const modelPreviewStyle = ref<CSSProperties>({

View File

@@ -33,7 +33,8 @@ export enum ServerFeatureFlag {
SHOW_SIGNIN_BUTTON = 'show_signin_button',
UNIFIED_CLOUD_AUTH = 'unified_cloud_auth',
CONSOLIDATED_BILLING_ENABLED = 'consolidated_billing_enabled',
SIGNUP_TURNSTILE = 'signup_turnstile'
SIGNUP_TURNSTILE = 'signup_turnstile',
SUPPORTS_MODEL_TYPE_TAGS = 'supports_model_type_tags'
}
/**
@@ -208,6 +209,12 @@ export function useFeatureFlags() {
remoteConfig.value.signup_turnstile,
'off'
)
},
get supportsModelTypeTags() {
return api.getServerFeature(
ServerFeatureFlag.SUPPORTS_MODEL_TYPE_TAGS,
false
)
}
})

View File

@@ -9,6 +9,20 @@ import { useAssetsStore } from '@/stores/assetsStore'
const mockAssetsByKey = vi.hoisted(() => new Map<string, AssetItem[]>())
const mockLoadingByKey = vi.hoisted(() => new Map<string, boolean>())
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
vi.mock('@/composables/useFeatureFlags', () => ({
useFeatureFlags: () => ({
flags: {
get supportsModelTypeTags() {
return mockSupportsModelTypeTags.value
},
get modelUploadButtonEnabled() {
return false
}
}
})
}))
vi.mock('@/i18n', () => ({
t: (key: string, params?: Record<string, string>) =>
@@ -214,6 +228,7 @@ describe('AssetBrowserModal', () => {
vi.resetAllMocks()
mockAssetsByKey.clear()
mockLoadingByKey.clear()
mockSupportsModelTypeTags.value = false
})
describe('Integration with useAssetBrowser', () => {
@@ -420,5 +435,20 @@ describe('AssetBrowserModal', () => {
'assetBrowser.allCategory:{"category":"Checkpoints"}'
)
})
it('strips the model_type: prefix from the title when the flag is on', async () => {
mockSupportsModelTypeTags.value = true
const assets = [
createTestAsset('asset1', 'Model A', 'model_type:checkpoints')
]
mockAssetsByKey.set('CheckpointLoaderSimple', assets)
renderModal({ nodeType: 'CheckpointLoaderSimple' })
await flushPromises()
expect(screen.getByTestId('modal-title').textContent).toBe(
'assetBrowser.allCategory:{"category":"Checkpoints"}'
)
})
})
})

View File

@@ -100,6 +100,7 @@ import SearchInput from '@/components/ui/search-input/SearchInput.vue'
import Button from '@/components/ui/button/Button.vue'
import BaseModalLayout from '@/components/widget/layout/BaseModalLayout.vue'
import LeftSidePanel from '@/components/widget/panel/LeftSidePanel.vue'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { usePrimeVueOverlayChildStyle } from '@/composables/usePopoverSizing'
import AssetFilterBar from '@/platform/assets/components/AssetFilterBar.vue'
import AssetGrid from '@/platform/assets/components/AssetGrid.vue'
@@ -109,12 +110,14 @@ import { useAssetBrowser } from '@/platform/assets/composables/useAssetBrowser'
import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
import { useModelUpload } from '@/platform/assets/composables/useModelUpload'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { getPrimaryCategoryTag } from '@/platform/assets/utils/assetMetadataUtils'
import { formatCategoryLabel } from '@/platform/assets/utils/categoryLabel'
import { useAssetsStore } from '@/stores/assetsStore'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
import { OnCloseKey } from '@/types/widgetTypes'
const { t } = useI18n()
const { flags } = useFeatureFlags()
const assetStore = useAssetsStore()
const modelToNodeStore = useModelToNodeStore()
const breakpoints = useBreakpoints(breakpointsTailwind)
@@ -191,9 +194,21 @@ const focusedAsset = ref<AssetDisplayItem | null>(null)
const isRightPanelOpen = ref(false)
const primaryCategoryTag = computed(() => {
const modelTypeMode = flags.supportsModelTypeTags
// A node-typed picker is FOR a category; title off that category rather
// than guessing from the first asset, whose first model_type value may be
// a different category it shares a root with.
if (modelTypeMode && props.nodeType) {
const mapped = modelToNodeStore.getCategoryForNodeType(props.nodeType)
if (mapped) return mapped
}
const assets = fetchedAssets.value ?? []
// Covered assets title off the model_type value they group under (so title
// and grouping cannot diverge); uncovered assets keep the legacy verbatim
// first tag.
const tagFromAssets = assets
.map((asset) => asset.tags?.find((tag) => tag !== 'models'))
.map((asset) => getPrimaryCategoryTag(asset, modelTypeMode))
.find((tag): tag is string => typeof tag === 'string' && tag.length > 0)
if (tagFromAssets) return tagFromAssets

View File

@@ -1,5 +1,5 @@
import { createTestingPinia } from '@pinia/testing'
import { describe, expect, it, vi } from 'vitest'
import { afterEach, describe, expect, it, vi } from 'vitest'
import { createI18n } from 'vue-i18n'
import { render, screen } from '@testing-library/vue'
@@ -14,6 +14,14 @@ vi.mock('@/composables/useCopyToClipboard', () => ({
})
}))
const mockDistribution = vi.hoisted(() => ({ isCloud: false }))
vi.mock('@/platform/distribution/types', async (importOriginal) => ({
...(await importOriginal<object>()),
get isCloud() {
return mockDistribution.isCloud
}
}))
const i18n = createI18n({
legacy: false,
locale: 'en',
@@ -41,6 +49,10 @@ describe('ModelInfoPanel', () => {
...overrides
})
afterEach(() => {
mockDistribution.isCloud = false
})
function renderPanel(asset: AssetDisplayItem) {
return render(ModelInfoPanel, {
props: { asset },
@@ -138,6 +150,18 @@ describe('ModelInfoPanel', () => {
).toBeInTheDocument()
})
it('shows an editable model type dropdown for a mutable asset on cloud', () => {
mockDistribution.isCloud = true
renderPanel(createMockAsset({ is_immutable: false }))
expect(screen.getByRole('combobox')).toBeInTheDocument()
})
it('keeps the model type read-only on core even for a mutable asset', () => {
mockDistribution.isCloud = false
renderPanel(createMockAsset({ is_immutable: false }))
expect(screen.queryByRole('combobox')).not.toBeInTheDocument()
})
it('renders base models field', () => {
const asset = createMockAsset({
user_metadata: { base_model: ['SDXL'] }

View File

@@ -71,7 +71,7 @@
</span>
</template>
<ModelInfoField :label="t('assetBrowser.modelInfo.modelType')">
<Select v-if="!isImmutable" v-model="selectedModelType">
<Select v-if="isModelTypeEditable" v-model="selectedModelType">
<SelectTrigger class="w-full">
<SelectValue
:placeholder="t('assetBrowser.modelInfo.selectModelType')"
@@ -215,6 +215,7 @@ import { useI18n } from 'vue-i18n'
import EditableText from '@/components/common/EditableText.vue'
import { useCopyToClipboard } from '@/composables/useCopyToClipboard'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import PropertiesAccordionItem from '@/components/rightSidePanel/layout/PropertiesAccordionItem.vue'
import Button from '@/components/ui/button/Button.vue'
import Select from '@/components/ui/select/Select.vue'
@@ -229,17 +230,19 @@ import TagsInputItemDelete from '@/components/ui/tags-input/TagsInputItemDelete.
import TagsInputItemText from '@/components/ui/tags-input/TagsInputItemText.vue'
import type { AssetDisplayItem } from '@/platform/assets/composables/useAssetBrowser'
import { useModelTypes } from '@/platform/assets/composables/useModelTypes'
import { isCloud } from '@/platform/distribution/types'
import type { AssetUserMetadata } from '@/platform/assets/schemas/assetSchema'
import {
buildModelTypeTagUpdate,
getAssetAdditionalTags,
getAssetBaseModels,
getAssetDescription,
getAssetDisplayName,
getAssetFilename,
getAssetModelType,
getAssetSourceUrl,
getAssetTriggerPhrases,
getAssetUserDescription,
getEditableModelType,
getSourceName
} from '@/platform/assets/utils/assetMetadataUtils'
import { useAssetsStore } from '@/stores/assetsStore'
@@ -265,6 +268,7 @@ const { asset, cacheKey, selectContentStyle } = defineProps<{
}>()
const assetsStore = useAssetsStore()
const { flags } = useFeatureFlags()
const { modelTypes } = useModelTypes()
const pendingUpdates = ref<AssetUserMetadata>({})
@@ -272,6 +276,9 @@ const pendingModelType = ref<string | undefined>(undefined)
const isEditingDisplayName = ref(false)
const isImmutable = computed(() => asset.is_immutable ?? true)
// Retagging a model rewrites its asset tags; core is filesystem-backed and does
// not yet move the file to match, so the model type is read-only off-cloud.
const isModelTypeEditable = computed(() => !isImmutable.value && isCloud)
const displayName = computed(
() => pendingUpdates.value.name ?? getAssetDisplayName(asset)
)
@@ -318,12 +325,17 @@ function handleDisplayNameEdit(newName: string) {
}
const debouncedSaveModelType = useDebounceFn((newModelType: string) => {
if (isImmutable.value) return
const currentModelType = getAssetModelType(asset)
if (!isModelTypeEditable.value) return
const currentModelType = getEditableModelType(
asset,
flags.supportsModelTypeTags
)
if (currentModelType === newModelType) return
const newTags = asset.tags
.filter((tag) => tag !== currentModelType)
.concat(newModelType)
const newTags = buildModelTypeTagUpdate(
asset,
newModelType,
flags.supportsModelTypeTags
)
assetsStore.updateAssetTags(asset, newTags, cacheKey)
}, 500)
@@ -345,7 +357,10 @@ const userDescription = computed({
})
const selectedModelType = computed({
get: () => pendingModelType.value ?? getAssetModelType(asset) ?? undefined,
get: () =>
pendingModelType.value ??
getEditableModelType(asset, flags.supportsModelTypeTags) ??
undefined,
set: (value: string | undefined) => {
if (!value) return
pendingModelType.value = value

View File

@@ -25,10 +25,22 @@ vi.mock('@/i18n', () => ({
d: (date: Date) => date.toLocaleDateString()
}))
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
vi.mock('@/composables/useFeatureFlags', () => ({
useFeatureFlags: () => ({
flags: {
get supportsModelTypeTags() {
return mockSupportsModelTypeTags.value
}
}
})
}))
describe('useAssetBrowser', () => {
beforeEach(() => {
setActivePinia(createPinia())
vi.restoreAllMocks()
mockSupportsModelTypeTags.value = false
})
// Test fixtures - minimal data focused on functionality being tested
@@ -138,6 +150,25 @@ describe('useAssetBrowser', () => {
})
})
it('strips the model_type: prefix from the badge when the flag is on', () => {
mockSupportsModelTypeTags.value = true
const apiAsset = createApiAsset({
tags: ['models', 'model_type:checkpoints', 'sdxl']
})
const { filteredAssets } = useAssetBrowser(ref([apiAsset]))
const result = filteredAssets.value[0]
expect(result.badges).toContainEqual({
label: 'checkpoints',
type: 'type'
})
expect(result.badges).not.toContainEqual({
label: 'model_type:checkpoints',
type: 'type'
})
})
it('handles tags with multiple slashes in badges', () => {
const apiAsset = createApiAsset({
tags: ['models', 'checkpoint/subfolder/model-name']
@@ -668,6 +699,34 @@ describe('useAssetBrowser', () => {
])
})
it('groups by model_type:* value and ignores other tags when the flag is on', () => {
mockSupportsModelTypeTags.value = true
const assets = [
createApiAsset({ tags: ['models', 'model_type:checkpoints', 'sdxl'] }),
createApiAsset({ tags: ['models', 'model_type:LLM'] })
]
const { navItems } = useAssetBrowser(ref(assets))
const typeGroup = navItems.value[2] as { items: { id: string }[] }
expect(typeGroup.items.map((i) => i.id)).toEqual(['LLM', 'checkpoints'])
})
it('ignores model_type: and groups by bare tags when the flag is off', () => {
const assets = [
createApiAsset({ tags: ['models', 'model_type:checkpoints'] }),
createApiAsset({ tags: ['models', 'model_type:LLM'] })
]
const { navItems } = useAssetBrowser(ref(assets))
const typeGroup = navItems.value[2] as { items: { id: string }[] }
expect(typeGroup.items.map((i) => i.id)).toEqual([
'model_type:LLM',
'model_type:checkpoints'
])
})
it('handles assets with no category tag', () => {
const assets = [
createApiAsset({ tags: ['models'] }), // No second tag

View File

@@ -19,10 +19,13 @@ import {
} from '@/platform/assets/utils/assetFilterUtils'
import {
getAssetBaseModels,
getAssetFilename
getAssetCategories,
getAssetFilename,
getAssetTypeBadges
} from '@/platform/assets/utils/assetMetadataUtils'
import { MODELS_TAG } from '@/platform/assets/services/assetService'
import { sortAssets } from '@/platform/assets/utils/assetSortUtils'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { useAssetDownloadStore } from '@/stores/assetDownloadStore'
import type { NavGroupData, NavItemData } from '@/types/navTypes'
@@ -43,18 +46,19 @@ export interface AssetDisplayItem extends AssetItem {
}
}
const displayItemCache = new WeakMap<AssetItem, AssetDisplayItem>()
const displayItemCache = new WeakMap<
AssetItem,
{ modelTypeMode: boolean; item: AssetDisplayItem }
>()
function buildDisplayItem(asset: AssetItem): AssetDisplayItem {
function buildDisplayItem(
asset: AssetItem,
modelTypeMode: boolean
): AssetDisplayItem {
const badges: AssetBadge[] = []
const typeTag = asset.tags.find((tag) => tag !== 'models')
if (typeTag) {
const badgeLabel = typeTag.includes('/')
? typeTag.substring(typeTag.indexOf('/') + 1)
: typeTag
badges.push({ label: badgeLabel, type: 'type' })
for (const typeBadge of getAssetTypeBadges(asset, modelTypeMode)) {
badges.push({ label: typeBadge, type: 'type' })
}
for (const model of getAssetBaseModels(asset)) {
@@ -75,12 +79,15 @@ function buildDisplayItem(asset: AssetItem): AssetDisplayItem {
}
}
function transformAssetForDisplay(asset: AssetItem): AssetDisplayItem {
function transformAssetForDisplay(
asset: AssetItem,
modelTypeMode: boolean
): AssetDisplayItem {
const cached = displayItemCache.get(asset)
if (cached) return cached
const built = buildDisplayItem(asset)
displayItemCache.set(asset, built)
return built
if (cached && cached.modelTypeMode === modelTypeMode) return cached.item
const item = buildDisplayItem(asset, modelTypeMode)
displayItemCache.set(asset, { modelTypeMode, item })
return item
}
/**
@@ -93,6 +100,7 @@ export function useAssetBrowser(
const assets = computed<AssetItem[]>(() => assetsSource.value ?? [])
const assetDownloadStore = useAssetDownloadStore()
const { sessionDownloadCount } = storeToRefs(assetDownloadStore)
const { flags } = useFeatureFlags()
// State
const searchQuery = ref('')
@@ -122,12 +130,10 @@ export function useAssetBrowser(
})
const typeCategories = computed<NavItemData[]>(() => {
const modelTypeMode = flags.supportsModelTypeTags
const categories = assets.value
.filter((asset) => asset.tags.includes(MODELS_TAG))
.flatMap((asset) =>
asset.tags.filter((tag) => tag !== MODELS_TAG && tag.length > 0)
)
.map((tag) => tag.split('/')[0])
.flatMap((asset) => getAssetCategories(asset, modelTypeMode))
return Array.from(new Set(categories))
.sort()
@@ -191,7 +197,9 @@ export function useAssetBrowser(
// Category-filtered assets for filter options (before search/format/base model filters)
const categoryFilteredAssets = computed(() => {
return assets.value.filter(filterByCategory(selectedCategory.value))
return assets.value.filter(
filterByCategory(selectedCategory.value, flags.supportsModelTypeTags)
)
})
const { availableFileFormats, availableBaseModels } = useAssetFilterOptions(
@@ -248,7 +256,10 @@ export function useAssetBrowser(
const sortedAssets = sortAssets(filtered, filters.value.sortBy)
// Transform to display format
return sortedAssets.map(transformAssetForDisplay)
const modelTypeMode = flags.supportsModelTypeTags
return sortedAssets.map((asset) =>
transformAssetForDisplay(asset, modelTypeMode)
)
})
function updateFilters(newFilters: AssetFilterState) {

View File

@@ -38,7 +38,10 @@ vi.mock('@/scripts/api', () => ({
api: {
fetchApi: vi.fn(),
addEventListener: vi.fn(),
apiURL: vi.fn((path: string) => path)
apiURL: vi.fn((path: string) => path),
getServerFeature: vi.fn(
(_name: string, defaultValue?: unknown) => defaultValue
)
}
}))
@@ -279,6 +282,43 @@ describe('useUploadModelWizard', () => {
expect(result?.modelType).toBe('checkpoints')
})
it('namespaces the tag but keeps user_metadata.model_type bare when the backend supports it', async () => {
const { assetService } =
await import('@/platform/assets/services/assetService')
const { api } = await import('@/scripts/api')
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
type: 'sync',
asset: {
id: 'asset-1',
name: 'model.safetensors',
tags: ['models', 'model_type:checkpoints']
}
})
vi.mocked(api.getServerFeature).mockImplementation((name, defaultValue) =>
name === 'supports_model_type_tags' ? true : defaultValue
)
try {
const wizard = setupUploadModelWizard(modelTypes, {
requiredModelType: 'checkpoints'
})
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
await wizard.uploadModel()
const uploadArg = vi.mocked(assetService.uploadAssetAsync).mock
.calls[0][0]
expect(uploadArg.tags).toEqual(['models', 'model_type:checkpoints'])
expect(uploadArg.user_metadata?.model_type).toBe('checkpoints')
// The namespaced returned tag must not trip the required-type guard.
expect(wizard.uploadTypeMismatch.value).toBeNull()
} finally {
vi.mocked(api.getServerFeature).mockImplementation(
(_name, defaultValue) => defaultValue
)
}
})
it('returns the synced asset filename for sync imports', async () => {
const { assetService } =
await import('@/platform/assets/services/assetService')
@@ -347,6 +387,65 @@ describe('useUploadModelWizard', () => {
})
})
it('treats a namespaced model_type: tag as satisfying the required type', async () => {
const { assetService } =
await import('@/platform/assets/services/assetService')
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
type: 'sync',
asset: {
id: 'asset-1',
name: 'model.safetensors',
tags: ['models', 'model_type:checkpoints']
}
})
const wizard = setupUploadModelWizard(
ref([
{ name: 'Checkpoint', value: 'checkpoints' },
{ name: 'LoRA', value: 'loras' }
]),
{ requiredModelType: 'checkpoints' }
)
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
const result = await wizard.uploadModel()
expect(result).not.toBeNull()
expect(wizard.uploadTypeMismatch.value).toBeNull()
})
it('strips the model_type: prefix from the imported-type label on a real mismatch', async () => {
const { assetService } =
await import('@/platform/assets/services/assetService')
vi.mocked(assetService.uploadAssetAsync).mockResolvedValue({
type: 'sync',
asset: {
id: 'asset-lora',
name: 'model.safetensors',
tags: ['models', 'model_type:loras']
}
})
const wizard = setupUploadModelWizard(
ref([
{ name: 'Checkpoint', value: 'checkpoints' },
{ name: 'LoRA', value: 'loras' }
]),
{ requiredModelType: 'checkpoints' }
)
wizard.wizardData.value.url = 'https://civitai.com/models/12345'
const result = await wizard.uploadModel()
expect(result).toBeNull()
expect(wizard.uploadTypeMismatch.value).toEqual({
importedModelType: 'loras',
importedModelTypeLabel: 'LoRA',
requiredModelType: 'checkpoints',
requiredModelTypeLabel: 'Checkpoint'
})
})
it('does not block sync imports as mismatches without a required model type', async () => {
const { assetService } =
await import('@/platform/assets/services/assetService')

View File

@@ -2,6 +2,7 @@ import type { Ref } from 'vue'
import { computed, ref, watch } from 'vue'
import { useI18n } from 'vue-i18n'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { st } from '@/i18n'
import { civitaiImportSource } from '@/platform/assets/importSources/civitaiImportSource'
import { huggingfaceImportSource } from '@/platform/assets/importSources/huggingfaceImportSource'
@@ -11,7 +12,11 @@ import type {
} from '@/platform/assets/schemas/assetSchema'
import { assetService } from '@/platform/assets/services/assetService'
import type { ImportSource } from '@/platform/assets/types/importSource'
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
import {
getAssetFilename,
stripModelTypePrefix,
toModelTypeTag
} from '@/platform/assets/utils/assetMetadataUtils'
import { validateSourceUrl } from '@/platform/assets/utils/importSourceUtil'
import { useAssetDownloadStore } from '@/stores/assetDownloadStore'
import { useAssetsStore } from '@/stores/assetsStore'
@@ -68,6 +73,7 @@ export function useUploadModelWizard(
options: UploadModelWizardOptions = {}
) {
const { t } = useI18n()
const { flags } = useFeatureFlags()
const assetsStore = useAssetsStore()
const assetDownloadStore = useAssetDownloadStore()
const modelToNodeStore = useModelToNodeStore()
@@ -271,19 +277,22 @@ export function useUploadModelWizard(
}
function getImportedModelType(asset: AssetItem): string | undefined {
const knownType = asset.tags.find(
(tag) =>
tag !== MODEL_ROOT_TAG &&
const subtypeTags = asset.tags
.filter((tag) => tag !== MODEL_ROOT_TAG)
.map(stripModelTypePrefix)
return (
subtypeTags.find((tag) =>
modelTypes.value.some((type) => type.value === tag)
) ?? subtypeTags[0]
)
return knownType ?? asset.tags.find((tag) => tag !== MODEL_ROOT_TAG)
}
function blockMismatchedImportedModel(
asset: AssetItem,
requiredType: string
): boolean {
if (asset.tags.includes(requiredType)) return false
if (asset.tags.map(stripModelTypePrefix).includes(requiredType))
return false
const importedType = getImportedModelType(asset)
uploadStatus.value = 'error'
@@ -317,7 +326,11 @@ export function useUploadModelWizard(
try {
const modelType = resolvedModelType.value
const tags = modelType ? ['models', modelType] : ['models']
const subtypeTag =
modelType && flags.supportsModelTypeTags
? toModelTypeTag(modelType)
: modelType
const tags = subtypeTag ? [MODEL_ROOT_TAG, subtypeTag] : [MODEL_ROOT_TAG]
const filename =
wizardData.value.metadata?.filename ||
wizardData.value.metadata?.name ||

View File

@@ -11,6 +11,8 @@ const zAsset = z.object({
tags: z.array(z.string()).optional().default([]),
preview_id: z.string().nullable().optional(),
display_name: z.string().optional(),
/** Path within the model's category folder, i.e. the value a loader widget expects. */
loader_path: z.string().nullish(),
preview_url: z.string().optional(),
thumbnail_url: z.string().optional(),
created_at: z.string().optional(),
@@ -27,11 +29,6 @@ const zAssetResponse = zListAssetsResponse
assets: z.array(zAsset)
})
const zModelFolder = z.object({
name: z.string(),
folders: z.array(z.string())
})
// Zod schema for ModelFile to align with interface
const zModelFile = z.object({
name: z.string(),
@@ -100,7 +97,6 @@ export type AssetItem = z.infer<typeof zAsset>
export type AssetResponse = z.infer<typeof zAssetResponse>
export type AssetMetadata = z.infer<typeof zAssetMetadata>
export type AsyncUploadResponse = z.infer<typeof zAsyncUploadResponse>
export type ModelFolder = z.infer<typeof zModelFolder>
export type ModelFile = z.infer<typeof zModelFile>
/** Payload for updating an asset via PUT /assets/:id */
@@ -132,4 +128,10 @@ export type TagsOperationResult = z.infer<typeof tagsOperationResultSchema>
export interface ModelFolderInfo {
name: string
folders: string[]
/**
* The folder's raw registered extension allowlist from
* `/experiment/models`. An empty array means match-all; absent on older
* backends.
*/
extensions?: string[]
}

View File

@@ -12,6 +12,7 @@ import { api } from '@/scripts/api'
const mockDistributionState = vi.hoisted(() => ({ isCloud: false }))
const mockSettingStoreGet = vi.hoisted(() => vi.fn(() => false))
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: true }))
vi.mock('@/platform/distribution/types', () => ({
get isCloud() {
@@ -19,6 +20,16 @@ vi.mock('@/platform/distribution/types', () => ({
}
}))
vi.mock('@/composables/useFeatureFlags', () => ({
useFeatureFlags: () => ({
flags: {
get supportsModelTypeTags() {
return mockSupportsModelTypeTags.value
}
}
})
}))
vi.mock('@/platform/settings/settingStore', () => ({
useSettingStore: vi.fn(() => ({
get: mockSettingStoreGet
@@ -40,7 +51,9 @@ vi.mock('@/stores/modelToNodeStore', () => {
vi.mock('@/scripts/api', () => ({
api: {
fetchApi: vi.fn()
fetchApi: vi.fn(),
addCustomEventListener: vi.fn(),
removeCustomEventListener: vi.fn()
}
}))
@@ -87,6 +100,7 @@ function validAsset(overrides: Partial<AssetItem> = {}): AssetItem {
return {
id: 'asset-1',
name: 'model.safetensors',
loader_path: overrides.name ?? 'model.safetensors',
tags: ['models'],
...overrides
}
@@ -416,32 +430,329 @@ describe(assetService.deleteAsset, () => {
})
})
describe(assetService.getAssetModelFolders, () => {
describe(assetService.getAssetModels, () => {
beforeEach(() => {
vi.clearAllMocks()
assetService.invalidateModelBuckets()
mockSupportsModelTypeTags.value = true
})
it('walks the models tag once, excluding missing assets', async () => {
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({ id: 'a', tags: ['models', 'model_type:checkpoints'] })
])
)
await assetService.getAssetModels('checkpoints')
expect(fetchApiMock).toHaveBeenCalledTimes(1)
const requestedUrl = fetchApiMock.mock.calls[0]?.[0] as string
const params = new URL(requestedUrl, 'http://localhost').searchParams
expect(params.get('include_tags')).toBe('models')
expect(params.get('exclude_tags')).toBe(MISSING_TAG)
})
it('buckets by bare tags when model_type tags are unsupported', async () => {
mockSupportsModelTypeTags.value = false
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'a',
name: 'a.safetensors',
tags: ['models', 'checkpoints']
})
])
)
const models = await assetService.getAssetModels('checkpoints')
expect(models).toEqual([{ name: 'a.safetensors', pathIndex: 0 }])
})
it('drops uncategorized model assets with a warning', async () => {
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'ok',
name: 'ok.safetensors',
tags: ['models', 'model_type:loras']
}),
validAsset({
id: 'uncat',
name: 'orphan.safetensors',
tags: ['models']
})
])
)
const loras = await assetService.getAssetModels('loras')
expect(loras).toEqual([{ name: 'ok.safetensors', pathIndex: 0 }])
expect(warn).toHaveBeenCalledWith(
expect.stringContaining('orphan.safetensors')
)
warn.mockRestore()
})
it('maps loader_path and drops unloadable assets without one', async () => {
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'nested',
name: 'model.safetensors',
loader_path: 'sdxl/model.safetensors',
tags: ['models', 'model_type:checkpoints']
}),
validAsset({
id: 'orphan',
name: 'orphan.safetensors',
loader_path: null,
tags: ['models', 'model_type:checkpoints']
}),
validAsset({
id: 'other-folder',
name: 'lora.safetensors',
tags: ['models', 'model_type:loras']
})
])
)
const models = await assetService.getAssetModels('checkpoints')
expect(models).toEqual([{ name: 'sdxl/model.safetensors', pathIndex: 0 }])
expect(warn).toHaveBeenCalledWith(
expect.stringContaining('orphan.safetensors')
)
warn.mockRestore()
})
it('drops assets whose loader path is traversal-shaped', async () => {
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {})
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'evil',
name: 'evil.safetensors',
loader_path: '../../secrets/evil.safetensors',
tags: ['models', 'model_type:checkpoints']
}),
validAsset({
id: 'ok',
name: 'fine.safetensors',
tags: ['models', 'model_type:checkpoints']
})
])
)
const models = await assetService.getAssetModels('checkpoints')
expect(models).toEqual([{ name: 'fine.safetensors', pathIndex: 0 }])
expect(warn).toHaveBeenCalledWith(expect.stringContaining('unsafe'))
warn.mockRestore()
})
it('groups slashed bare tags by their top-level segment', async () => {
mockSupportsModelTypeTags.value = false
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'slashed',
name: 'model1.safetensors',
tags: ['models', 'Chatterbox/subfolder1/model1']
})
])
)
const models = await assetService.getAssetModels('Chatterbox')
expect(models).toEqual([{ name: 'model1.safetensors', pathIndex: 0 }])
})
it('falls back to filename metadata then name on bare-tag backends', async () => {
mockSupportsModelTypeTags.value = false
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'cloud-hash',
name: 'blake3-content-hash',
loader_path: null,
user_metadata: { filename: 'sdxl/cloud-model.safetensors' },
tags: ['models', 'checkpoints']
}),
validAsset({
id: 'bare',
name: 'plain.safetensors',
loader_path: null,
tags: ['models', 'checkpoints']
})
])
)
const models = await assetService.getAssetModels('checkpoints')
expect(models).toEqual([
{ name: 'sdxl/cloud-model.safetensors', pathIndex: 0 },
{ name: 'plain.safetensors', pathIndex: 0 }
])
})
it('orders each folder subdirectories-first then files, alphabetically', async () => {
const checkpointAsset = (id: string, loaderPath: string) =>
validAsset({
id,
name: loaderPath.split('/').pop()!,
loader_path: loaderPath,
tags: ['models', 'model_type:checkpoints']
})
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
checkpointAsset('1', 'sdxl/base.safetensors'),
checkpointAsset('2', 'v1-5.safetensors'),
checkpointAsset('3', 'sdxl/refiner.safetensors'),
checkpointAsset('4', 'anything.safetensors'),
checkpointAsset('5', 'dynamicrafter/model.safetensors')
])
)
const models = await assetService.getAssetModels('checkpoints')
expect(models.map((m) => m.name)).toEqual([
'dynamicrafter/model.safetensors',
'sdxl/base.safetensors',
'sdxl/refiner.safetensors',
'anything.safetensors',
'v1-5.safetensors'
])
})
it('does not let a stale in-flight walk overwrite an invalidated cache', async () => {
let resolveStaleWalk!: (response: Response) => void
fetchApiMock.mockReturnValueOnce(
new Promise<Response>((resolve) => {
resolveStaleWalk = resolve
})
)
const staleRead = assetService.getAssetModels('checkpoints')
assetService.invalidateModelBuckets()
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'fresh',
name: 'fresh.safetensors',
tags: ['models', 'model_type:checkpoints']
})
])
)
const freshModels = await assetService.getAssetModels('checkpoints')
expect(freshModels.map((m) => m.name)).toEqual(['fresh.safetensors'])
resolveStaleWalk(
buildAssetListResponse([
validAsset({
id: 'stale',
name: 'stale.safetensors',
tags: ['models', 'model_type:checkpoints']
})
])
)
await staleRead
const cachedModels = await assetService.getAssetModels('checkpoints')
expect(cachedModels.map((m) => m.name)).toEqual(['fresh.safetensors'])
expect(fetchApiMock).toHaveBeenCalledTimes(2)
})
it('places multi-category assets in every folder from a single walk', async () => {
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({
id: 'shared',
name: 'dual_use.safetensors',
loader_path: 'dual_use.safetensors',
tags: [
'models',
'model_type:checkpoints',
'model_type:diffusion_models'
]
})
])
)
const checkpoints = await assetService.getAssetModels('checkpoints')
const diffusion = await assetService.getAssetModels('diffusion_models')
expect(checkpoints).toEqual([
{ name: 'dual_use.safetensors', pathIndex: 0 }
])
expect(diffusion).toEqual([{ name: 'dual_use.safetensors', pathIndex: 0 }])
// Both folder reads resolve from a single memoized models walk.
expect(fetchApiMock).toHaveBeenCalledTimes(1)
})
})
describe(assetService.onModelsScanned, () => {
beforeEach(() => {
vi.clearAllMocks()
})
it('requests missing-tag exclusion and returns alphabetical unique folders without include_public', async () => {
it('invokes the callback when the scan event fires and unsubscribes cleanly', () => {
const callback = vi.fn()
const unsubscribe = assetService.onModelsScanned(callback)
const [eventType, handler] = vi.mocked(api.addCustomEventListener).mock
.calls[0]!
expect(eventType).toBe('assets.seed.fast_complete')
handler!(new CustomEvent(eventType))
expect(callback).toHaveBeenCalledOnce()
unsubscribe()
expect(api.removeCustomEventListener).toHaveBeenCalledWith(
eventType,
handler
)
})
})
describe(assetService.seedModelAssets, () => {
beforeEach(() => {
vi.clearAllMocks()
})
it('POSTs the models root to the seed endpoint', async () => {
fetchApiMock.mockResolvedValueOnce(
buildAssetListResponse([
validAsset({ id: 'a', tags: ['models', 'loras'] }),
validAsset({ id: 'b', tags: ['models', 'checkpoints'] }),
validAsset({ id: 'c', tags: ['models', 'configs'] }),
validAsset({ id: 'e', tags: ['models', 'loras'] })
])
buildResponse({ status: 'started' }, { status: 202 })
)
const folders = await assetService.getAssetModelFolders()
await assetService.seedModelAssets()
expect(folders).toEqual([
{ name: 'checkpoints', folders: [] },
{ name: 'loras', folders: [] }
])
expect(fetchApiMock).toHaveBeenCalledWith('/assets/seed', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ roots: ['models'] })
})
})
const requestedUrl = fetchApiMock.mock.calls[0]?.[0] as string
const params = new URL(requestedUrl, 'http://localhost').searchParams
expect(params.has('include_public')).toBe(false)
expect(params.get('exclude_tags')).toBe(MISSING_TAG)
it('treats an already-running scan (409) as success', async () => {
fetchApiMock.mockResolvedValueOnce(
buildResponse({ status: 'already_running' }, { ok: false, status: 409 })
)
await expect(assetService.seedModelAssets()).resolves.toBeUndefined()
})
it('throws on other error statuses', async () => {
fetchApiMock.mockResolvedValueOnce(
buildResponse({}, { ok: false, status: 500 })
)
await expect(assetService.seedModelAssets()).rejects.toThrow('500')
})
})

View File

@@ -1,9 +1,11 @@
import { fromZodError } from 'zod-validation-error'
import { z } from 'zod'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { st } from '@/i18n'
import {
assetFilenameSchema,
assetItemSchema,
assetResponseSchema,
asyncUploadResponseSchema,
@@ -17,9 +19,9 @@ import type {
AssetUpdatePayload,
AsyncUploadResponse,
ModelFile,
ModelFolder,
TagsOperationResult
} from '@/platform/assets/schemas/assetSchema'
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
import { isCloud } from '@/platform/distribution/types'
import { useSettingStore } from '@/platform/settings/settingStore'
import { api } from '@/scripts/api'
@@ -180,6 +182,7 @@ function getLocalizedErrorMessage(errorCode: string): string {
}
const ASSETS_ENDPOINT = '/assets'
const ASSETS_SEED_ENDPOINT = '/assets/seed'
const ASSETS_DOWNLOAD_ENDPOINT = '/assets/download'
const ASSETS_EXPORT_ENDPOINT = '/assets/export'
const EXPERIMENTAL_WARNING = `EXPERIMENTAL: If you are seeing this please make sure "Comfy.Assets.UseAssetAPI" is set to "false" in your ComfyUI Settings.\n`
@@ -187,6 +190,8 @@ const DEFAULT_LIMIT = 500
const INPUT_ASSETS_WITH_PUBLIC_LIMIT = 500
export const MODELS_TAG = 'models'
/** Prefix for the namespaced tag that carries a model's folder category, e.g. `model_type:checkpoints`. */
const MODEL_TYPE_TAG_PREFIX = 'model_type:'
export const INPUT_TAG = 'input'
export const OUTPUT_TAG = 'output'
/** Asset tag used by the backend for placeholder records that are not installed. */
@@ -209,6 +214,48 @@ function normalizeAssetTags(tags: string[]): string[] {
return tags.map((tag) => tag.trim()).filter(Boolean)
}
/**
* Resolves the model folder a tag represents, or undefined when the tag is not
* a folder category. `supports_model_type_tags` backends carry the category as
* a namespaced `model_type:<folder>` tag; older backends mint bare tags, which
* may carry subfolder paths (e.g. `Chatterbox/sub/model`) and group by their
* top-level segment, matching the asset browser's legacy grouping.
*/
function modelFolderFromTag(
tag: string,
modelTypeMode: boolean
): string | undefined {
if (modelTypeMode) {
return tag.startsWith(MODEL_TYPE_TAG_PREFIX)
? tag.slice(MODEL_TYPE_TAG_PREFIX.length)
: undefined
}
if (tag === MODELS_TAG || tag.length === 0) return undefined
return tag.split('/')[0]
}
/**
* Orders loader paths as subdirectories before files at every level,
* alphabetical within each group. The asset API returns models in storage
* order, which would otherwise interleave root-level files with folder
* contents in the sidebar tree.
*/
function compareLoaderPaths(a: string, b: string): number {
const aSegments = a.split('/')
const bSegments = b.split('/')
const sharedDepth = Math.min(aSegments.length, bSegments.length)
for (let i = 0; i < sharedDepth; i++) {
const aIsFile = i === aSegments.length - 1
const bIsFile = i === bSegments.length - 1
if (aIsFile !== bIsFile) return aIsFile ? 1 : -1
const order = aSegments[i].localeCompare(bSegments[i], undefined, {
numeric: true
})
if (order !== 0) return order
}
return 0
}
async function withCallerAbort<T>(
promise: Promise<T>,
signal?: AbortSignal
@@ -269,6 +316,26 @@ function createAssetService() {
let inputAssetsIncludingPublicRequestId = 0
let pendingInputAssetsIncludingPublic: Promise<AssetItem[]> | null = null
/**
* Model assets bucketed by folder category, built from a single walk of the
* `models` tag rather than a fetch per category. Shared by the folder list
* and per-folder listings so the sidebar loads every model in one pass.
*/
let modelBuckets: Map<string, AssetItem[]> | null = null
let modelBucketsRequestId = 0
let pendingModelBuckets: Promise<Map<string, AssetItem[]>> | null = null
/**
* Discards the cached model buckets so the next read re-walks the models
* tag. Bumping the request id keeps a walk that was already in flight from
* repopulating the cache with pre-invalidation data.
*/
function invalidateModelBuckets(): void {
modelBucketsRequestId++
modelBuckets = null
pendingModelBuckets = null
}
/** Invalidates the cached public-inclusive input assets without aborting in-flight readers. */
function invalidateInputAssetsIncludingPublic(): void {
inputAssetsIncludingPublicRequestId++
@@ -330,51 +397,156 @@ function createAssetService() {
return validateAssetResponse(data)
}
/**
* Gets a list of model folder keys from the asset API
*
* Logic:
* 1. Extract directory names directly from asset tags
* 2. Filter out blacklisted directories
* 3. Return alphabetically sorted directories with assets
*
* @returns The list of model folder keys
* Walks every `models`-tagged asset once and buckets each into the folder
* categories carried by its `model_type:` tags. A single asset lands in every
* category it is tagged with (e.g. a shared-root model in both `checkpoints`
* and `diffusion_models`). Which folders are actually shown is decided by
* `/experiment/models`; models with no category tag are dropped with a warning
* rather than hidden silently.
*/
async function getAssetModelFolders(): Promise<ModelFolder[]> {
const data = await handleAssetRequest(
{ includeTags: [MODELS_TAG] },
'model folders'
)
async function buildModelBuckets(): Promise<Map<string, AssetItem[]>> {
const assets = await getAllAssetsByTag(MODELS_TAG, true)
const modelTypeMode = useFeatureFlags().flags.supportsModelTypeTags
const buckets = new Map<string, AssetItem[]>()
// Blacklist directories we don't want to show
const blacklistedDirectories = new Set(['configs'])
for (const asset of assets) {
const folders = asset.tags
.map((tag) => modelFolderFromTag(tag, modelTypeMode))
.filter((folder): folder is string => folder !== undefined)
const folderTags = data.assets
.flatMap((asset) => asset.tags)
.filter((tag) => tag !== MODELS_TAG && !blacklistedDirectories.has(tag))
const discoveredFolders = new Set<string>(folderTags)
if (folders.length === 0) {
console.warn(
`Asset ${asset.id} (${asset.name}) is tagged '${MODELS_TAG}' but has no model category; skipping.`
)
continue
}
// Return only discovered folders in alphabetical order
const sortedFolders = Array.from(discoveredFolders).toSorted()
return sortedFolders.map((name) => ({ name, folders: [] }))
// On loader_path-contract backends a null loader_path marks an
// unloadable asset (e.g. an orphan): it must not mint a widget value,
// and `name` is deprecated for path semantics.
if (modelTypeMode && !asset.loader_path) {
console.warn(
`Asset ${asset.id} (${asset.name}) has no loader_path; skipping.`
)
continue
}
// The loader value flows into viewMetadata URLs and widget values, so a
// traversal-shaped path must not pass through even if the backend's own
// validation ever regresses.
const loaderValue = asset.loader_path ?? getAssetFilename(asset)
if (!assetFilenameSchema.safeParse(loaderValue).success) {
console.warn(
`Asset ${asset.id} (${asset.name}) has an unsafe loader path ('${loaderValue}'); skipping.`
)
continue
}
for (const folder of folders) {
const bucket = buckets.get(folder)
if (bucket) bucket.push(asset)
else buckets.set(folder, [asset])
}
}
for (const bucket of buckets.values()) {
bucket.sort((a, b) =>
compareLoaderPaths(
a.loader_path ?? getAssetFilename(a),
b.loader_path ?? getAssetFilename(b)
)
)
}
return buckets
}
/** Returns the memoized model buckets, walking the models tag on first read. */
async function loadModelBuckets(): Promise<Map<string, AssetItem[]>> {
if (modelBuckets) return modelBuckets
if (pendingModelBuckets) return pendingModelBuckets
const requestId = ++modelBucketsRequestId
const walk = async () => {
try {
const buckets = await buildModelBuckets()
if (requestId === modelBucketsRequestId) {
modelBuckets = buckets
}
return buckets
} finally {
if (requestId === modelBucketsRequestId) {
pendingModelBuckets = null
}
}
}
pendingModelBuckets = walk()
return pendingModelBuckets
}
/**
* Gets a list of models in the specified folder from the asset API
* Gets the models in the specified folder from the single models walk.
* @param folder The folder to list models from, such as 'checkpoints'
* @returns The list of model filenames within the specified folder
*/
async function getAssetModels(folder: string): Promise<ModelFile[]> {
const data = await handleAssetRequest(
{ includeTags: [MODELS_TAG, folder] },
`models for ${folder}`
)
return data.assets.map((asset) => ({
name: asset.name,
const buckets = await loadModelBuckets()
return (buckets.get(folder) ?? []).map((asset) => ({
// `loader_path` is the category-relative path the loader widget expects
// and the source for the sidebar tree. Backends that predate it (bare-tag
// mode; today's cloud) fall back to the filename metadata — the same
// value the asset browser serializes — rather than `name`, which is a
// content hash on cloud.
name: asset.loader_path ?? getAssetFilename(asset),
// Asset records carry no root identity, so every model reports root 0.
// Known limitation on multi-root categories (extra_model_paths.yaml):
// preview reads target root 0 (wrong file or 404 for secondary-root
// files), and same-relative-path files in different roots collapse
// onto one sidebar row. Metadata is unaffected unless relative paths
// collide (/view_metadata searches roots in order without an index),
// as are loader widget values; lifting this needs the backend to carry
// root identity on assets.
pathIndex: 0
}))
}
/**
* Asks the backend to rescan the model roots on disk so newly added files
* become assets. Fire-and-forget: the scan's fast (insert) phase already
* writes the category tags and filenames the sidebar needs and is announced
* by an `assets.seed.fast_complete` websocket event. A 409 means a scan is
* already running, which will emit the same event, so it is not an error.
*/
async function seedModelAssets(): Promise<void> {
const res = await api.fetchApi(ASSETS_SEED_ENDPOINT, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ roots: ['models'] })
})
if (!res.ok && res.status !== 409) {
throw new Error(
`Unable to start asset scan: Server returned ${res.status}`
)
}
}
/**
* Subscribes to the backend's scan fast-phase completion broadcast — the
* moment newly scanned files' tags and loader paths become queryable. The
* wire-level event (`assets.seed.fast_complete`) is owned here; consumers
* receive a callback and an unsubscribe function.
*/
function onModelsScanned(callback: () => void | Promise<void>): () => void {
const handler = () => {
void callback()
}
api.addCustomEventListener('assets.seed.fast_complete', handler)
return () => {
api.removeCustomEventListener('assets.seed.fast_complete', handler)
}
}
/**
* Checks if a widget input should use the asset browser based on both input name and node comfyClass
*
@@ -971,8 +1143,10 @@ function createAssetService() {
}
return {
getAssetModelFolders,
getAssetModels,
invalidateModelBuckets,
onModelsScanned,
seedModelAssets,
isAssetAPIEnabled,
isAssetBrowserEligible,
shouldUseAssetBrowser,

View File

@@ -45,7 +45,7 @@ describe('assetFilterUtils properties', () => {
it('filterByCategory("all") accepts every asset', () => {
fc.assert(
fc.property(arbAssetItem, (asset) => {
expect(filterByCategory('all')(asset)).toBe(true)
expect(filterByCategory('all', false)(asset)).toBe(true)
})
)
})
@@ -56,7 +56,7 @@ describe('assetFilterUtils properties', () => {
fc.array(arbAssetItem, { maxLength: 30 }),
fc.stringMatching(/^[a-z]{1,8}$/),
(assets, category) => {
const filter = filterByCategory(category)
const filter = filterByCategory(category, false)
const result = assets.filter(filter)
expect(result.length).toBeLessThanOrEqual(assets.length)
for (const item of result) {

View File

@@ -26,19 +26,54 @@ function createAsset(
describe('filterByCategory', () => {
it.for([
{ category: 'all', tags: ['checkpoint'], expected: true },
{ category: 'checkpoint', tags: ['checkpoint'], expected: true },
{ category: 'lora', tags: ['checkpoint'], expected: false },
{ category: 'all', tags: ['checkpoint'], mode: false, expected: true },
{
category: 'checkpoint',
tags: ['checkpoint'],
mode: false,
expected: true
},
{ category: 'lora', tags: ['checkpoint'], mode: false, expected: false },
{
category: 'checkpoint',
tags: ['models', 'checkpoint/xl'],
mode: false,
expected: true
},
{ category: 'xl', tags: ['models', 'checkpoint/xl'], expected: false }
{
category: 'xl',
tags: ['models', 'checkpoint/xl'],
mode: false,
expected: false
},
{
category: 'checkpoints',
tags: ['models', 'model_type:checkpoints'],
mode: true,
expected: true
},
{
category: 'LLM',
tags: ['models', 'model_type:LLM'],
mode: true,
expected: true
},
{
category: 'sdxl',
tags: ['models', 'model_type:checkpoints', 'sdxl'],
mode: true,
expected: false
},
{
category: 'checkpoints',
tags: ['models', 'model_type:checkpoints'],
mode: false,
expected: false
}
])(
'category=$category with tags=$tags returns $expected',
({ category, tags, expected }) => {
const filter = filterByCategory(category)
'category=$category tags=$tags mode=$mode returns $expected',
({ category, tags, mode, expected }) => {
const filter = filterByCategory(category, mode)
const asset = createAsset('model.safetensors', { tags })
expect(filter(asset)).toBe(expected)
}

View File

@@ -1,21 +1,14 @@
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import type { OwnershipOption } from '@/platform/assets/types/filterTypes'
import { getAssetBaseModels } from '@/platform/assets/utils/assetMetadataUtils'
import {
getAssetBaseModels,
getAssetCategories
} from '@/platform/assets/utils/assetMetadataUtils'
export function filterByCategory(category: string) {
export function filterByCategory(category: string, modelTypeMode: boolean) {
return (asset: AssetItem) => {
if (category === 'all') return true
// Check if any tag matches the category (for exact matches)
if (asset.tags.includes(category)) return true
// Check if any tag's top-level folder matches the category
return asset.tags.some((tag) => {
if (typeof tag === 'string' && tag.includes('/')) {
return tag.split('/')[0] === category
}
return false
})
return getAssetCategories(asset, modelTypeMode).includes(category)
}
}

View File

@@ -2,22 +2,34 @@ import { afterEach, describe, expect, it, vi } from 'vitest'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import {
MISSING_TAG,
MODELS_TAG
} from '@/platform/assets/services/assetService'
import {
buildModelTypeTagUpdate,
getAssetAdditionalTags,
getAssetBaseModel,
getAssetBaseModels,
getAssetCardTitle,
getAssetCategories,
getAssetDescription,
getAssetDisplayFilename,
getAssetDisplayName,
getAssetFilename,
getAssetMetadataDimensions,
getAssetModelType,
getAssetNodeCategoryCandidates,
getAssetSourceUrl,
getPrimaryCategoryTag,
getAssetStoredFilename,
getAssetTriggerPhrases,
getAssetTypeBadges,
getAssetUserDescription,
getEditableModelType,
getSourceName,
resolveDisplayImageDimensions
resolveDisplayImageDimensions,
stripModelTypePrefix,
toModelTypeTag
} from '@/platform/assets/utils/assetMetadataUtils'
const { isCloudRef } = vi.hoisted(() => ({
@@ -274,7 +286,17 @@ describe('assetMetadataUtils', () => {
tags: ['models'],
expected: null
},
{ name: 'returns null when tags empty', tags: [], expected: null }
{ name: 'returns null when tags empty', tags: [], expected: null },
{
name: 'never returns a raw model_type: literal (no round-trip into edit widgets)',
tags: ['models', 'model_type:checkpoints'],
expected: null
},
{
name: 'skips model_type: tags in favour of the bare twin',
tags: ['models', 'model_type:checkpoints', 'checkpoints'],
expected: 'checkpoints'
}
])('$name', ({ tags, expected }) => {
const asset = { ...mockAsset, tags }
expect(getAssetModelType(asset)).toBe(expected)
@@ -540,3 +562,353 @@ describe('assetMetadataUtils', () => {
})
})
})
describe('getAssetCategories', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('uses model_type:* values as the group and disregards other tags in model_type mode', () => {
expect(
getAssetCategories(
asset(['models', 'model_type:checkpoints', 'sdxl']),
true
)
).toEqual(['checkpoints'])
})
it('preserves the model_type value casing', () => {
expect(
getAssetCategories(asset(['models', 'model_type:LLM']), true)
).toEqual(['LLM'])
})
it('routes an uncovered asset by its bare tags in model_type mode', () => {
expect(getAssetCategories(asset(['models', 'checkpoints']), true)).toEqual([
'checkpoints'
])
})
it('ignores model_type: and uses bare-tag grouping when mode is off', () => {
expect(
getAssetCategories(
asset(['models', 'model_type:checkpoints', 'sdxl']),
false
)
).toEqual(['model_type:checkpoints', 'sdxl'])
})
it('never surfaces namespace residue as a category for uncovered assets in mode', () => {
expect(
getAssetCategories(asset(['models', 'model_type:', 'sdxl']), true)
).toEqual(['sdxl'])
})
})
describe('getPrimaryCategoryTag', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('uses the model_type value a covered asset groups under', () => {
expect(
getPrimaryCategoryTag(asset(['models', 'sdxl', 'model_type:vae']), true)
).toBe('vae')
})
it('keeps the legacy verbatim tag for an uncovered hierarchical asset', () => {
expect(
getPrimaryCategoryTag(asset(['models', 'Chatterbox/sub/model']), true)
).toBe('Chatterbox/sub/model')
})
it('skips namespace residue instead of titling off a raw model_type: tag', () => {
expect(
getPrimaryCategoryTag(asset(['models', 'model_type:']), true)
).toBeUndefined()
})
it('returns the legacy first non-models tag when mode is off', () => {
expect(
getPrimaryCategoryTag(asset(['models', 'model_type:vae']), false)
).toBe('model_type:vae')
})
})
describe('getAssetNodeCategoryCandidates', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('orders the most specific (deepest) tag ahead of a flat model_type value', () => {
expect(
getAssetNodeCategoryCandidates(
asset(['models', 'model_type:LLM', 'LLM/Qwen-VL/Qwen3-0.6B']),
true
)
).toEqual(['LLM/Qwen-VL/Qwen3-0.6B', 'LLM'])
})
it('strips the model_type: prefix when it is the only candidate', () => {
expect(
getAssetNodeCategoryCandidates(asset(['models', 'model_type:vae']), true)
).toEqual(['vae'])
})
it('keeps a model_type value ahead of an equally-deep bare tag', () => {
expect(
getAssetNodeCategoryCandidates(
asset(['models', 'model_type:checkpoints', 'sdxl']),
true
)
).toEqual(['checkpoints', 'sdxl'])
})
it('demotes an unrelated deeper bare tag below the model_type value', () => {
expect(
getAssetNodeCategoryCandidates(
asset(['models', 'model_type:vae', 'foo/bar']),
true
)
).toEqual(['vae', 'foo/bar'])
})
it('keeps unrelated bare tags as trailing fallbacks rather than dropping them', () => {
expect(
getAssetNodeCategoryCandidates(
asset(['models', 'model_type:LLM', 'LLM/Qwen-VL', 'foo/bar/baz']),
true
)
).toEqual(['LLM/Qwen-VL', 'LLM', 'foo/bar/baz'])
})
it('keeps a hierarchical tag intact', () => {
expect(
getAssetNodeCategoryCandidates(
asset(['models', 'chatterbox/chatterbox_vc']),
true
)
).toEqual(['chatterbox/chatterbox_vc'])
})
it('returns no candidates when only reserved tags are present', () => {
expect(
getAssetNodeCategoryCandidates(asset(['models', 'missing']), true)
).toEqual([])
})
it('uses the first non-reserved tag verbatim when mode is off', () => {
expect(
getAssetNodeCategoryCandidates(asset(['models', 'model_type:vae']), false)
).toEqual(['model_type:vae'])
expect(
getAssetNodeCategoryCandidates(asset(['models', 'checkpoints']), false)
).toEqual(['checkpoints'])
})
})
describe('getAssetTypeBadges', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('strips the model_type: prefix in model_type mode (no raw leak)', () => {
expect(
getAssetTypeBadges(
asset(['models', 'model_type:checkpoints', 'sdxl']),
true
)
).toEqual(['checkpoints'])
})
it('badges the model_type value even when a bare tag comes first, matching the grouping', () => {
expect(
getAssetTypeBadges(asset(['models', 'foo', 'model_type:bar']), true)
).toEqual(['bar'])
})
it('badges every category a shared multi-type asset groups under', () => {
expect(
getAssetTypeBadges(
asset([
'models',
'model_type:checkpoints',
'model_type:diffusion_models'
]),
true
)
).toEqual(['checkpoints', 'diffusion_models'])
})
it('falls back to the bare tag for an uncovered asset in model_type mode', () => {
expect(getAssetTypeBadges(asset(['models', 'sdxl']), true)).toEqual([
'sdxl'
])
})
it('returns no badge rather than a blank one for a malformed empty model_type: tag', () => {
expect(getAssetTypeBadges(asset(['models', 'model_type:']), true)).toEqual(
[]
)
})
it('leaks the literal model_type: tag when mode is off', () => {
expect(
getAssetTypeBadges(asset(['models', 'model_type:checkpoints']), false)
).toEqual(['model_type:checkpoints'])
})
it('shows the segment after the first slash for a bare hierarchical tag', () => {
expect(
getAssetTypeBadges(asset(['models', 'checkpoint/xl']), false)
).toEqual(['xl'])
})
it('returns no badges when only the models tag is present', () => {
expect(getAssetTypeBadges(asset(['models']), true)).toEqual([])
})
})
describe('stripModelTypePrefix', () => {
it('removes the model_type: prefix when present', () => {
expect(stripModelTypePrefix('model_type:checkpoints')).toBe('checkpoints')
})
it('leaves a tag without the prefix unchanged', () => {
expect(stripModelTypePrefix('checkpoints')).toBe('checkpoints')
expect(stripModelTypePrefix('checkpoint/xl')).toBe('checkpoint/xl')
})
})
describe('toModelTypeTag', () => {
it('prefixes a folder_name with the model_type namespace', () => {
expect(toModelTypeTag('checkpoints')).toBe('model_type:checkpoints')
expect(toModelTypeTag('ultralytics_bbox')).toBe(
'model_type:ultralytics_bbox'
)
})
})
describe('getEditableModelType', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('returns the stripped model_type value in model_type mode', () => {
expect(
getEditableModelType(
asset(['models', 'checkpoints', 'model_type:checkpoints']),
true
)
).toBe('checkpoints')
})
it('falls back to the bare tag for an uncovered asset in model_type mode', () => {
expect(getEditableModelType(asset(['models', 'sam2']), true)).toBe('sam2')
})
it('uses the legacy first-non-models tag when mode is off', () => {
expect(
getEditableModelType(asset(['models', 'checkpoints', 'sdxl']), false)
).toBe('checkpoints')
})
it('returns null when only the models tag is present', () => {
expect(getEditableModelType(asset(['models']), true)).toBeNull()
})
})
describe('buildModelTypeTagUpdate', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it('swaps the bare subtype tag when mode is off', () => {
expect(
buildModelTypeTagUpdate(asset(['models', 'checkpoints']), 'loras', false)
).toEqual(['models', 'loras'])
})
it('preserves user labels and swaps only the subtype tag when mode is off', () => {
expect(
buildModelTypeTagUpdate(
asset(['models', 'checkpoints', 'sdxl']),
'loras',
false
)
).toEqual(['models', 'sdxl', 'loras'])
})
it('writes only the model_type form for a covered asset, leaving the bare twin for the backend', () => {
expect(
buildModelTypeTagUpdate(
asset(['models', 'checkpoints', 'model_type:checkpoints']),
'loras',
true
)
).toEqual(['models', 'checkpoints', 'model_type:loras'])
})
it('replaces every existing model_type form for a shared-path dual-tagged asset', () => {
expect(
buildModelTypeTagUpdate(
asset([
'models',
'diffusion_models',
'model_type:diffusion_models',
'model_type:unet_gguf'
]),
'loras',
true
)
).toEqual(['models', 'diffusion_models', 'model_type:loras'])
})
it('drops the bare current type for an uncovered asset in model_type mode', () => {
expect(
buildModelTypeTagUpdate(asset(['models', 'sam2']), 'loras', true)
).toEqual(['models', 'model_type:loras'])
})
it('keeps user labels untouched in model_type mode', () => {
expect(
buildModelTypeTagUpdate(
asset(['models', 'checkpoints', 'model_type:checkpoints', 'sdxl']),
'loras',
true
)
).toEqual(['models', 'checkpoints', 'sdxl', 'model_type:loras'])
})
})
describe('reserved tag mirrors', () => {
const asset = (tags: string[]): AssetItem => ({
id: 'a',
name: 'model.safetensors',
tags
})
it("treats assetService's canonical reserved tags as reserved (locals must not drift)", () => {
expect(getAssetCategories(asset([MODELS_TAG, 'x']), false)).toEqual(['x'])
expect(
getAssetNodeCategoryCandidates(
asset([MODELS_TAG, MISSING_TAG, 'x']),
true
)
).toEqual(['x'])
expect(getAssetTypeBadges(asset([MODELS_TAG, 'x']), false)).toEqual(['x'])
expect(getAssetModelType(asset([MODELS_TAG]))).toBeNull()
})
})

View File

@@ -2,6 +2,11 @@ import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { isCloud } from '@/platform/distribution/types'
import { isCivitaiUrl } from '@/utils/formatUtil'
// Reserved tag literals (mirror assetService's MODELS_TAG/MISSING_TAG). Kept
// local so this leaf util doesn't pull the heavier assetService -> i18n chain.
const MODELS_TAG = 'models'
const MISSING_TAG = 'missing'
/**
* Type-safe utilities for extracting metadata from assets.
* These utilities check user_metadata first, then metadata, then fallback.
@@ -140,16 +145,236 @@ export function getSourceName(url: string): string {
return 'Source'
}
export const MODEL_TYPE_TAG_PREFIX = 'model_type:'
/**
* Extracts the model type from asset tags
* Extracts the model type from asset tags as a bare (non-namespaced) value.
* Never returns a raw `model_type:*` literal: this value feeds edit widgets
* whose save path writes tags back verbatim, so a namespaced tag leaking
* through here would round-trip the prefixed literal into the tag set.
* @param asset - The asset to extract model type from
* @returns The model type string or null if not present
*/
export function getAssetModelType(asset: AssetItem): string | null {
const typeTag = asset.tags?.find((tag) => tag && tag !== 'models')
const typeTag = asset.tags?.find(
(tag) => tag && tag !== MODELS_TAG && !tag.startsWith(MODEL_TYPE_TAG_PREFIX)
)
return typeTag ?? null
}
/**
* Builds the namespaced subtype tag the backend stores in `model_type:` mode.
* The argument is a discovery folder_name (e.g. `checkpoints`,
* `ultralytics_bbox`); the backend keeps the bare directory-path twin in sync.
*/
export function toModelTypeTag(folderName: string): string {
return `${MODEL_TYPE_TAG_PREFIX}${folderName}`
}
/** Strips the `model_type:` prefix off each namespaced tag, dropping non-`model_type:` tags. */
function getModelTypeTagValues(asset: AssetItem): string[] {
return asset.tags
.filter((tag) => tag.startsWith(MODEL_TYPE_TAG_PREFIX))
.map((tag) => tag.slice(MODEL_TYPE_TAG_PREFIX.length))
.filter((tag) => tag.length > 0)
}
/**
* Resolves the folder_name shown as the asset's current model type in the edit
* dropdown. In `modelTypeMode` the stripped `model_type:` value is authoritative
* (covered assets); an uncovered asset with no `model_type:` tag falls back to
* its bare subtype tag, mirroring the read-side grouping. Outside the mode this
* is the legacy first-non-`models` tag.
*/
export function getEditableModelType(
asset: AssetItem,
modelTypeMode: boolean
): string | null {
if (modelTypeMode) {
const [modelType] = getModelTypeTagValues(asset)
if (modelType) return modelType
}
return getAssetModelType(asset)
}
/**
* Computes the tag set for re-typing a model asset to `newFolderName`. In
* `modelTypeMode` only the `model_type:` form is written — the backend keeps the
* bare directory-path twin in sync, so existing `model_type:` tags are dropped
* (covered assets) or the bare current type is dropped (uncovered assets) and
* the new `model_type:<folder_name>` is added. Outside the mode it swaps the
* legacy bare subtype tag, preserving the pre-namespace behavior.
*/
export function buildModelTypeTagUpdate(
asset: AssetItem,
newFolderName: string,
modelTypeMode: boolean
): string[] {
if (!modelTypeMode) {
const currentType = getAssetModelType(asset)
return asset.tags.filter((tag) => tag !== currentType).concat(newFolderName)
}
const modelTypeTags = asset.tags.filter((tag) =>
tag.startsWith(MODEL_TYPE_TAG_PREFIX)
)
const currentBareType = getAssetModelType(asset)
const tagsToRemove =
modelTypeTags.length > 0
? new Set(modelTypeTags)
: new Set(currentBareType ? [currentBareType] : [])
return asset.tags
.filter((tag) => !tagsToRemove.has(tag))
.concat(toModelTypeTag(newFolderName))
}
/** Legacy grouping: each non-`models` tag's top-level path segment. */
function getBareTagCategories(asset: AssetItem): string[] {
return asset.tags
.filter((tag) => tag !== MODELS_TAG && tag.length > 0)
.map((tag) => tag.split('/')[0])
}
/**
* Resolves the category keys a model asset is grouped under.
*
* `modelTypeMode` reflects whether the backend declares the `model_type:` tag
* scheme (the `supports_model_type_tags` capability). When true, an asset's
* `model_type:*` values are authoritative; an asset with no `model_type:` tag
* still routes by its bare tags. When false (the default) categories come from
* the legacy bare-tag top-level grouping and `model_type:` is ignored.
*/
export function getAssetCategories(
asset: AssetItem,
modelTypeMode: boolean
): string[] {
if (modelTypeMode) {
const modelTypes = getModelTypeTagValues(asset)
if (modelTypes.length > 0) return modelTypes
// Uncovered assets route by bare tags, but namespace residue (e.g. a
// malformed empty `model_type:`) must not surface as a raw category.
return getBareTagCategories(asset).filter(
(category) => !category.startsWith(MODEL_TYPE_TAG_PREFIX)
)
}
return getBareTagCategories(asset)
}
/**
* Resolves the primary tag a browser surface titles itself after. In
* `modelTypeMode` a covered asset uses its first `model_type:*` value — the
* key it groups under — while an uncovered asset keeps the legacy selection
* (first non-`models` tag, verbatim, hierarchical paths intact). Outside the
* mode this is exactly the legacy selection.
*/
export function getPrimaryCategoryTag(
asset: AssetItem,
modelTypeMode: boolean
): string | undefined {
if (modelTypeMode) {
const [modelType] = getModelTypeTagValues(asset)
if (modelType) return modelType
return asset.tags.find(
(tag) => tag !== MODELS_TAG && !tag.startsWith(MODEL_TYPE_TAG_PREFIX)
)
}
return asset.tags.find((tag) => tag !== MODELS_TAG)
}
/** Number of `parent/child` segments in a tag, used to pick the most specific. */
function pathDepth(tag: string): number {
return tag.split('/').length
}
/** Removes the `model_type:` namespace prefix from a tag when present. */
export function stripModelTypePrefix(tag: string): string {
return tag.startsWith(MODEL_TYPE_TAG_PREFIX)
? tag.slice(MODEL_TYPE_TAG_PREFIX.length)
: tag
}
/**
* Resolves the labels shown as an asset card's type badges.
*
* In `modelTypeMode` a covered asset badges every `model_type:*` value — the
* same keys it groups under (`getAssetCategories`) — so a shared-root asset
* tagged with several categories carries each of them; whichever category
* view the card appears in is represented on the card. Uncovered assets (and
* legacy mode) keep the original single selection: first non-`models` tag,
* with bare hierarchical tags showing the segment after the first `/`.
*/
export function getAssetTypeBadges(
asset: AssetItem,
modelTypeMode: boolean
): string[] {
if (modelTypeMode) {
const modelTypes = getModelTypeTagValues(asset)
if (modelTypes.length > 0) return modelTypes
}
const typeTag = asset.tags.find(
(tag) =>
tag !== MODELS_TAG &&
!(modelTypeMode && tag.startsWith(MODEL_TYPE_TAG_PREFIX))
)
if (!typeTag) return []
return [
typeTag.includes('/') ? typeTag.slice(typeTag.indexOf('/') + 1) : typeTag
]
}
/**
* Ordered node-category candidates for an asset, most specific first.
*
* Callers resolve a node provider by trying each candidate in order and taking
* the first that maps to a provider. The full (possibly hierarchical) value is
* kept so `modelToNodeStore`'s `parent/child` fallback still works.
*
* In `modelTypeMode` (backend declares `supports_model_type_tags`) candidates
* come in two tiers. Tier 1: the stripped `model_type:*` values plus bare tags
* *related* to one of them (equal to it, or extending it as a `parent/child`
* path), ordered by descending depth — so a resolvable `LLM/Qwen-VL/...` twin
* wins over a flat `model_type:LLM`, while ties keep `model_type:*` values
* ahead of bare tags. Tier 2: unrelated bare tags (e.g. a user-added
* `foo/bar`), tried only when nothing authoritative resolves — they can no
* longer pre-empt a resolvable `model_type:*` value however deep they are.
* An uncovered asset (no `model_type:` tag) routes by all its bare tags,
* deepest first. Outside `modelTypeMode` the legacy first-non-reserved tag is
* used verbatim.
*/
export function getAssetNodeCategoryCandidates(
asset: AssetItem,
modelTypeMode: boolean
): string[] {
if (!modelTypeMode) {
const legacy = asset.tags.find(
(tag) => tag !== MODELS_TAG && tag !== MISSING_TAG
)
return legacy ? [legacy] : []
}
const bareTags = asset.tags.filter(
(tag) =>
tag !== MODELS_TAG &&
tag !== MISSING_TAG &&
!tag.startsWith(MODEL_TYPE_TAG_PREFIX)
)
const byDepthDesc = (a: string, b: string) => pathDepth(b) - pathDepth(a)
const modelTypes = getModelTypeTagValues(asset)
if (modelTypes.length === 0) return bareTags.toSorted(byDepthDesc)
const isRelated = (tag: string) =>
modelTypes.some((type) => tag === type || tag.startsWith(`${type}/`))
return [
...[...modelTypes, ...bareTags.filter(isRelated)].sort(byDepthDesc),
...bareTags.filter((tag) => !isRelated(tag)).sort(byDepthDesc)
]
}
/**
* Extracts user description from asset user_metadata
* @param asset - The asset to extract user description from

View File

@@ -1,14 +1,25 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { resolveModelNodeFromAsset } from '@/platform/assets/utils/resolveModelNodeFromAsset'
const mockGetNodeProvider = vi.hoisted(() => vi.fn())
const mockSupportsModelTypeTags = vi.hoisted(() => ({ value: false }))
vi.mock('@/stores/modelToNodeStore', () => ({
useModelToNodeStore: () => ({ getNodeProvider: mockGetNodeProvider })
}))
vi.mock('@/composables/useFeatureFlags', () => ({
useFeatureFlags: () => ({
flags: {
get supportsModelTypeTags() {
return mockSupportsModelTypeTags.value
}
}
})
}))
function createMockAsset(overrides: Partial<AssetItem> = {}): AssetItem {
return {
id: 'asset-123',
@@ -49,6 +60,11 @@ describe('resolveModelNodeFromAsset', () => {
beforeEach(() => {
vi.clearAllMocks()
vi.spyOn(console, 'error').mockImplementation(() => {})
mockSupportsModelTypeTags.value = false
})
afterEach(() => {
vi.restoreAllMocks()
})
describe('valid assets', () => {
@@ -68,6 +84,48 @@ describe('resolveModelNodeFromAsset', () => {
}
})
it('strips the model_type: prefix when resolving the provider in model_type mode', () => {
mockSupportsModelTypeTags.value = true
mockProvider(createMockNodeProvider())
const result = resolveModelNodeFromAsset(
createMockAsset({ tags: ['models', 'model_type:vae'] })
)
expect(result.success).toBe(true)
expect(mockGetNodeProvider).toHaveBeenCalledWith('vae')
})
it('skips an unresolvable incidental tag and resolves via the model_type value', () => {
mockSupportsModelTypeTags.value = true
mockGetNodeProvider.mockImplementation((category: string) =>
category === 'vae' ? createMockNodeProvider() : undefined
)
const result = resolveModelNodeFromAsset(
createMockAsset({ tags: ['models', 'model_type:vae', 'foo/bar'] })
)
expect(result.success).toBe(true)
expect(mockGetNodeProvider).toHaveBeenCalledWith('foo/bar')
expect(mockGetNodeProvider).toHaveBeenCalledWith('vae')
})
it('prefers the deepest resolvable path over a flat model_type value', () => {
mockSupportsModelTypeTags.value = true
mockGetNodeProvider.mockImplementation((category: string) =>
category === 'LLM/Qwen-VL/Qwen3-0.6B'
? createMockNodeProvider()
: undefined
)
const result = resolveModelNodeFromAsset(
createMockAsset({
tags: ['models', 'model_type:LLM', 'LLM/Qwen-VL/Qwen3-0.6B']
})
)
expect(result.success).toBe(true)
expect(mockGetNodeProvider).toHaveBeenCalledWith('LLM/Qwen-VL/Qwen3-0.6B')
})
it('falls back to metadata.filename when user_metadata.filename missing', () => {
mockProvider(createMockNodeProvider())
const result = resolveModelNodeFromAsset(
@@ -201,7 +259,7 @@ describe('resolveModelNodeFromAsset', () => {
if (!result.success) {
expect(result.error.code).toBe('NO_PROVIDER')
expect(result.error.message).toContain('checkpoints')
expect(result.error.details?.category).toBe('checkpoints')
expect(result.error.details?.candidates).toEqual(['checkpoints'])
}
})
})

View File

@@ -4,7 +4,11 @@ import {
MISSING_TAG,
MODELS_TAG
} from '@/platform/assets/services/assetService'
import { getAssetFilename } from '@/platform/assets/utils/assetMetadataUtils'
import {
getAssetFilename,
getAssetNodeCategoryCandidates
} from '@/platform/assets/utils/assetMetadataUtils'
import { useFeatureFlags } from '@/composables/useFeatureFlags'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
import type { ModelNodeProvider } from '@/stores/modelToNodeStore'
@@ -81,10 +85,12 @@ export function resolveModelNodeFromAsset(
}
}
const category = validAsset.tags.find(
(tag) => tag !== MODELS_TAG && tag !== MISSING_TAG
const { flags } = useFeatureFlags()
const candidates = getAssetNodeCategoryCandidates(
validAsset,
flags.supportsModelTypeTags
)
if (!category) {
if (candidates.length === 0) {
console.error(
`Asset ${validAsset.id} has no valid category tag. Available tags: ${validAsset.tags.join(', ')} (expected tag other than '${MODELS_TAG}' or '${MISSING_TAG}')`
)
@@ -99,19 +105,31 @@ export function resolveModelNodeFromAsset(
}
}
const provider = useModelToNodeStore().getNodeProvider(category)
if (!provider) {
console.error(`No node provider registered for category: ${category}`)
const modelToNodeStore = useModelToNodeStore()
const resolved = candidates
.map((category) => ({
category,
provider: modelToNodeStore.getNodeProvider(category)
}))
.find((candidate) => candidate.provider !== undefined)
if (!resolved?.provider) {
// Known gap (out of scope for FE-1076): flat `model_type:LLM`-style tags
// whose loaders are only registered hierarchically land here until the
// backend emits a subtype-carrying tag.
console.error(
`No node provider registered for category: ${candidates.join(', ')}`
)
return {
success: false,
error: {
code: 'NO_PROVIDER',
message: `No node provider registered for category: ${category}`,
message: `No node provider registered for category: ${candidates.join(', ')}`,
assetId: validAsset.id,
details: { category }
details: { candidates }
}
}
}
return { success: true, value: { provider, filename } }
return { success: true, value: { provider: resolved.provider, filename } }
}

View File

@@ -1215,6 +1215,15 @@ export const CORE_SETTINGS: SettingParams[] = [
defaultValue: isCloud ? true : false,
experimental: true
},
{
id: 'Comfy.ModelLibrary.UseAssetBrowser',
name: 'Use the asset browser for the model library',
type: 'hidden',
tooltip:
'When enabled alongside the asset API, the model library opens the asset browser. Otherwise it opens the sidebar tree.',
defaultValue: isCloud ? true : false,
experimental: true
},
{
id: 'Comfy.VersionCompatibility.DisableWarnings',
name: 'Disable version compatibility warnings',

View File

@@ -428,6 +428,7 @@ const zSettings = z.object({
'Comfy.VueNodes.Enabled': z.boolean(),
'Comfy.AppBuilder.VueNodeSwitchDismissed': z.boolean(),
'Comfy.Assets.UseAssetAPI': z.boolean(),
'Comfy.ModelLibrary.UseAssetBrowser': z.boolean(),
'Comfy.Queue.QPOV2': z.boolean(),
'Comfy.Queue.ShowRunProgressBar': z.boolean(),
'Comfy-Desktop.AutoUpdate': z.boolean(),

View File

@@ -875,6 +875,113 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
})
})
describe('pagination safety', () => {
it('stops instead of looping when the backend ignores offset', async () => {
const store = useAssetsStore()
const nodeType = 'CheckpointLoaderSimple'
// A backend that ignores offset returns the same full page every time.
const fullPage = Array.from({ length: 500 }, (_, i) =>
createMockAsset(`asset-${i}`)
)
vi.mocked(assetService.getAssetsForNodeType).mockResolvedValue(fullPage)
await store.updateModelsForNodeType(nodeType)
await vi.waitFor(() => {
expect(
vi.mocked(assetService.getAssetsForNodeType)
).toHaveBeenCalledTimes(2)
})
expect(
vi.mocked(assetService.getAssetsForNodeType)
).toHaveBeenCalledTimes(2)
expect(store.getAssets(nodeType)).toHaveLength(500)
})
it('continues past an all-duplicate page whose content differs from the previous page', async () => {
const store = useAssetsStore()
const nodeType = 'CheckpointLoaderSimple'
// Concurrent writes can shift pagination windows so a page is all
// already-seen assets without the backend ignoring offset; later pages
// can still hold unseen assets.
const fullPage = Array.from({ length: 500 }, (_, i) =>
createMockAsset(`asset-${i}`)
)
const samePageReordered = [...fullPage].reverse()
const finalPage = [createMockAsset('late-arrival')]
let callCount = 0
vi.mocked(assetService.getAssetsForNodeType).mockImplementation(
async () => {
callCount++
if (callCount === 1) return fullPage
if (callCount === 2) return samePageReordered
return finalPage
}
)
await store.updateModelsForNodeType(nodeType)
await vi.waitFor(() => {
expect(
vi.mocked(assetService.getAssetsForNodeType)
).toHaveBeenCalledTimes(3)
})
expect(store.getAssets(nodeType).map((a) => a.id)).toContain(
'late-arrival'
)
})
it('terminates when an offset-ignoring backend alternates page orderings', async () => {
const store = useAssetsStore()
const nodeType = 'CheckpointLoaderSimple'
// Same full page served forever with a nondeterministic ordering: no
// page ever contributes a new ID, and no two consecutive pages are
// identical. The walk must still stop.
const fullPage = Array.from({ length: 500 }, (_, i) =>
createMockAsset(`asset-${i}`)
)
const reversed = [...fullPage].reverse()
let callCount = 0
vi.mocked(assetService.getAssetsForNodeType).mockImplementation(
async () => {
callCount++
return callCount % 2 === 1 ? fullPage : reversed
}
)
await store.updateModelsForNodeType(nodeType)
expect(callCount).toBeLessThanOrEqual(5)
expect(store.getAssets(nodeType)).toHaveLength(500)
})
})
describe('refresh error surfacing', () => {
it('surfaces a failed refresh on the committed state consumers read', async () => {
const store = useAssetsStore()
const nodeType = 'CheckpointLoaderSimple'
vi.mocked(assetService.getAssetsForNodeType).mockResolvedValueOnce([
createMockAsset('existing')
])
await store.updateModelsForNodeType(nodeType)
expect(store.getError(nodeType)).toBeUndefined()
vi.mocked(assetService.getAssetsForNodeType).mockRejectedValueOnce(
new Error('backend down')
)
await store.updateModelsForNodeType(nodeType)
expect(store.getAssets(nodeType).map((a) => a.id)).toEqual(['existing'])
expect(store.getError(nodeType)?.message).toBe('backend down')
})
})
describe('concurrent request handling', () => {
it('should short-circuit concurrent calls to prevent duplicate work', async () => {
const store = useAssetsStore()
@@ -924,6 +1031,34 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
vi.mocked(assetService.getAssetsForNodeType)
).toHaveBeenCalledTimes(2)
})
it('keeps a newer request single-flighted when a stale request finishes after invalidation', async () => {
const store = useAssetsStore()
const nodeType = 'CheckpointLoaderSimple'
let resolveFirst!: (assets: AssetItem[]) => void
const firstFetch = new Promise<AssetItem[]>((resolve) => {
resolveFirst = resolve
})
vi.mocked(assetService.getAssetsForNodeType)
.mockReturnValueOnce(firstFetch)
.mockReturnValue(new Promise<AssetItem[]>(() => {}))
const staleRequest = store.updateModelsForNodeType(nodeType)
store.invalidateCategory('checkpoints')
void store.updateModelsForNodeType(nodeType)
resolveFirst([createMockAsset('stale')])
await staleRequest
// The stale request's teardown must not evict the newer request's
// single-flight entry: a third call short-circuits instead of starting
// a duplicate walk.
void store.updateModelsForNodeType(nodeType)
expect(
vi.mocked(assetService.getAssetsForNodeType)
).toHaveBeenCalledTimes(2)
})
})
describe('shallowReactive state reactivity', () => {
@@ -1435,6 +1570,35 @@ describe('assetsStore - Model Assets Cache (Cloud)', () => {
})
})
describe('assetsStore - Model Assets Cache (non-cloud)', () => {
beforeEach(() => {
setActivePinia(createTestingPinia({ stubActions: false }))
mockIsCloud.value = false
vi.clearAllMocks()
})
it('caches model assets fetched by tag on non-cloud builds', async () => {
const store = useAssetsStore()
vi.mocked(assetService.getAssetsByTag).mockResolvedValue([
{
id: 'm1',
name: 'sd_xl_base_1.0.safetensors',
tags: ['checkpoints', 'models']
},
{ id: 'm2', name: 'lora.safetensors', tags: ['loras', 'models'] }
])
await store.updateModelsForTag('models')
expect(assetService.getAssetsByTag).toHaveBeenCalledWith(
'models',
true,
expect.anything()
)
expect(store.getAssets('tag:models')).toHaveLength(2)
})
})
describe('assetsStore - Deletion State and Input Mapping', () => {
beforeEach(() => {
setActivePinia(createTestingPinia({ stubActions: false }))

View File

@@ -394,421 +394,444 @@ export const useAssetsStore = defineStore('assets', () => {
* Multiple node types sharing the same category share the same cache entry.
* Public API accepts nodeType for backwards compatibility but translates
* to category internally using modelToNodeStore.getCategoryForNodeType().
* Cloud-only feature - empty Maps in desktop builds
*
* Runs on every distribution; whether anything fetches through it is
* decided by consumers via `assetService.isAssetAPIEnabled()`, which stays
* the authoritative off-cloud gate.
*/
const getModelState = () => {
if (isCloud) {
const modelStateByCategory = ref(new Map<string, ModelPaginationState>())
const modelStateByCategory = ref(new Map<string, ModelPaginationState>())
const assetsArrayCache = new Map<
string,
{ source: Map<string, AssetItem>; array: AssetItem[] }
>()
const assetsArrayCache = new Map<
string,
{ source: Map<string, AssetItem>; array: AssetItem[] }
>()
const pendingRequestByCategory = new Map<string, ModelPaginationState>()
const pendingPromiseByCategory = new Map<string, Promise<void>>()
const pendingRequestByCategory = new Map<string, ModelPaginationState>()
const pendingPromiseByCategory = new Map<string, Promise<void>>()
function createState(
existingAssets?: Map<string, AssetItem>
): ModelPaginationState {
const assets = new Map(existingAssets)
return reactive({
assets,
offset: 0,
hasMore: true,
isLoading: true
})
function createState(
existingAssets?: Map<string, AssetItem>
): ModelPaginationState {
const assets = new Map(existingAssets)
return reactive({
assets,
offset: 0,
hasMore: true,
isLoading: true
})
}
function isStale(category: string, state: ModelPaginationState): boolean {
const committed = modelStateByCategory.value.get(category)
const pending = pendingRequestByCategory.get(category)
return committed !== state && pending !== state
}
const EMPTY_ASSETS: AssetItem[] = []
/**
* Resolve a key to a category. Handles both nodeType and tag:xxx formats.
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
* @returns The category or undefined if not resolvable
*/
function resolveCategory(key: string): string | undefined {
if (key.startsWith('tag:')) {
return key
}
return modelToNodeStore.getCategoryForNodeType(key)
}
/**
* Get assets by nodeType or tag key.
* Translates nodeType to category internally for cache lookup.
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
*/
function getAssets(key: string): AssetItem[] {
const category = resolveCategory(key)
if (!category) return EMPTY_ASSETS
const state = modelStateByCategory.value.get(category)
const assetsMap = state?.assets
if (!assetsMap) return EMPTY_ASSETS
const cached = assetsArrayCache.get(category)
if (cached && cached.source === assetsMap) {
return cached.array
}
function isStale(category: string, state: ModelPaginationState): boolean {
const committed = modelStateByCategory.value.get(category)
const pending = pendingRequestByCategory.get(category)
return committed !== state && pending !== state
const array = Array.from(assetsMap.values())
assetsArrayCache.set(category, { source: assetsMap, array })
return array
}
function isLoading(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.get(category)?.isLoading ?? false
}
function getError(key: string): Error | undefined {
const category = resolveCategory(key)
if (!category) return undefined
return modelStateByCategory.value.get(category)?.error
}
function hasMore(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.get(category)?.hasMore ?? false
}
function hasAssetKey(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.has(category)
}
/**
* Check if a category exists in the cache.
* Checks both direct category keys and tag-prefixed keys.
* @param category The category to check (e.g., 'checkpoints', 'loras')
*/
function hasCategory(category: string): boolean {
return (
modelStateByCategory.value.has(category) ||
modelStateByCategory.value.has(`tag:${category}`)
)
}
/**
* Internal helper to fetch and cache assets for a category.
* Loads first batch immediately, then progressively loads remaining batches.
* Keeps existing data visible until new data is successfully fetched.
*
* Concurrent calls for the same category are short-circuited: if a request
* is already in progress (tracked via pendingRequestByCategory), subsequent
* calls return immediately to avoid redundant work.
*/
async function updateModelsForCategory(
category: string,
fetcher: (options: PaginationOptions) => Promise<AssetItem[]>
): Promise<void> {
if (pendingPromiseByCategory.has(category)) {
return pendingPromiseByCategory.get(category)!
}
const EMPTY_ASSETS: AssetItem[] = []
const existingState = modelStateByCategory.value.get(category)
const state = createState(existingState?.assets)
/**
* Resolve a key to a category. Handles both nodeType and tag:xxx formats.
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
* @returns The category or undefined if not resolvable
*/
function resolveCategory(key: string): string | undefined {
if (key.startsWith('tag:')) {
return key
}
return modelToNodeStore.getCategoryForNodeType(key)
const seenIds = new Set<string>()
const seenPageSignatures = new Set<string>()
let consecutiveNoProgressPages = 0
const hasExistingData = modelStateByCategory.value.has(category)
if (hasExistingData) {
pendingRequestByCategory.set(category, state)
} else {
// Also track in pending map for initial loads to prevent concurrent calls
pendingRequestByCategory.set(category, state)
modelStateByCategory.value.set(category, state)
}
/**
* Get assets by nodeType or tag key.
* Translates nodeType to category internally for cache lookup.
* @param key Either a nodeType (e.g., 'CheckpointLoaderSimple') or tag key (e.g., 'tag:models')
*/
function getAssets(key: string): AssetItem[] {
const category = resolveCategory(key)
if (!category) return EMPTY_ASSETS
async function loadBatches(): Promise<void> {
while (state.hasMore) {
try {
const newAssets = await fetcher({
limit: MODEL_BATCH_SIZE,
offset: state.offset
})
const state = modelStateByCategory.value.get(category)
const assetsMap = state?.assets
if (!assetsMap) return EMPTY_ASSETS
if (isStale(category, state)) return
const cached = assetsArrayCache.get(category)
if (cached && cached.source === assetsMap) {
return cached.array
}
const array = Array.from(assetsMap.values())
assetsArrayCache.set(category, { source: assetsMap, array })
return array
}
function isLoading(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.get(category)?.isLoading ?? false
}
function getError(key: string): Error | undefined {
const category = resolveCategory(key)
if (!category) return undefined
return modelStateByCategory.value.get(category)?.error
}
function hasMore(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.get(category)?.hasMore ?? false
}
function hasAssetKey(key: string): boolean {
const category = resolveCategory(key)
if (!category) return false
return modelStateByCategory.value.has(category)
}
/**
* Check if a category exists in the cache.
* Checks both direct category keys and tag-prefixed keys.
* @param category The category to check (e.g., 'checkpoints', 'loras')
*/
function hasCategory(category: string): boolean {
return (
modelStateByCategory.value.has(category) ||
modelStateByCategory.value.has(`tag:${category}`)
)
}
/**
* Internal helper to fetch and cache assets for a category.
* Loads first batch immediately, then progressively loads remaining batches.
* Keeps existing data visible until new data is successfully fetched.
*
* Concurrent calls for the same category are short-circuited: if a request
* is already in progress (tracked via pendingRequestByCategory), subsequent
* calls return immediately to avoid redundant work.
*/
async function updateModelsForCategory(
category: string,
fetcher: (options: PaginationOptions) => Promise<AssetItem[]>
): Promise<void> {
if (pendingPromiseByCategory.has(category)) {
return pendingPromiseByCategory.get(category)!
}
const existingState = modelStateByCategory.value.get(category)
const state = createState(existingState?.assets)
const seenIds = new Set<string>()
const hasExistingData = modelStateByCategory.value.has(category)
if (hasExistingData) {
pendingRequestByCategory.set(category, state)
} else {
// Also track in pending map for initial loads to prevent concurrent calls
pendingRequestByCategory.set(category, state)
modelStateByCategory.value.set(category, state)
}
async function loadBatches(): Promise<void> {
while (state.hasMore) {
try {
const newAssets = await fetcher({
limit: MODEL_BATCH_SIZE,
offset: state.offset
})
if (isStale(category, state)) return
const isFirstBatch = state.offset === 0
if (isFirstBatch) {
assetsArrayCache.delete(category)
if (hasExistingData) {
pendingRequestByCategory.delete(category)
modelStateByCategory.value.set(category, state)
}
const isFirstBatch = state.offset === 0
if (isFirstBatch) {
assetsArrayCache.delete(category)
if (hasExistingData) {
pendingRequestByCategory.delete(category)
modelStateByCategory.value.set(category, state)
}
// Merge new assets into existing map and track seen IDs
for (const asset of newAssets) {
seenIds.add(asset.id)
state.assets.set(asset.id, asset)
}
state.assets = new Map(state.assets)
state.offset += newAssets.length
state.hasMore = newAssets.length === MODEL_BATCH_SIZE
if (isFirstBatch) {
state.isLoading = false
}
if (state.hasMore) {
await new Promise((resolve) => setTimeout(resolve, 50))
}
} catch (err) {
if (isStale(category, state)) return
console.error(`Error loading batch for ${category}:`, err)
state.error = err instanceof Error ? err : new Error(String(err))
state.hasMore = false
state.isLoading = false
pendingRequestByCategory.delete(category)
return
}
}
const staleIds = [...state.assets.keys()].filter(
(id) => !seenIds.has(id)
)
for (const id of staleIds) {
state.assets.delete(id)
// Merge new assets into existing map and track seen IDs
const uniqueIdsBefore = seenIds.size
for (const asset of newAssets) {
seenIds.add(asset.id)
state.assets.set(asset.id, asset)
}
state.assets = new Map(state.assets)
// Termination guards for backends that do not honour `offset`.
// A page whose exact ID sequence was already served means the
// walk is cycling, however the pages are ordered — stop. A single
// all-duplicate page with fresh content (concurrent writes
// shifting pagination windows) keeps going, but a run of them
// with no new IDs is treated as exhausted so reordered responses
// can never loop forever.
const batchSignature = newAssets.map((asset) => asset.id).join(',')
const isRepeatedPage =
newAssets.length > 0 && seenPageSignatures.has(batchSignature)
seenPageSignatures.add(batchSignature)
const madeProgress = seenIds.size > uniqueIdsBefore
consecutiveNoProgressPages = madeProgress
? 0
: consecutiveNoProgressPages + 1
state.offset += newAssets.length
state.hasMore =
newAssets.length === MODEL_BATCH_SIZE &&
!isRepeatedPage &&
consecutiveNoProgressPages < 3
if (isFirstBatch) {
state.isLoading = false
}
if (state.hasMore) {
await new Promise((resolve) => setTimeout(resolve, 50))
}
} catch (err) {
if (isStale(category, state)) return
console.error(`Error loading batch for ${category}:`, err)
state.error = err instanceof Error ? err : new Error(String(err))
state.hasMore = false
state.isLoading = false
// A refresh that fails before its first batch never replaces the
// committed state, so mirror the error onto the state consumers
// actually read (getError) instead of only the discarded one.
const committed = modelStateByCategory.value.get(category)
if (committed && committed !== state) {
committed.error = state.error
}
if (pendingRequestByCategory.get(category) === state) {
pendingRequestByCategory.delete(category)
}
return
}
assetsArrayCache.delete(category)
}
const staleIds = [...state.assets.keys()].filter(
(id) => !seenIds.has(id)
)
for (const id of staleIds) {
state.assets.delete(id)
}
assetsArrayCache.delete(category)
if (pendingRequestByCategory.get(category) === state) {
pendingRequestByCategory.delete(category)
}
}
const promise = loadBatches().finally(() => {
// Guard both cleanups: an invalidateCategory during an awaited fetch
// lets a newer request register its own entries before this one's
// teardown runs, and an unconditional delete would evict the newer
// request's entry and break single-flighting.
const promise = loadBatches().finally(() => {
if (pendingPromiseByCategory.get(category) === promise) {
pendingPromiseByCategory.delete(category)
})
pendingPromiseByCategory.set(category, promise)
await promise
}
/**
* Fetch and cache model assets for a specific node type.
* Translates nodeType to category internally - multiple node types
* sharing the same category will share the same cache entry.
* @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple')
*/
async function updateModelsForNodeType(nodeType: string): Promise<void> {
const category = modelToNodeStore.getCategoryForNodeType(nodeType)
if (!category) return
// Use category as cache key but fetch using nodeType for API compatibility
await updateModelsForCategory(category, (opts) =>
assetService.getAssetsForNodeType(nodeType, opts)
)
}
/**
* Fetch and cache model assets for a specific tag
* @param tag The tag to fetch assets for (e.g., 'models')
*/
async function updateModelsForTag(tag: string): Promise<void> {
const category = `tag:${tag}`
await updateModelsForCategory(category, (opts) =>
assetService.getAssetsByTag(tag, true, opts)
)
}
/**
* Invalidate the cache for a specific category.
* Forces a refetch on next access.
* @param category The category to invalidate (e.g., 'checkpoints', 'loras')
*/
function invalidateCategory(category: string): void {
modelStateByCategory.value.delete(category)
assetsArrayCache.delete(category)
pendingRequestByCategory.delete(category)
pendingPromiseByCategory.delete(category)
}
/**
* Optimistically update an asset in the cache
* @param assetId The asset ID to update
* @param updates Partial asset data to merge
* @param cacheKey Optional cache key to target (nodeType or 'tag:xxx')
*/
function updateAssetInCache(
assetId: string,
updates: Partial<AssetItem>,
cacheKey?: string
) {
const category = cacheKey ? resolveCategory(cacheKey) : undefined
if (cacheKey && !category) return
const categoriesToCheck = category
? [category]
: Array.from(modelStateByCategory.value.keys())
for (const cat of categoriesToCheck) {
const state = modelStateByCategory.value.get(cat)
if (!state?.assets) continue
const existingAsset = state.assets.get(assetId)
if (existingAsset) {
const updatedAsset = { ...existingAsset, ...updates }
state.assets.set(assetId, updatedAsset)
assetsArrayCache.delete(cat)
if (cacheKey) return
}
}
}
})
pendingPromiseByCategory.set(category, promise)
await promise
}
/**
* Update asset metadata with optimistic cache update
* @param asset The asset to update
* @param userMetadata The user_metadata to save
* @param cacheKey Optional cache key to target for optimistic update
*/
async function updateAssetMetadata(
asset: AssetItem,
userMetadata: Record<string, unknown>,
cacheKey?: string
) {
const originalMetadata = asset.user_metadata
updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey)
/**
* Fetch and cache model assets for a specific node type.
* Translates nodeType to category internally - multiple node types
* sharing the same category will share the same cache entry.
* @param nodeType The node type to fetch assets for (e.g., 'CheckpointLoaderSimple')
*/
async function updateModelsForNodeType(nodeType: string): Promise<void> {
const category = modelToNodeStore.getCategoryForNodeType(nodeType)
if (!category) return
try {
const updatedAsset = await assetService.updateAsset(asset.id, {
user_metadata: userMetadata
})
updateAssetInCache(asset.id, updatedAsset, cacheKey)
} catch (error) {
console.error('Failed to update asset metadata:', error)
updateAssetInCache(
asset.id,
{ user_metadata: originalMetadata },
cacheKey
)
// Use category as cache key but fetch using nodeType for API compatibility
await updateModelsForCategory(category, (opts) =>
assetService.getAssetsForNodeType(nodeType, opts)
)
}
/**
* Fetch and cache model assets for a specific tag
* @param tag The tag to fetch assets for (e.g., 'models')
*/
async function updateModelsForTag(tag: string): Promise<void> {
const category = `tag:${tag}`
await updateModelsForCategory(category, (opts) =>
assetService.getAssetsByTag(tag, true, opts)
)
}
/**
* Invalidate the cache for a specific category.
* Forces a refetch on next access.
* @param category The category to invalidate (e.g., 'checkpoints', 'loras')
*/
function invalidateCategory(category: string): void {
modelStateByCategory.value.delete(category)
assetsArrayCache.delete(category)
pendingRequestByCategory.delete(category)
pendingPromiseByCategory.delete(category)
}
/**
* Optimistically update an asset in the cache
* @param assetId The asset ID to update
* @param updates Partial asset data to merge
* @param cacheKey Optional cache key to target (nodeType or 'tag:xxx')
*/
function updateAssetInCache(
assetId: string,
updates: Partial<AssetItem>,
cacheKey?: string
) {
const category = cacheKey ? resolveCategory(cacheKey) : undefined
if (cacheKey && !category) return
const categoriesToCheck = category
? [category]
: Array.from(modelStateByCategory.value.keys())
for (const cat of categoriesToCheck) {
const state = modelStateByCategory.value.get(cat)
if (!state?.assets) continue
const existingAsset = state.assets.get(assetId)
if (existingAsset) {
const updatedAsset = { ...existingAsset, ...updates }
state.assets.set(assetId, updatedAsset)
assetsArrayCache.delete(cat)
if (cacheKey) return
}
}
/**
* Update asset tags using add/remove endpoints
* @param asset The asset to update (used to read current tags)
* @param newTags The desired tags array
* @param cacheKey Optional cache key to target for optimistic update
*/
async function updateAssetTags(
asset: AssetItem,
newTags: string[],
cacheKey?: string
) {
const originalTags = asset.tags
const tagsToAdd = difference(newTags, originalTags)
const tagsToRemove = difference(originalTags, newTags)
if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return
updateAssetInCache(asset.id, { tags: newTags }, cacheKey)
let removedTagsOnServer: string[] = []
try {
let removeResult: TagsOperationResult | undefined
if (tagsToRemove.length > 0) {
removeResult = await assetService.removeAssetTags(
asset.id,
tagsToRemove
)
removedTagsOnServer = removeResult.removed ?? tagsToRemove
}
const addResult =
tagsToAdd.length > 0
? await assetService.addAssetTags(asset.id, tagsToAdd)
: undefined
const finalTags = (addResult ?? removeResult)?.total_tags
if (finalTags) {
updateAssetInCache(asset.id, { tags: finalTags }, cacheKey)
}
} catch (error) {
console.error('Failed to update asset tags:', error)
updateAssetInCache(asset.id, { tags: originalTags }, cacheKey)
if (removedTagsOnServer.length > 0) {
try {
await assetService.addAssetTags(asset.id, removedTagsOnServer)
} catch (compensationError) {
console.error(
'Failed to restore tags after partial failure; invalidating cache to force refetch:',
compensationError
)
const categoriesToInvalidate = new Set<string>()
const resolved = cacheKey ? resolveCategory(cacheKey) : undefined
if (resolved) {
categoriesToInvalidate.add(resolved)
}
for (const [
category,
state
] of modelStateByCategory.value.entries()) {
if (state.assets?.has(asset.id)) {
categoriesToInvalidate.add(category)
}
}
for (const category of categoriesToInvalidate) {
invalidateCategory(category)
}
}
}
}
}
/**
* Invalidate model caches for a given category (e.g., 'checkpoints', 'loras')
* Clears the category cache and tag-based caches so next access triggers refetch
* @param category The model category to invalidate (e.g., 'checkpoints')
*/
function invalidateModelsForCategory(category: string): void {
invalidateCategory(category)
invalidateCategory(`tag:${category}`)
invalidateCategory('tag:models')
}
return {
getAssets,
isLoading,
getError,
hasMore,
hasAssetKey,
hasCategory,
updateModelsForNodeType,
updateModelsForTag,
invalidateCategory,
updateAssetMetadata,
updateAssetTags,
invalidateModelsForCategory
}
}
const emptyAssets: AssetItem[] = []
/**
* Update asset metadata with optimistic cache update
* @param asset The asset to update
* @param userMetadata The user_metadata to save
* @param cacheKey Optional cache key to target for optimistic update
*/
async function updateAssetMetadata(
asset: AssetItem,
userMetadata: Record<string, unknown>,
cacheKey?: string
) {
const originalMetadata = asset.user_metadata
updateAssetInCache(asset.id, { user_metadata: userMetadata }, cacheKey)
try {
const updatedAsset = await assetService.updateAsset(asset.id, {
user_metadata: userMetadata
})
updateAssetInCache(asset.id, updatedAsset, cacheKey)
} catch (error) {
console.error('Failed to update asset metadata:', error)
updateAssetInCache(
asset.id,
{ user_metadata: originalMetadata },
cacheKey
)
}
}
/**
* Update asset tags using add/remove endpoints
* @param asset The asset to update (used to read current tags)
* @param newTags The desired tags array
* @param cacheKey Optional cache key to target for optimistic update
*/
async function updateAssetTags(
asset: AssetItem,
newTags: string[],
cacheKey?: string
) {
const originalTags = asset.tags
const tagsToAdd = difference(newTags, originalTags)
const tagsToRemove = difference(originalTags, newTags)
if (tagsToAdd.length === 0 && tagsToRemove.length === 0) return
updateAssetInCache(asset.id, { tags: newTags }, cacheKey)
let removedTagsOnServer: string[] = []
try {
let removeResult: TagsOperationResult | undefined
if (tagsToRemove.length > 0) {
removeResult = await assetService.removeAssetTags(
asset.id,
tagsToRemove
)
removedTagsOnServer = removeResult.removed ?? tagsToRemove
}
const addResult =
tagsToAdd.length > 0
? await assetService.addAssetTags(asset.id, tagsToAdd)
: undefined
const finalTags = (addResult ?? removeResult)?.total_tags
if (finalTags) {
updateAssetInCache(asset.id, { tags: finalTags }, cacheKey)
}
} catch (error) {
console.error('Failed to update asset tags:', error)
updateAssetInCache(asset.id, { tags: originalTags }, cacheKey)
if (removedTagsOnServer.length > 0) {
try {
await assetService.addAssetTags(asset.id, removedTagsOnServer)
} catch (compensationError) {
console.error(
'Failed to restore tags after partial failure; invalidating cache to force refetch:',
compensationError
)
const categoriesToInvalidate = new Set<string>()
const resolved = cacheKey ? resolveCategory(cacheKey) : undefined
if (resolved) {
categoriesToInvalidate.add(resolved)
}
for (const [
category,
state
] of modelStateByCategory.value.entries()) {
if (state.assets?.has(asset.id)) {
categoriesToInvalidate.add(category)
}
}
for (const category of categoriesToInvalidate) {
invalidateCategory(category)
}
}
}
}
}
/**
* Invalidate model caches for a given category (e.g., 'checkpoints', 'loras')
* Clears the category cache and tag-based caches so next access triggers refetch
* @param category The model category to invalidate (e.g., 'checkpoints')
*/
function invalidateModelsForCategory(category: string): void {
invalidateCategory(category)
invalidateCategory(`tag:${category}`)
invalidateCategory('tag:models')
}
return {
getAssets: () => emptyAssets,
isLoading: () => false,
getError: () => undefined,
hasMore: () => false,
hasAssetKey: () => false,
hasCategory: () => false,
updateModelsForNodeType: async () => {},
invalidateCategory: () => {},
updateModelsForTag: async () => {},
updateAssetMetadata: async () => {},
updateAssetTags: async () => {},
invalidateModelsForCategory: () => {}
getAssets,
isLoading,
getError,
hasMore,
hasAssetKey,
hasCategory,
updateModelsForNodeType,
updateModelsForTag,
invalidateCategory,
updateAssetMetadata,
updateAssetTags,
invalidateModelsForCategory
}
}

View File

@@ -5,7 +5,12 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
import { assetService } from '@/platform/assets/services/assetService'
import { useSettingStore } from '@/platform/settings/settingStore'
import { api } from '@/scripts/api'
import { useModelStore } from '@/stores/modelStore'
import {
ResourceState,
effectiveModelExtensions,
matchesModelExtension,
useModelStore
} from '@/stores/modelStore'
// Mock the api
vi.mock('@/scripts/api', () => ({
@@ -15,6 +20,7 @@ vi.mock('@/scripts/api', () => ({
viewMetadata: vi.fn(),
apiURL: vi.fn((path: string) => `http://localhost:8188${path}`),
addEventListener: vi.fn(),
addCustomEventListener: vi.fn(),
removeEventListener: vi.fn()
}
}))
@@ -22,8 +28,10 @@ vi.mock('@/scripts/api', () => ({
// Mock the assetService
vi.mock('@/platform/assets/services/assetService', () => ({
assetService: {
getAssetModelFolders: vi.fn(),
getAssetModels: vi.fn()
getAssetModels: vi.fn(),
invalidateModelBuckets: vi.fn(),
onModelsScanned: vi.fn(),
seedModelAssets: vi.fn()
}
}))
@@ -57,16 +65,15 @@ function enableMocks(useAssetAPI = false) {
{ name: 'vae', folders: ['/path/to/vae'] }
])
// Mock asset API - also returns objects with name and folders properties
vi.mocked(assetService.getAssetModelFolders).mockResolvedValue([
{ name: 'checkpoints', folders: ['/path/to/checkpoints'] },
{ name: 'vae', folders: ['/path/to/vae'] }
])
// Asset API supplies only the per-folder model contents; folders come from
// api.getModelFolders in both paths.
vi.mocked(assetService.getAssetModels).mockResolvedValue([
{ name: 'sdxl.safetensors', pathIndex: 0 },
{ name: 'sdv15.safetensors', pathIndex: 0 },
{ name: 'noinfo.safetensors', pathIndex: 0 }
])
vi.mocked(assetService.seedModelAssets).mockResolvedValue(undefined)
vi.mocked(assetService.onModelsScanned).mockReturnValue(() => {})
vi.mocked(api.viewMetadata).mockImplementation((_, model) => {
if (model === 'noinfo.safetensors') {
@@ -209,6 +216,193 @@ describe('useModelStore', () => {
expect(api.getModelFolders).toHaveBeenCalledTimes(2)
expect(api.getModels).not.toHaveBeenCalled()
})
it('kicks off a backend scan when models come from the asset API', async () => {
enableMocks(true)
store = useModelStore()
await store.refresh()
expect(assetService.seedModelAssets).toHaveBeenCalledTimes(1)
})
it('does not scan on the legacy listing path', async () => {
enableMocks(false)
store = useModelStore()
await store.refresh()
expect(assetService.seedModelAssets).not.toHaveBeenCalled()
})
})
describe('concurrent folder loads', () => {
it('does not let a stale folder response overwrite a fresher one', async () => {
enableMocks()
let resolveStale!: (value: { name: string; folders: string[] }[]) => void
vi.mocked(api.getModelFolders).mockReturnValueOnce(
new Promise((resolve) => {
resolveStale = resolve
})
)
store = useModelStore()
const staleLoad = store.loadModelFolders()
vi.mocked(api.getModelFolders).mockResolvedValueOnce([
{ name: 'fresh-folder', folders: ['/fresh'] }
])
await store.loadModelFolders()
expect(store.modelFolders.map((f) => f.directory)).toEqual([
'fresh-folder'
])
resolveStale([{ name: 'stale-folder', folders: ['/stale'] }])
await staleLoad
expect(store.modelFolders.map((f) => f.directory)).toEqual([
'fresh-folder'
])
})
})
it('eager-loading before boot loads the folder structure first', async () => {
enableMocks()
store = useModelStore()
await store.loadModels()
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
expect(api.getModels).toHaveBeenCalledWith('checkpoints')
expect(api.getModels).toHaveBeenCalledWith('vae')
})
describe('refreshModelFolder races', () => {
it('keeps the newer refresh when an older one for the same folder finishes last', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
await store.getLoadedModelFolder('checkpoints')
let resolveOld!: (value: { name: string; pathIndex: number }[]) => void
vi.mocked(api.getModels).mockReturnValueOnce(
new Promise((resolve) => {
resolveOld = resolve
})
)
const oldRefresh = store.refreshModelFolder('checkpoints')
vi.mocked(api.getModels).mockResolvedValueOnce([
{ name: 'newer.safetensors', pathIndex: 0 }
])
await store.refreshModelFolder('checkpoints')
resolveOld([{ name: 'older.safetensors', pathIndex: 0 }])
await oldRefresh
const folder = await store.getLoadedModelFolder('checkpoints')
expect(folder!.models['0/newer.safetensors']).toBeDefined()
expect(folder!.models['0/older.safetensors']).toBeUndefined()
})
it('does not resurrect a stale folder over a fresher structure', async () => {
enableMocks()
store = useModelStore()
await store.loadModelFolders()
await store.getLoadedModelFolder('checkpoints')
let resolveStaleContents!: (
value: { name: string; pathIndex: number }[]
) => void
vi.mocked(api.getModels).mockReturnValueOnce(
new Promise((resolve) => {
resolveStaleContents = resolve
})
)
const staleRefresh = store.refreshModelFolder('checkpoints')
// A full reload rebuilds the folder structure mid-refresh.
await store.loadModelFolders()
const freshFolder = await store.getLoadedModelFolder('checkpoints')
resolveStaleContents([{ name: 'stale.safetensors', pathIndex: 0 }])
await staleRefresh
const current = await store.getLoadedModelFolder('checkpoints')
expect(current).toBe(freshFolder)
expect(current!.models['0/stale.safetensors']).toBeUndefined()
})
})
describe('scan fast-phase completion', () => {
it('re-loads folders whose eager load was still in flight when the reload fired', async () => {
enableMocks(true)
store = useModelStore()
await store.loadModelFolders()
// Eager load starts but its fetch never lands before the scan event.
let resolveEager!: (value: { name: string; pathIndex: number }[]) => void
vi.mocked(assetService.getAssetModels).mockReturnValueOnce(
new Promise((resolve) => {
resolveEager = resolve
})
)
const eagerLoad = store.getLoadedModelFolder('checkpoints')
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
.calls[0]?.[0]
await scanCallback!()
// The rebuilt folder must have been re-loaded, not left uninitialized
// while the original request finishes into a detached folder object.
const folder = store.modelFolders.find(
(f) => f.directory === 'checkpoints'
)
expect(folder!.state).toBe(ResourceState.Loaded)
resolveEager([{ name: 'detached.safetensors', pathIndex: 0 }])
await eagerLoad
const current = await store.getLoadedModelFolder('checkpoints')
expect(current!.models['0/detached.safetensors']).toBeUndefined()
})
it('re-loads previously loaded folders when the event fires', async () => {
enableMocks(true)
store = useModelStore()
await store.loadModelFolders()
await store.getLoadedModelFolder('checkpoints')
expect(assetService.getAssetModels).toHaveBeenCalledTimes(1)
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
.calls[0]?.[0]
expect(scanCallback).toBeDefined()
await scanCallback!()
await vi.waitFor(() => {
expect(assetService.getAssetModels).toHaveBeenCalledTimes(2)
})
expect(assetService.invalidateModelBuckets).toHaveBeenCalled()
expect(assetService.seedModelAssets).not.toHaveBeenCalled()
})
it('logs instead of rejecting when the post-scan reload fails', async () => {
const error = vi.spyOn(console, 'error').mockImplementation(() => {})
enableMocks(true)
vi.mocked(api.getModelFolders).mockRejectedValue(
new Error('transient network failure')
)
store = useModelStore()
const scanCallback = vi.mocked(assetService.onModelsScanned).mock
.calls[0]?.[0]
await scanCallback!()
await vi.waitFor(() => {
expect(error).toHaveBeenCalledWith(
expect.stringContaining('reload'),
expect.any(Error)
)
})
error.mockRestore()
})
})
describe('API switching functionality', () => {
@@ -218,28 +412,117 @@ describe('useModelStore', () => {
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
// Folders come from /experiment/models; legacy path also serves models.
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
expect(api.getModels).toHaveBeenCalledWith('checkpoints')
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(0)
expect(assetService.getAssetModels).toHaveBeenCalledTimes(0)
expect(folderStore).toBeDefined()
expect(Object.keys(folderStore!.models)).toHaveLength(3)
})
it('should use asset API for complete workflow when UseAssetAPI setting is true', async () => {
it('should use asset API for model contents but /experiment/models for folders when UseAssetAPI is true', async () => {
enableMocks(true) // useAssetAPI = true
store = useModelStore()
await store.loadModelFolders()
const folderStore = await store.getLoadedModelFolder('checkpoints')
// Both APIs return objects with .name property, modelStore extracts folder.name in both cases
expect(assetService.getAssetModelFolders).toHaveBeenCalledTimes(1)
// Folders always come from /experiment/models; only contents use the asset API.
expect(api.getModelFolders).toHaveBeenCalledTimes(1)
expect(assetService.getAssetModels).toHaveBeenCalledWith('checkpoints')
expect(api.getModelFolders).toHaveBeenCalledTimes(0)
expect(api.getModels).toHaveBeenCalledTimes(0)
expect(folderStore).toBeDefined()
expect(Object.keys(folderStore!.models)).toHaveLength(3)
})
it('filters asset-path folder contents by the folder extensions', async () => {
enableMocks(true)
vi.mocked(api.getModelFolders).mockResolvedValue([
{ name: 'checkpoints', folders: ['/p'], extensions: ['.safetensors'] }
])
vi.mocked(assetService.getAssetModels).mockResolvedValue([
{ name: 'keep.safetensors', pathIndex: 0 },
{ name: 'notes.txt', pathIndex: 0 }
])
store = useModelStore()
await store.loadModelFolders()
const folder = await store.getLoadedModelFolder('checkpoints')
const names = Object.values(folder!.models).map((m) => m.file_name)
expect(names).toEqual(['keep.safetensors'])
})
it('hides non-model noise in match-all folders on the asset path', async () => {
enableMocks(true)
vi.mocked(api.getModelFolders).mockResolvedValue([
{ name: 'LLM', folders: ['/p'], extensions: [] }
])
vi.mocked(assetService.getAssetModels).mockResolvedValue([
{ name: 'model.safetensors', pathIndex: 0 },
{ name: 'README.md', pathIndex: 0 }
])
store = useModelStore()
await store.loadModelFolders()
const folder = await store.getLoadedModelFolder('LLM')
const names = Object.values(folder!.models).map((m) => m.file_name)
expect(names).toEqual(['model.safetensors'])
})
it('leaves the legacy listing unfiltered', async () => {
enableMocks(false)
vi.mocked(api.getModelFolders).mockResolvedValue([
{ name: 'checkpoints', folders: ['/p'], extensions: ['.safetensors'] }
])
vi.mocked(api.getModels).mockResolvedValue([
{ name: 'keep.safetensors', pathIndex: 0 },
{ name: 'legacy-visible.gguf', pathIndex: 0 }
])
store = useModelStore()
await store.loadModelFolders()
const folder = await store.getLoadedModelFolder('checkpoints')
const names = Object.values(folder!.models).map((m) => m.file_name)
expect(names).toEqual(['keep.safetensors', 'legacy-visible.gguf'])
})
})
})
describe(matchesModelExtension, () => {
it('keeps files whose extension is in the folder list', () => {
expect(
matchesModelExtension('a.safetensors', ['.safetensors', '.ckpt'])
).toBe(true)
expect(matchesModelExtension('a.txt', ['.safetensors'])).toBe(false)
})
it('matches case-insensitively and on subpaths', () => {
expect(
matchesModelExtension('sub/dir/A.SAFETENSORS', ['.safetensors'])
).toBe(true)
})
it('is permissive when there are no real extensions', () => {
// Unfiltered folders (empty) and the `folder`/`''` sentinels show everything.
expect(matchesModelExtension('readme.md', [])).toBe(true)
expect(matchesModelExtension('anything', ['folder'])).toBe(true)
expect(matchesModelExtension('anything', [''])).toBe(true)
})
})
describe(effectiveModelExtensions, () => {
it('uses a registered allowlist verbatim', () => {
expect(effectiveModelExtensions(['.gguf'])).toEqual(['.gguf'])
})
it('substitutes the default list for match-all folders', () => {
const effective = effectiveModelExtensions([])
expect(effective).toContain('.safetensors')
expect(matchesModelExtension('readme.md', effective)).toBe(false)
})
it('treats an absent field (older backends) like match-all', () => {
expect(effectiveModelExtensions(undefined)).toEqual(
effectiveModelExtensions([])
)
})
})

View File

@@ -1,8 +1,9 @@
import { defineStore } from 'pinia'
import { computed, ref } from 'vue'
import { computed, onScopeDispose, ref } from 'vue'
import type { ModelFile } from '@/platform/assets/schemas/assetSchema'
import { assetService } from '@/platform/assets/services/assetService'
import { isCloud } from '@/platform/distribution/types'
import { useSettingStore } from '@/platform/settings/settingStore'
import { api } from '@/scripts/api'
@@ -100,6 +101,11 @@ export class ComfyModelDef {
if (this.has_loaded_metadata || this.is_load_requested) {
return
}
// viewMetadata reads the safetensors header off local disk; on Cloud the
// model bytes live in object storage so there is nothing to read.
if (isCloud) {
return
}
this.is_load_requested = true
try {
const metadata = await api.viewMetadata(this.directory, this.file_name)
@@ -156,6 +162,66 @@ export enum ResourceState {
Loaded
}
/**
* Resolves the preview image for a model: embedded metadata thumbnail when
* loaded, otherwise the server-rendered `.webp` preview. The preview endpoint
* reads a rendered thumbnail off local disk, which is unavailable on Cloud
* (model bytes live in object storage), so Cloud resolves to no preview.
*/
export function getModelPreviewUrl(model: ComfyModelDef): string {
if (model.image) return model.image
if (isCloud) return ''
const extension = model.file_name.split('.').pop()
const filename = model.file_name.replace(`.${extension}`, '.webp')
const encodedFilename = encodeURIComponent(filename).replace(/%2F/g, '/')
return `/api/experiment/models/preview/${model.directory}/${model.path_index}/${encodedFilename}`
}
/**
* FE-owned copy of core's default `supported_pt_extensions`, applied to
* match-all folders (empty registered allowlist) so they don't surface
* README/config noise. Accepted to go stale across core version bumps; the
* whole surface is expected to be short-lived.
*/
const DEFAULT_MODEL_EXTENSIONS = [
'.ckpt',
'.pt',
'.pt2',
'.bin',
'.pth',
'.safetensors',
'.pkl',
'.sft'
]
/**
* Resolves a folder's display allowlist from its raw registered `extensions`
* (`/experiment/models`): non-empty is used verbatim; an empty array
* (match-all) or an absent field (older backends) takes the FE default list,
* reproducing the legacy sidebar's global-set behavior so nothing that used
* to be hidden starts showing.
*/
export function effectiveModelExtensions(
extensions: string[] | undefined
): string[] {
return extensions?.length ? extensions : DEFAULT_MODEL_EXTENSIONS
}
/**
* Whether a model file belongs in a folder given its display allowlist. An
* empty list, or a list with no real (`.`-prefixed) extensions (the
* `'folder'`/`''` sentinels), leaves the folder unfiltered.
*/
export function matchesModelExtension(
fileName: string,
extensions: string[]
): boolean {
const realExtensions = extensions.filter((ext) => ext.startsWith('.'))
if (realExtensions.length === 0) return true
const lower = fileName.toLowerCase()
return realExtensions.some((ext) => lower.endsWith(ext.toLowerCase()))
}
export class ModelFolder {
/** Models in this folder */
models: Record<string, ComfyModelDef> = {}
@@ -163,7 +229,8 @@ export class ModelFolder {
constructor(
public directory: string,
private getModelsFunc: (folder: string) => Promise<ModelFile[]>
private getModelsFunc: (folder: string) => Promise<ModelFile[]>,
public readonly extensions: string[] = []
) {}
get key(): string {
@@ -180,6 +247,7 @@ export class ModelFolder {
this.state = ResourceState.Loading
const models = await this.getModelsFunc(this.directory)
for (const model of models) {
if (!matchesModelExtension(model.name, this.extensions)) continue
this.models[`${model.pathIndex}/${model.name}`] = new ComfyModelDef(
model.name,
this.directory,
@@ -212,22 +280,33 @@ export const useModelStore = defineStore('models', () => {
: (folder) => api.getModels(folder)
}
let modelFoldersRequestId = 0
/**
* Loads the model folders from the server
* Loads the model folders from the server.
*
* The folder list (and its registration order) always comes from
* `/experiment/models`, the source of truth for which model folders exist;
* only the per-folder contents differ between the asset API and legacy paths.
* Concurrent loads (manual refresh racing the scan-complete reload) commit
* only the newest request so a slow stale response cannot overwrite a
* fresher folder structure.
*/
async function loadModelFolders() {
const useAssetAPI: boolean = settingStore.get('Comfy.Assets.UseAssetAPI')
const resData = useAssetAPI
? await assetService.getAssetModelFolders()
: await api.getModelFolders()
const requestId = ++modelFoldersRequestId
const resData = await api.getModelFolders()
if (requestId !== modelFoldersRequestId) return
modelFolderNames.value = resData.map((folder) => folder.name)
modelFolderByName.value = {}
const useAssetAPI: boolean = settingStore.get('Comfy.Assets.UseAssetAPI')
const getModelsFunc = createGetModelsFunc()
for (const folderName of modelFolderNames.value) {
modelFolderByName.value[folderName] = new ModelFolder(
folderName,
getModelsFunc
for (const folder of resData) {
modelFolderByName.value[folder.name] = new ModelFolder(
folder.name,
getModelsFunc,
// Display filtering applies to the asset walk only; the legacy
// listing keeps its historical server-side (global-set) filtering.
useAssetAPI ? effectiveModelExtensions(folder.extensions) : []
)
}
}
@@ -240,9 +319,15 @@ export const useModelStore = defineStore('models', () => {
}
/**
* Loads all model folders' contents from the server
* Loads all model folders' contents from the server. Loads the folder
* structure first when it has not arrived yet — eager loading can run
* before app boot's own loadModelFolders call resolves, and iterating an
* empty folder list would silently load nothing.
*/
async function loadModels() {
if (modelFolderNames.value.length === 0) {
await loadModelFolders()
}
return Promise.all(modelFolders.value.map((folder) => folder.load()))
}
@@ -253,25 +338,45 @@ export const useModelStore = defineStore('models', () => {
* a newly-introduced folder type is picked up without dropping other
* folders' loaded contents.
*/
const folderRefreshIds = new Map<string, number>()
async function refreshModelFolder(folderName: string) {
assetService.invalidateModelBuckets()
if (!(folderName in modelFolderByName.value)) {
await refresh()
return
}
const folder = new ModelFolder(folderName, createGetModelsFunc())
const requestId = modelFoldersRequestId
const refreshId = (folderRefreshIds.get(folderName) ?? 0) + 1
folderRefreshIds.set(folderName, refreshId)
const folder = new ModelFolder(
folderName,
createGetModelsFunc(),
modelFolderByName.value[folderName].extensions
)
await folder.load()
// A full reload may have rebuilt the folder structure while this folder
// refreshed, and a newer refresh of the same folder may have already
// committed; committing then would resurrect a stale folder object.
if (requestId !== modelFoldersRequestId) return
if (folderRefreshIds.get(folderName) !== refreshId) return
modelFolderByName.value[folderName] = folder
}
/**
* Refreshes the folder structure and re-loads any folder whose contents
* had previously been loaded. Used by manual refresh actions ("r" key,
* sidebar refresh button) to pick up on-disk changes without losing the
* currently-visible contents.
* Re-fetches the folder structure and re-loads any folder whose contents
* had previously been loaded, picking up server-side changes without
* losing the currently-visible contents.
*/
async function refresh() {
async function reloadModels() {
assetService.invalidateModelBuckets()
// Loading counts as previously loaded: a scan-complete reload can land
// while the eager load is still in flight, and replacing those folder
// objects without re-loading them would strand the sidebar on
// uninitialized folders whose original loads finish into detached
// objects.
const previouslyLoaded = modelFolders.value
.filter((folder) => folder.state === ResourceState.Loaded)
.filter((folder) => folder.state !== ResourceState.Uninitialized)
.map((folder) => folder.directory)
await loadModelFolders()
await Promise.all(
@@ -281,6 +386,43 @@ export const useModelStore = defineStore('models', () => {
)
}
/**
* Asks the backend to rescan the model roots so files added on disk since
* startup become assets. Skipped on Cloud (models are ingested via uploads,
* not scanned from disk) and on the legacy listing path (which reads the
* filesystem live on every request).
*/
async function requestModelScan() {
if (isCloud) return
if (!settingStore.get('Comfy.Assets.UseAssetAPI')) return
try {
await assetService.seedModelAssets()
} catch (error) {
console.warn('Unable to start model asset scan', error)
}
}
/**
* Manual refresh ("r" key, sidebar refresh button): kicks off a backend
* rescan and immediately re-loads the currently known server state; the
* scan completion subscription below re-loads again with whatever the
* scan discovered. The scan is deliberately not awaited so it runs
* concurrently with the reload.
*/
async function refresh() {
void requestModelScan()
await reloadModels()
}
const unsubscribeModelsScanned = assetService.onModelsScanned(async () => {
try {
await reloadModels()
} catch (error) {
console.error('Failed to reload the model library after a scan', error)
}
})
onScopeDispose(unsubscribeModelsScanned)
return {
models,
modelFolders,

View File

@@ -5,12 +5,24 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
import { useSidebarTabStore } from '@/stores/workspace/sidebarTabStore'
const { mockGetSetting, mockRegisterCommand, mockRegisterCommands } =
vi.hoisted(() => ({
const {
mockGetSetting,
mockRegisterCommand,
mockRegisterCommands,
mockBrowseModelAssets,
registeredCommands,
commandStoreCommands
} = vi.hoisted(() => {
const registeredCommands: { id: string; function: () => unknown }[] = []
return {
mockGetSetting: vi.fn(),
mockRegisterCommand: vi.fn(),
mockRegisterCommands: vi.fn()
}))
mockRegisterCommand: vi.fn((command) => registeredCommands.push(command)),
mockRegisterCommands: vi.fn(),
mockBrowseModelAssets: vi.fn(),
registeredCommands,
commandStoreCommands: [] as { id: string; function: () => unknown }[]
}
})
vi.mock('@/platform/settings/settingStore', () => ({
useSettingStore: () => ({
@@ -21,7 +33,7 @@ vi.mock('@/platform/settings/settingStore', () => ({
vi.mock('@/stores/commandStore', () => ({
useCommandStore: () => ({
registerCommand: mockRegisterCommand,
commands: []
commands: commandStoreCommands
})
}))
@@ -99,8 +111,18 @@ describe('useSidebarTabStore', () => {
mockGetSetting.mockReset()
mockRegisterCommand.mockClear()
mockRegisterCommands.mockClear()
mockBrowseModelAssets.mockClear()
registeredCommands.length = 0
commandStoreCommands.length = 0
})
const toggleModelLibrary = async () => {
const toggleCommand = registeredCommands.find(
(command) => command.id === 'Workspace.ToggleSidebarTab.model-library'
)
await toggleCommand?.function()
}
it('registers the job history tab when QPO V2 is enabled', () => {
mockGetSetting.mockImplementation((key: string) =>
key === 'Comfy.Queue.QPOV2' ? true : undefined
@@ -160,4 +182,63 @@ describe('useSidebarTabStore', () => {
])
expect(mockRegisterCommand).toHaveBeenCalledTimes(6)
})
describe('model library view selection', () => {
it('toggles the sidebar tab when the asset view is disabled', async () => {
mockGetSetting.mockImplementation((key: string) =>
key === 'Comfy.ModelLibrary.UseAssetBrowser' ? false : undefined
)
commandStoreCommands.push({
id: 'Comfy.BrowseModelAssets',
function: mockBrowseModelAssets
})
const store = useSidebarTabStore()
store.registerCoreSidebarTabs()
await toggleModelLibrary()
expect(store.activeSidebarTabId).toBe('model-library')
expect(mockBrowseModelAssets).not.toHaveBeenCalled()
})
it('opens the asset browser when the browser and asset API are enabled', async () => {
mockGetSetting.mockImplementation((key: string) =>
key === 'Comfy.ModelLibrary.UseAssetBrowser' ||
key === 'Comfy.Assets.UseAssetAPI'
? true
: undefined
)
commandStoreCommands.push({
id: 'Comfy.BrowseModelAssets',
function: mockBrowseModelAssets
})
const store = useSidebarTabStore()
store.registerCoreSidebarTabs()
await toggleModelLibrary()
expect(mockBrowseModelAssets).toHaveBeenCalledOnce()
expect(store.activeSidebarTabId).toBeNull()
})
it('falls back to the sidebar tree when the asset API is disabled', async () => {
mockGetSetting.mockImplementation((key: string) =>
key === 'Comfy.ModelLibrary.UseAssetBrowser' ? true : false
)
commandStoreCommands.push({
id: 'Comfy.BrowseModelAssets',
function: mockBrowseModelAssets
})
const store = useSidebarTabStore()
store.registerCoreSidebarTabs()
await toggleModelLibrary()
expect(store.activeSidebarTabId).toBe('model-library')
expect(mockBrowseModelAssets).not.toHaveBeenCalled()
})
})
})

View File

@@ -76,8 +76,13 @@ export const useSidebarTabStore = defineStore('sidebarTab', () => {
const settingStore = useSettingStore()
const commandStore = useCommandStore()
// The asset browser cannot function without the asset API, so the
// browser routing derives from both settings: with the API disabled
// the browser setting is inert and the tab always opens the sidebar
// tree, rather than prompt-correcting the combination.
if (
tab.id === 'model-library' &&
settingStore.get('Comfy.ModelLibrary.UseAssetBrowser') &&
settingStore.get('Comfy.Assets.UseAssetAPI')
) {
await commandStore.commands