From 501a7e49e536e44e5412371db2459a93c600b6de Mon Sep 17 00:00:00 2001 From: bymyself Date: Sun, 17 Aug 2025 09:31:13 -0700 Subject: [PATCH] feat: Add preInit lifecycle hook for early extension initialization - Add preInit hook to ComfyExtension interface that runs before canvas creation - Implement preInit invocation in app.ts after extension loading - Create example extension demonstrating header provider registration - Add comprehensive tests for hook execution order and error handling This enables extensions to register services and cross-cutting concerns before any other initialization occurs. Co-Authored-By: Claude --- src/examples/headerRegistrationExtension.ts | 167 ++++++++++ src/scripts/app.ts | 3 + src/types/comfy.ts | 7 + tests-ui/tests/extension/preInitHook.test.ts | 328 +++++++++++++++++++ 4 files changed, 505 insertions(+) create mode 100644 src/examples/headerRegistrationExtension.ts create mode 100644 tests-ui/tests/extension/preInitHook.test.ts diff --git a/src/examples/headerRegistrationExtension.ts b/src/examples/headerRegistrationExtension.ts new file mode 100644 index 0000000000..f15f9d20a7 --- /dev/null +++ b/src/examples/headerRegistrationExtension.ts @@ -0,0 +1,167 @@ +import { headerRegistry } from '@/services/headerRegistry' +import type { ComfyExtension } from '@/types/comfy' +import type { + HeaderMap, + HeaderProviderContext, + IHeaderProvider +} from '@/types/headerTypes' + +/** + * Example extension showing how to register header providers + * during the pre-init lifecycle hook. + * + * The pre-init hook is the earliest extension lifecycle hook, + * called before the canvas is created. This makes it perfect + * for registering cross-cutting concerns like header providers. + */ + +// Example: Authentication token provider +class AuthTokenProvider implements IHeaderProvider { + async provideHeaders(_context: HeaderProviderContext): Promise { + // This could fetch tokens from a secure store, refresh them, etc. + const token = await this.getAuthToken() + + if (token) { + return { + Authorization: `Bearer ${token}` + } + } + + return {} + } + + private async getAuthToken(): Promise { + // Example: Get token from localStorage or a secure store + // In a real implementation, this might refresh tokens, handle expiration, etc. + return localStorage.getItem('auth_token') + } +} + +// Example: API key provider for specific domains +class ApiKeyProvider implements IHeaderProvider { + private apiKeys: Record = { + 'api.example.com': 'example-api-key', + 'api.another.com': 'another-api-key' + } + + provideHeaders(context: HeaderProviderContext): HeaderMap { + const url = new URL(context.url) + const apiKey = this.apiKeys[url.hostname] + + if (apiKey) { + return { + 'X-API-Key': apiKey + } + } + + return {} + } +} + +// Example: Custom header provider for debugging +class DebugHeaderProvider implements IHeaderProvider { + provideHeaders(_context: HeaderProviderContext): HeaderMap { + if (process.env.NODE_ENV === 'development') { + return { + 'X-Debug-Mode': 'true', + 'X-Request-ID': crypto.randomUUID() + } + } + + return {} + } +} + +export const headerRegistrationExtension: ComfyExtension = { + name: 'HeaderRegistration', + + /** + * Pre-init hook - called before canvas creation. + * This is the perfect place to register header providers. + */ + async preInit(_app) { + console.log( + '[HeaderRegistration] Registering header providers in pre-init hook' + ) + + // Register auth token provider with high priority + const authRegistration = headerRegistry.registerHeaderProvider( + new AuthTokenProvider(), + { + priority: 100 + } + ) + + // Register API key provider + const apiKeyRegistration = headerRegistry.registerHeaderProvider( + new ApiKeyProvider(), + { + priority: 90 + } + ) + + // Register debug header provider with lower priority + const debugRegistration = headerRegistry.registerHeaderProvider( + new DebugHeaderProvider(), + { + priority: 10 + } + ) + + // Store registrations for potential cleanup later + // Extensions can store their data on the app instance + const extensionData = { + headerRegistrations: [ + authRegistration, + apiKeyRegistration, + debugRegistration + ] + } + + // Store a reference on the extension itself for potential cleanup + ;(headerRegistrationExtension as any).registrations = + extensionData.headerRegistrations + }, + + /** + * Standard init hook - called after canvas creation. + * At this point, header providers are already active. + */ + async init(_app) { + console.log( + '[HeaderRegistration] Headers are now being injected into all HTTP requests' + ) + }, + + /** + * Setup hook - called after app is fully initialized. + * We could add UI elements here to manage headers. + */ + async setup(_app) { + // Example: Add a command to test header injection + const { useCommandStore } = await import('@/stores/commandStore') + + useCommandStore().registerCommand({ + id: 'header-registration.test', + icon: 'pi pi-globe', + label: 'Test Header Injection', + function: async () => { + try { + // Make a test request to see headers in action + const response = await fetch('/api/test') + console.log('[HeaderRegistration] Test request completed', { + status: response.status, + headers: response.headers + }) + } catch (error) { + console.error('[HeaderRegistration] Test request failed', error) + } + } + }) + } +} + +// Extension usage: +// 1. Import this extension in your extension index +// 2. Register it with app.registerExtension(headerRegistrationExtension) +// 3. Header providers will be automatically registered before any network activity diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 88b0406dd5..c1fe20b5aa 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -799,6 +799,9 @@ export class ComfyApp { await useWorkspaceStore().workflow.syncWorkflows() await useExtensionService().loadExtensions() + // Call preInit hook before any other initialization + await useExtensionService().invokeExtensionsAsync('preInit') + this.#addProcessKeyHandler() this.#addConfigureHandler() this.#addApiUpdateHandlers() diff --git a/src/types/comfy.ts b/src/types/comfy.ts index 6227533578..85b04164c2 100644 --- a/src/types/comfy.ts +++ b/src/types/comfy.ts @@ -70,6 +70,13 @@ export interface ComfyExtension { * Badges to add to the about page */ aboutPageBadges?: AboutPageBadge[] + /** + * Allows the extension to run code before the app is initialized. This is the earliest lifecycle hook. + * Called before the canvas is created and before any other extension hooks. + * Useful for registering services, header providers, or other cross-cutting concerns. + * @param app The ComfyUI app instance + */ + preInit?(app: ComfyApp): Promise | void /** * Allows any initialisation, e.g. loading resources. Called after the canvas is created but before nodes are added * @param app The ComfyUI app instance diff --git a/tests-ui/tests/extension/preInitHook.test.ts b/tests-ui/tests/extension/preInitHook.test.ts new file mode 100644 index 0000000000..b683624363 --- /dev/null +++ b/tests-ui/tests/extension/preInitHook.test.ts @@ -0,0 +1,328 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { reactive } from 'vue' + +import { ComfyApp } from '@/scripts/app' +import type { ComfyExtension } from '@/types/comfy' + +// Create mock extension service +const mockExtensionService = { + loadExtensions: vi.fn().mockResolvedValue(undefined), + registerExtension: vi.fn(), + invokeExtensions: vi.fn(), + invokeExtensionsAsync: vi.fn().mockResolvedValue(undefined), + enabledExtensions: [] as ComfyExtension[] +} + +// Mock extension service +vi.mock('@/services/extensionService', () => ({ + useExtensionService: () => mockExtensionService +})) + +// Mock dependencies +vi.mock('@/stores/toastStore', () => ({ + useToastStore: () => ({ + add: vi.fn() + }) +})) + +vi.mock('@/stores/workspaceStore', () => ({ + useWorkspaceStore: () => ({ + workflow: { + syncWorkflows: vi.fn().mockResolvedValue(undefined) + } + }) +})) + +vi.mock('@/services/subgraphService', () => ({ + useSubgraphService: () => ({ + registerNewSubgraph: vi.fn() + }) +})) + +// Mock LiteGraph +vi.mock('@/lib/litegraph/src/litegraph', async (importOriginal) => { + const actual = (await importOriginal()) as any + return { + ...actual, + LGraph: vi.fn().mockImplementation(() => ({ + events: { + addEventListener: vi.fn() + }, + start: vi.fn(), + stop: vi.fn(), + registerNodeType: vi.fn(), + createNode: vi.fn() + })), + LGraphCanvas: vi.fn().mockImplementation((canvasEl) => ({ + state: reactive({}), + draw: vi.fn(), + canvas: canvasEl + })), + LiteGraph: { + ...actual.LiteGraph, + alt_drag_do_clone_nodes: false, + macGesturesRequireMac: true + } + } +}) + +// Mock other required methods +vi.mock('@/stores/extensionStore', () => ({ + useExtensionStore: () => ({ + disabledExtensions: new Set() + }) +})) + +vi.mock('@/utils/app.utils', () => ({ + makeUUID: vi.fn(() => 'test-uuid') +})) + +// Mock API +vi.mock('@/scripts/api', () => ({ + api: { + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + apiURL: vi.fn((path) => `/api${path}`), + connect: vi.fn(), + init: vi.fn().mockResolvedValue(undefined), + getSystemStats: vi.fn().mockResolvedValue({}), + getNodeDefs: vi.fn().mockResolvedValue({}) + } +})) + +describe('Extension Pre-Init Hook', () => { + let app: ComfyApp + let mockCanvas: HTMLCanvasElement + let callOrder: string[] + + beforeEach(() => { + vi.clearAllMocks() + callOrder = [] + + // Reset mock extension service + mockExtensionService.enabledExtensions = [] + mockExtensionService.invokeExtensionsAsync.mockReset() + mockExtensionService.invokeExtensionsAsync.mockImplementation( + async (method: keyof ComfyExtension) => { + // Call the appropriate hook on all registered extensions + for (const ext of mockExtensionService.enabledExtensions) { + const hookFn = ext[method] + if (typeof hookFn === 'function') { + try { + await hookFn.call(ext, app) + } catch (error) { + console.error(`Error in extension ${ext.name} ${method}`, error) + } + } + } + } + ) + + // Create mock canvas + mockCanvas = document.createElement('canvas') + mockCanvas.getContext = vi.fn().mockReturnValue({ + scale: vi.fn(), + save: vi.fn(), + restore: vi.fn(), + clearRect: vi.fn(), + fillRect: vi.fn(), + strokeRect: vi.fn() + }) + + // Create mock DOM elements + const createMockElement = (id: string) => { + const el = document.createElement('div') + el.id = id + document.body.appendChild(el) + return el + } + + createMockElement('comfyui-body-top') + createMockElement('comfyui-body-left') + createMockElement('comfyui-body-right') + createMockElement('comfyui-body-bottom') + createMockElement('graph-canvas-container') + + app = new ComfyApp() + // Mock app methods that are called during setup + app.registerNodes = vi.fn().mockResolvedValue(undefined) + + // Mock addEventListener for canvas element + mockCanvas.addEventListener = vi.fn() + mockCanvas.removeEventListener = vi.fn() + + // Mock window methods + window.addEventListener = vi.fn() + window.ResizeObserver = vi.fn().mockImplementation(() => ({ + observe: vi.fn(), + disconnect: vi.fn() + })) + + // Mock WebSocket + const mockWebSocket = vi.fn().mockImplementation(() => ({ + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + send: vi.fn(), + close: vi.fn(), + readyState: 1, + CONNECTING: 0, + OPEN: 1, + CLOSING: 2, + CLOSED: 3 + })) + ;(mockWebSocket as any).CONNECTING = 0 + ;(mockWebSocket as any).OPEN = 1 + ;(mockWebSocket as any).CLOSING = 2 + ;(mockWebSocket as any).CLOSED = 3 + global.WebSocket = mockWebSocket as any + }) + + afterEach(() => { + // Clean up DOM elements + document.body.innerHTML = '' + }) + + it('should call preInit hook before init hook', async () => { + const testExtension: ComfyExtension = { + name: 'TestExtension', + preInit: vi.fn(async () => { + callOrder.push('preInit') + }), + init: vi.fn(async () => { + callOrder.push('init') + }), + setup: vi.fn(async () => { + callOrder.push('setup') + }) + } + + // Register the extension + mockExtensionService.enabledExtensions.push(testExtension) + + // Run app setup + await app.setup(mockCanvas) + + // Verify all hooks were called + expect(testExtension.preInit).toHaveBeenCalledWith(app) + expect(testExtension.init).toHaveBeenCalledWith(app) + expect(testExtension.setup).toHaveBeenCalledWith(app) + + // Verify correct order + expect(callOrder).toEqual(['preInit', 'init', 'setup']) + }) + + it('should call preInit before canvas creation', async () => { + const events: string[] = [] + + const testExtension: ComfyExtension = { + name: 'CanvasTestExtension', + preInit: vi.fn(async () => { + events.push('preInit') + // Canvas should not exist yet + expect(app.canvas).toBeUndefined() + }), + init: vi.fn(async () => { + events.push('init') + // Canvas should exist by init + expect(app.canvas).toBeDefined() + }) + } + + mockExtensionService.enabledExtensions.push(testExtension) + + await app.setup(mockCanvas) + + expect(events).toEqual(['preInit', 'init']) + }) + + it('should handle async preInit hooks', async () => { + const preInitComplete = vi.fn() + + const testExtension: ComfyExtension = { + name: 'AsyncExtension', + preInit: vi.fn(async () => { + // Simulate async operation + await new Promise((resolve) => setTimeout(resolve, 10)) + preInitComplete() + }), + init: vi.fn() + } + + mockExtensionService.enabledExtensions.push(testExtension) + + await app.setup(mockCanvas) + + // Ensure async preInit completed before init + expect(preInitComplete).toHaveBeenCalled() + expect(testExtension.init).toHaveBeenCalled() + + // Verify order - preInit should be called before init + const preInitCallOrder = (preInitComplete as any).mock + .invocationCallOrder[0] + const initCallOrder = (testExtension.init as any).mock + .invocationCallOrder[0] + expect(preInitCallOrder).toBeLessThan(initCallOrder) + }) + + it('should call preInit for multiple extensions in registration order', async () => { + const extension1: ComfyExtension = { + name: 'Extension1', + preInit: vi.fn(() => { + callOrder.push('ext1-preInit') + }) + } + + const extension2: ComfyExtension = { + name: 'Extension2', + preInit: vi.fn(() => { + callOrder.push('ext2-preInit') + }) + } + + const extension3: ComfyExtension = { + name: 'Extension3', + preInit: vi.fn(() => { + callOrder.push('ext3-preInit') + }) + } + + mockExtensionService.enabledExtensions.push( + extension1, + extension2, + extension3 + ) + + await app.setup(mockCanvas) + + expect(callOrder).toContain('ext1-preInit') + expect(callOrder).toContain('ext2-preInit') + expect(callOrder).toContain('ext3-preInit') + }) + + it('should handle errors in preInit gracefully', async () => { + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + + const errorExtension: ComfyExtension = { + name: 'ErrorExtension', + preInit: vi.fn(async () => { + throw new Error('PreInit error') + }), + init: vi.fn() // Should still be called + } + + mockExtensionService.enabledExtensions.push(errorExtension) + + await app.setup(mockCanvas) + + // Error should be logged + expect(consoleError).toHaveBeenCalledWith( + expect.stringContaining('Error in extension ErrorExtension'), + expect.any(Error) + ) + + // Other hooks should still be called + expect(errorExtension.init).toHaveBeenCalled() + + consoleError.mockRestore() + }) +})