test: harden WS host override cleanup and add scheme-prefix + debug log

Move window.name reset into afterEach so it cannot leak on assertion
failure, add a named test documenting scheme-prefixed override fallback,
log when the override is active, and note the bare-host expectation in
the type JSDoc.

src/scripts/api.createSocket.test.ts

@@ -0,0 +1,110 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { api } from '@/scripts/api'
+
+describe('createSocket WebSocket host', () => {
+  let socketUrls: string[]
+
+  beforeEach(() => {
+    vi.useFakeTimers()
+    socketUrls = []
+    api.socket = null
+    api.api_host = 'localhost:8188'
+
+    vi.stubGlobal('WebSocket', function (this: WebSocket, url: string) {
+      socketUrls.push(url)
+      Object.assign(this, {
+        readyState: 1,
+        binaryType: '',
+        send: vi.fn(),
+        close: vi.fn(),
+        addEventListener: vi.fn(),
+        removeEventListener: vi.fn()
+      })
+    })
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.unstubAllGlobals()
+    delete window.__COMFY_API_WS_HOST__
+    window.name = ''
+    api.socket = null
+  })
+
+  it('dials the configured host when the override is set', () => {
+    window.__COMFY_API_WS_HOST__ = 'ws.example.com'
+
+    api.init()
+
+    expect(socketUrls).toHaveLength(1)
+    expect(socketUrls[0]).toContain('://ws.example.com')
+    expect(socketUrls[0]).not.toContain('localhost:8188')
+  })
+
+  it('preserves query parameters when the override is set', () => {
+    window.__COMFY_API_WS_HOST__ = 'ws.example.com'
+    window.name = 'session-123'
+
+    api.init()
+
+    expect(socketUrls[0]).toContain('clientId=session-123')
+  })
+
+  it('falls back to api_host when the override is unset', () => {
+    api.init()
+
+    expect(socketUrls).toHaveLength(1)
+    expect(socketUrls[0]).toContain('://localhost:8188')
+  })
+
+  it('falls back to api_host when the override is an empty string', () => {
+    window.__COMFY_API_WS_HOST__ = ''
+
+    api.init()
+
+    expect(socketUrls[0]).toContain('://localhost:8188')
+  })
+
+  it('falls back to api_host when the override is whitespace only', () => {
+    window.__COMFY_API_WS_HOST__ = '   '
+
+    api.init()
+
+    expect(socketUrls[0]).toContain('://localhost:8188')
+  })
+
+  it('ignores a prototype-pollution gadget and uses api_host', () => {
+    // @ts-expect-error - simulating a prototype-pollution attack
+    Object.prototype.__COMFY_API_WS_HOST__ = 'evil.example.com'
+    try {
+      api.init()
+
+      expect(socketUrls[0]).toContain('://localhost:8188')
+      expect(socketUrls[0]).not.toContain('evil.example.com')
+    } finally {
+      // @ts-expect-error - cleaning up the simulated attack
+      delete Object.prototype.__COMFY_API_WS_HOST__
+    }
+  })
+
+  it('falls back to polling without throwing when the host is unusable', () => {
+    vi.stubGlobal('WebSocket', function () {
+      throw new DOMException('invalid url', 'SyntaxError')
+    })
+    window.__COMFY_API_WS_HOST__ = 'https://not a host'
+
+    expect(() => api.init()).not.toThrow()
+    expect(api.socket).toBeNull()
+  })
+
+  it('falls back to polling when the override includes a scheme prefix', () => {
+    vi.stubGlobal('WebSocket', function () {
+      throw new DOMException('invalid url', 'SyntaxError')
+    })
+    window.__COMFY_API_WS_HOST__ = 'wss://ws.example.com'
+
+    expect(() => api.init()).not.toThrow()
+    expect(api.socket).toBeNull()
+  })
+})

src/scripts/api.ts

@@ -320,6 +320,28 @@ export class PromptExecutionError extends Error {
   }
 }
 
+/**
+ * Reads the deploy-time WebSocket host override. Requires an own string
+ * property so a DOM-clobbering element (`id="__COMFY_API_WS_HOST__"`) or a
+ * prototype-pollution gadget cannot redirect the authenticated socket, and
+ * treats blank values as unset so the socket falls back to `api_host`.
+ */
+function getApiWebSocketHostOverride(): string {
+  if (!Object.hasOwn(window, '__COMFY_API_WS_HOST__')) {
+    return ''
+  }
+  const override = window.__COMFY_API_WS_HOST__
+  if (typeof override !== 'string') {
+    return ''
+  }
+  const trimmed = override.trim()
+  if (!trimmed) {
+    return ''
+  }
+  console.debug('[ComfyUI] WebSocket host override active:', trimmed)
+  return trimmed
+}
+
 export class ComfyApi extends EventTarget {
   private _registered = new Set()
   /**
@@ -677,11 +699,20 @@ export class ComfyApi extends EventTarget {
     }
 
     const protocol = window.location.protocol === 'https:' ? 'wss' : 'ws'
-    const baseUrl = `${protocol}://${this.api_host}${this.api_base}/ws`
+    const wsHost = getApiWebSocketHostOverride() || this.api_host
+    const baseUrl = `${protocol}://${wsHost}${this.api_base}/ws`
     const query = params.toString()
     const wsUrl = query ? `${baseUrl}?${query}` : baseUrl
 
-    this.socket = new WebSocket(wsUrl)
+    try {
+      this.socket = new WebSocket(wsUrl)
+    } catch (error) {
+      console.error('Failed to open WebSocket connection:', error)
+      if (!isReconnect) {
+        this._pollQueue()
+      }
+      return
+    }
     this.socket.binaryType = 'arraybuffer'
 
     this.socket.addEventListener('open', () => {

src/vite-env.d.ts

@@ -15,6 +15,16 @@ declare module '~icons/*' {
 declare global {
   interface Window {
     __COMFYUI_FRONTEND_VERSION__: string
+
+    /**
+     * Overrides the host the API WebSocket dials, falling back to
+     * `location.host` when unset. Injected at deploy time by static hosts
+     * whose rewrites are HTTP-only (e.g. Vercel) and therefore cannot proxy
+     * the `/ws` upgrade, letting the socket target a WebSocket-capable host
+     * while HTTP calls keep flowing through the rewrites. Expects a bare host
+     * with no scheme (e.g. `ws.example.com`, not `wss://ws.example.com`).
+     */
+    __COMFY_API_WS_HOST__?: string
   }
 
   interface ImportMetaEnv {