fix: require explicit Hugging Face gated responses

Author: jaeone94

browser_tests/assets/missing/missing_models_gated_unsupported.json

@@ -14,7 +14,6 @@ import {
 
 const FAKE_MODEL_NAME = 'fake_model.safetensors'
 const GATED_MODEL_REPO_URL = 'https://huggingface.co/comfy-e2e/gated-test'
-const UNSUPPORTED_MODEL_NAME = 'gated_model.bin'
 
 function getModelLabel(group: Locator, modelName: string = FAKE_MODEL_NAME) {
   return group.getByRole('button', { name: modelName, exact: true })
@@ -171,7 +170,11 @@ test.describe('Errors tab - Missing models', { tag: '@ui' }, () => {
             if (route.request().method() === 'HEAD') {
               return route.fulfill({
                 status: 403,
-                headers: { 'Access-Control-Allow-Origin': '*' }
+                headers: {
+                  'Access-Control-Allow-Origin': '*',
+                  'Access-Control-Expose-Headers': 'X-Error-Code',
+                  'X-Error-Code': 'GatedRepo'
+                }
               })
             }
 
@@ -224,62 +227,5 @@ test.describe('Errors tab - Missing models', { tag: '@ui' }, () => {
         await expect(accessPage).toHaveURL(GATED_MODEL_REPO_URL)
       })
     })
-
-    test.describe('Gated model with an unsupported extension', () => {
-      test('Should not show gated guidance when no row can act on it', async ({
-        comfyPage
-      }) => {
-        let releaseGatedHead: () => void
-        const gatedHeadServed = new Promise<void>((resolve) => {
-          releaseGatedHead = resolve
-        })
-
-        await comfyPage.page
-          .context()
-          .route('https://huggingface.co/**', async (route) => {
-            if (route.request().method() !== 'HEAD') return route.abort()
-
-            if (route.request().url().endsWith('.bin')) {
-              await route.fulfill({
-                status: 403,
-                headers: { 'Access-Control-Allow-Origin': '*' }
-              })
-              releaseGatedHead()
-              return
-            }
-
-            await gatedHeadServed
-            return route.fulfill({
-              status: 200,
-              headers: {
-                'Access-Control-Allow-Origin': '*',
-                'Access-Control-Expose-Headers': 'content-length',
-                'content-length': '1024'
-              }
-            })
-          })
-        await loadWorkflowAndOpenErrorsTab(
-          comfyPage,
-          'missing/missing_models_gated_unsupported'
-        )
-
-        const modelsGroup = comfyPage.page.getByTestId(
-          TestIds.dialogs.missingModelsGroup
-        )
-        await expect(
-          getModelLabel(modelsGroup, UNSUPPORTED_MODEL_NAME)
-        ).toBeVisible()
-        await expect(
-          modelsGroup.getByText('checkpoints · 1 KB', { exact: true })
-        ).toBeVisible()
-
-        await expect(
-          comfyPage.page.getByTestId(TestIds.dialogs.missingModelGatedAccess)
-        ).toHaveCount(0)
-        await expect(
-          comfyPage.page.getByTestId(TestIds.dialogs.missingModelGatedHint)
-        ).toHaveCount(0)
-      })
-    })
   })
 })

browser_tests/tests/propertiesPanel/errorsTabMissingModels.spec.ts

@@ -1,7 +1,7 @@
 import { createPinia, setActivePinia } from 'pinia'
 import { render, screen, waitFor } from '@testing-library/vue'
 import userEvent from '@testing-library/user-event'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { nextTick } from 'vue'
 import { createI18n } from 'vue-i18n'
 
@@ -190,6 +190,10 @@ describe('MissingModelRow', () => {
     })
   })
 
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
   it('opens the model import dialog from the cloud row', async () => {
     const user = userEvent.setup()
     renderRow(makeModel([{ nodeId: '1', widgetName: 'ckpt_name' }]))
@@ -468,14 +472,20 @@ describe('MissingModelRow', () => {
     )
   })
 
-  it('delegates the gated repo action to the Desktop bridge when available', async () => {
+  it('delegates the gated repo action to the Desktop bridge with its receiver', async () => {
     mockIsCloud.value = false
     const user = userEvent.setup()
-    const openModelAccessPage = vi.fn().mockResolvedValue(true)
-    window.__comfyDesktop2 = {
-      isRemote: () => false,
-      openModelAccessPage
+    let receiver: unknown
+    let openedUrl: string | undefined
+    const bridge = {
+      isRemote: () => true,
+      async openModelAccessPage(this: unknown, url: string) {
+        receiver = this
+        openedUrl = url
+        return true
+      }
     }
+    window.__comfyDesktop2 = bridge
     const model = makeModel([{ nodeId: '1', widgetName: 'ckpt_name' }])
     model.representative.url =
       'https://huggingface.co/bfl/FLUX.1/resolve/main/model.safetensors'
@@ -490,9 +500,8 @@ describe('MissingModelRow', () => {
 
     await user.click(screen.getByTestId('missing-model-gated-access'))
 
-    expect(openModelAccessPage).toHaveBeenCalledWith(
-      'https://huggingface.co/bfl/FLUX.1'
-    )
+    expect(receiver).toBe(bridge)
+    expect(openedUrl).toBe('https://huggingface.co/bfl/FLUX.1')
     expect(mockOpenGatedRepoPage).not.toHaveBeenCalled()
   })
 
@@ -559,7 +568,6 @@ describe('MissingModelRow', () => {
       'Failed to open model access page in Desktop:',
       error
     )
-    consoleErrorSpy.mockRestore()
   })
 
   it('shows unknown category metadata for models without a directory', () => {

src/platform/missingModel/components/MissingModelRow.test.ts

@@ -93,7 +93,11 @@ describe('fetchModelMetadata', () => {
   it.for([401, 403, 451])(
     'returns gatedRepoUrl for gated HuggingFace HEAD requests (%s)',
     async (status) => {
-      fetchMock.mockResolvedValueOnce({ ok: false, status })
+      fetchMock.mockResolvedValueOnce({
+        ok: false,
+        status,
+        headers: new Headers({ 'x-error-code': 'GatedRepo' })
+      })
 
       const metadata = await fetchModelMetadata(
         `https://huggingface.co/bfl/FLUX.1/resolve/main/gated-${status}-${testId}.safetensors`
@@ -103,10 +107,14 @@ describe('fetchModelMetadata', () => {
     }
   )
 
-  it.for([404, 500])(
-    'does not treat HuggingFace %s as gated',
+  it.for([401, 403, 451, 404, 500])(
+    'does not treat HuggingFace %s as gated without the GatedRepo error code',
     async (status) => {
-      fetchMock.mockResolvedValueOnce({ ok: false, status })
+      fetchMock.mockResolvedValueOnce({
+        ok: false,
+        status,
+        headers: new Headers()
+      })
 
       const metadata = await fetchModelMetadata(
         `https://huggingface.co/org/model/resolve/main/not-gated-${status}-${testId}.safetensors`
@@ -117,7 +125,11 @@ describe('fetchModelMetadata', () => {
   )
 
   it('does not treat non-HuggingFace hosts as gated', async () => {
-    fetchMock.mockResolvedValueOnce({ ok: false, status: 403 })
+    fetchMock.mockResolvedValueOnce({
+      ok: false,
+      status: 403,
+      headers: new Headers({ 'x-error-code': 'GatedRepo' })
+    })
 
     const metadata = await fetchModelMetadata(
       `https://huggingface.co.evil.com/org/model/resolve/main/gated-${testId}.safetensors`

src/platform/missingModel/missingModelDownload.test.ts

@@ -166,6 +166,7 @@ async function fetchCivitaiMetadata(url: string): Promise<ModelMetadata> {
 }
 
 const GATED_STATUS_CODES = new Set([401, 403, 451])
+const HUGGING_FACE_GATED_ERROR_CODE = 'GatedRepo'
 
 async function fetchHeadMetadata(url: string): Promise<ModelMetadata> {
   try {
@@ -174,7 +175,8 @@ async function fetchHeadMetadata(url: string): Promise<ModelMetadata> {
     if (!response.ok) {
       if (
         isHuggingFaceRepoUrl(url) &&
-        GATED_STATUS_CODES.has(response.status)
+        GATED_STATUS_CODES.has(response.status) &&
+        response.headers.get('x-error-code') === HUGGING_FACE_GATED_ERROR_CODE
       ) {
         return { fileSize: null, gatedRepoUrl: downloadUrlToHfRepoUrl(url) }
       }