diff --git a/OpenVINO/pyproject.toml b/OpenVINO/pyproject.toml index 0b206d515..e6645deba 100644 --- a/OpenVINO/pyproject.toml +++ b/OpenVINO/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "aipg-openvino-utils" version = "1.0.0" -requires-python = ">=3.10" +requires-python = "==3.12.*" dependencies = [ "openvino", ] diff --git a/WebUI/build/build-config.json b/WebUI/build/build-config.json index 77e3f323d..4d5f0dc40 100644 --- a/WebUI/build/build-config.json +++ b/WebUI/build/build-config.json @@ -5,6 +5,11 @@ "afterPack": "build/scripts/after-pack.cjs", "copyright": "Copyright © Intel. All rights reserved", "extraResources": [ + { + "from": "../backend_shared", + "to": "backend_shared", + "filter": ["*.py"] + }, { "from": "../service", "to": "service", diff --git a/WebUI/electron/main.ts b/WebUI/electron/main.ts index 43913ff49..1839a137a 100644 --- a/WebUI/electron/main.ts +++ b/WebUI/electron/main.ts @@ -1247,16 +1247,6 @@ function initEventHandle() { return pathsManager.scanAll() }) - ipcMain.handle('refreshLLMModles', (_event) => { - // Old ipexllm backend removed - return empty array - return [] - }) - - ipcMain.handle('getDownloadedLLMs', (_event) => { - // Old ipexllm backend removed - return empty array - return [] - }) - ipcMain.handle('getDownloadedGGUFLLMs', (_event) => { return pathsManager.scanGGUFLLMModels() }) diff --git a/WebUI/electron/preload.ts b/WebUI/electron/preload.ts index 4e650f7f7..db8cb46b9 100644 --- a/WebUI/electron/preload.ts +++ b/WebUI/electron/preload.ts @@ -88,11 +88,9 @@ contextBridge.exposeInMainWorld('electronAPI', { getInitSetting: () => ipcRenderer.invoke('getInitSetting'), updateModelPaths: (modelPaths: ModelPaths) => ipcRenderer.invoke('updateModelPaths', modelPaths), restorePathsSettings: () => ipcRenderer.invoke('restorePathsSettings'), - refreshLLMModles: () => ipcRenderer.invoke('refreshLLMModles'), loadModels: () => ipcRenderer.invoke('loadModels'), zoomIn: () => ipcRenderer.invoke('zoomIn'), zoomOut: () => ipcRenderer.invoke('zoomOut'), - getDownloadedLLMs: () => ipcRenderer.invoke('getDownloadedLLMs'), getDownloadedGGUFLLMs: () => ipcRenderer.invoke('getDownloadedGGUFLLMs'), getDownloadedOpenVINOLLMModels: () => ipcRenderer.invoke('getDownloadedOpenVINOLLMModels'), getDownloadedEmbeddingModels: () => ipcRenderer.invoke('getDownloadedEmbeddingModels'), diff --git a/WebUI/electron/subprocesses/llamaCppBackendService.ts b/WebUI/electron/subprocesses/llamaCppBackendService.ts index 1921821cb..be30f66c8 100644 --- a/WebUI/electron/subprocesses/llamaCppBackendService.ts +++ b/WebUI/electron/subprocesses/llamaCppBackendService.ts @@ -7,6 +7,7 @@ import { app, net, type BrowserWindow } from 'electron' import { appLoggerInstance } from '../logging/logger.ts' import { packagedResourcesRoot } from '../aipgRoot.ts' import { createEnhancedErrorDetails, type ApiService, type ErrorDetails } from './service.ts' +import { terminateProcessTree, waitForServerReadyOrThrow } from './processLifecycle.ts' import { vulkanDeviceSelectorEnv, withSelectedDevice, @@ -1341,28 +1342,10 @@ export class LlamaCppBackendService implements ApiService { private async stopLlamaLlmServer(): Promise { if (this.llamaLlmProcess) { this.appLogger.info(`Stopping LLM server for model: ${this.currentLlmModel}`, this.name) - this.llamaLlmProcess.process.kill('SIGTERM') - - // Wait a bit for graceful shutdown, then force kill if needed - await new Promise((resolve) => { - const currentProcess = this.llamaLlmProcess - const timeout = setTimeout(() => { - if (currentProcess) { - this.appLogger.warn(`Force killing LLM server process`, this.name) - currentProcess.process.kill('SIGKILL') - } - resolve() - }, 5000) - - if (currentProcess) { - currentProcess.process.on('exit', () => { - clearTimeout(timeout) - resolve() - }) - } else { - clearTimeout(timeout) - resolve() - } + await terminateProcessTree(this.llamaLlmProcess.process, { + name: this.name, + label: 'LLM', + appLogger: this.appLogger, }) this.llamaLlmProcess = null @@ -1377,28 +1360,10 @@ export class LlamaCppBackendService implements ApiService { `Stopping embedding server for model: ${this.currentEmbeddingModel}`, this.name, ) - this.llamaEmbeddingProcess.process.kill('SIGTERM') - - // Wait a bit for graceful shutdown, then force kill if needed - await new Promise((resolve) => { - const currentProcess = this.llamaEmbeddingProcess - const timeout = setTimeout(() => { - if (currentProcess) { - this.appLogger.warn(`Force killing embedding server process`, this.name) - currentProcess.process.kill('SIGKILL') - } - resolve() - }, 5000) - - if (currentProcess) { - currentProcess.process.on('exit', () => { - clearTimeout(timeout) - resolve() - }) - } else { - clearTimeout(timeout) - resolve() - } + await terminateProcessTree(this.llamaEmbeddingProcess.process, { + name: this.name, + label: 'embedding', + appLogger: this.appLogger, }) this.llamaEmbeddingProcess = null @@ -1441,58 +1406,12 @@ export class LlamaCppBackendService implements ApiService { process: ChildProcess, getStartupError?: () => string | null, ): Promise { - const maxAttempts = this.llamaCppBuildVariant === 'ssd-offload' ? 500 : 120 - const delayMs = 1000 - - for (let attempt = 0; attempt < maxAttempts; attempt++) { - // Abort early with an actionable message if the server has reported a - // fatal startup error (e.g. ran out of memory for the context size). - const startupError = getStartupError?.() - if (startupError) { - this.appLogger.error(startupError, this.name) - throw new Error(startupError) - } - - // Check if process has exited before attempting health check - if (!process || process.killed) { - this.appLogger.warn( - `Process for ${this.name} is not alive, aborting health check`, - this.name, - ) - throw new Error(`Process exited before server became ready`) - } - - try { - const response = await fetch(healthUrl, { - method: 'GET', - signal: AbortSignal.timeout(1000), - }) - - if (response.ok) { - // Double-check process is still alive before accepting success - if (!process || process.killed) { - this.appLogger.warn( - `Process for ${this.name} exited after health check succeeded, marking as failed`, - this.name, - ) - throw new Error(`Process exited after health check succeeded`) - } - this.appLogger.info(`Server ready at ${healthUrl}`, this.name) - return - } - } catch (_error) { - // Server not ready yet, continue waiting - // But check if process is still alive - if (!process || process.killed) { - this.appLogger.warn(`Process for ${this.name} exited during health check wait`, this.name) - throw new Error(`Process exited during server startup`) - } - } - - await new Promise((resolve) => setTimeout(resolve, delayMs)) - } - - throw new Error(`Server failed to start within ${(maxAttempts * delayMs) / 1000} seconds`) + await waitForServerReadyOrThrow(healthUrl, process, { + name: this.name, + maxAttempts: this.llamaCppBuildVariant === 'ssd-offload' ? 500 : 120, + getStartupError, + appLogger: this.appLogger, + }) } private async detectStorageTargets(): Promise { diff --git a/WebUI/electron/subprocesses/openVINOBackendService.ts b/WebUI/electron/subprocesses/openVINOBackendService.ts index ef1544108..03eab4a31 100644 --- a/WebUI/electron/subprocesses/openVINOBackendService.ts +++ b/WebUI/electron/subprocesses/openVINOBackendService.ts @@ -12,6 +12,10 @@ import getPort, { portNumbers } from 'get-port' import { ensureManagedPython, installBackend, uvPipInstallToTarget } from './uvBasedBackends/uv.ts' import { binary, extract } from './tools.ts' import { getBundledBackendVersionSync, resolveModels } from '../remoteUpdates.ts' +import { + terminateProcessTree as killProcessTree, + waitForServerReadyOrThrow, +} from './processLifecycle.ts' import { getMissingPackages, hasAptGet, @@ -2136,39 +2140,7 @@ export class OpenVINOBackendService implements ApiService { * whole tree; we then wait for the OS to reap it and release the handles. */ private async terminateProcessTree(proc: ChildProcess, label: string): Promise { - const waitForExit = (ms: number): Promise => - new Promise((resolve) => { - if (proc.exitCode !== null || proc.signalCode !== null) { - resolve(true) - return - } - const timeout = setTimeout(() => resolve(false), ms) - proc.once('exit', () => { - clearTimeout(timeout) - resolve(true) - }) - }) - - proc.kill('SIGTERM') - let exited = await waitForExit(2000) - if (exited) return - - this.appLogger.warn(`Force killing OVMS ${label} process tree`, this.name) - if (process.platform === 'win32' && proc.pid !== undefined) { - try { - await execAsync(`taskkill /PID ${proc.pid} /T /F`) - } catch (e) { - // taskkill exits non-zero when the process is already gone — not fatal. - this.appLogger.warn(`taskkill for OVMS ${label} reported: ${e}`, this.name) - } - } else { - proc.kill('SIGKILL') - } - - exited = await waitForExit(5000) - if (!exited) { - this.appLogger.warn(`OVMS ${label} not confirmed exited after force kill`, this.name) - } + await killProcessTree(proc, { name: this.name, label, appLogger: this.appLogger }) } private async stopOvmsLlmServer(): Promise { @@ -2676,91 +2648,12 @@ export class OpenVINOBackendService implements ApiService { childProcess: ChildProcess, maxAttempts = 120, ): Promise { - const delayMs = 1000 - - // Track whether the process has exited (process.killed only reflects - // signals sent by Node.js — not OS-level kills like OOM or SIGSEGV). - let processExited = false - let exitCode: number | null = null - let exitSignal: string | null = null - const stderrChunks: string[] = [] - - const onExit = (code: number | null, signal: string | null) => { - processExited = true - exitCode = code - exitSignal = signal - } - childProcess.on('exit', onExit) - childProcess.stderr?.on('data', (data: Buffer) => { - stderrChunks.push(data.toString()) - // Keep only the last 20 lines of stderr for diagnostics - if (stderrChunks.length > 20) stderrChunks.shift() + await waitForServerReadyOrThrow(healthUrl, childProcess, { + name: this.name, + maxAttempts, + captureExitDiagnostics: true, + appLogger: this.appLogger, }) - - const buildExitErrorMessage = (): string => { - const reason = exitSignal - ? `killed by signal ${exitSignal}` - : exitCode !== null - ? `exit code ${exitCode}` - : 'exit code null (killed by OS signal, possibly OOM)' - const lastStderr = stderrChunks.join('').trim() - const stderrSuffix = lastStderr ? `\nLast stderr output:\n${lastStderr.slice(-2000)}` : '' - return `OVMS process crashed during startup (${reason})${stderrSuffix}` - } - - try { - for (let attempt = 0; attempt < maxAttempts; attempt++) { - // Check if process has exited (covers OS kills, crashes, OOM, etc.) - if (processExited || childProcess.killed) { - const msg = buildExitErrorMessage() - this.appLogger.warn( - `Process for ${this.name} is not alive, aborting health check: ${msg}`, - this.name, - ) - throw new Error(msg) - } - - try { - const response = await fetch(healthUrl, { - method: 'GET', - signal: AbortSignal.timeout(1000), - }) - - if (response.ok) { - if (processExited || childProcess.killed) { - const msg = buildExitErrorMessage() - this.appLogger.warn( - `Process for ${this.name} exited after health check succeeded: ${msg}`, - this.name, - ) - throw new Error(msg) - } - this.appLogger.info(`Server ready at ${healthUrl}`, this.name) - return - } - } catch (error) { - // Re-throw our own errors (from the process-exit check above) - if (error instanceof Error && error.message.startsWith('OVMS process crashed')) { - throw error - } - // Server not ready yet — check if the process died while we were waiting - if (processExited || childProcess.killed) { - const msg = buildExitErrorMessage() - this.appLogger.warn( - `Process for ${this.name} exited during health check wait: ${msg}`, - this.name, - ) - throw new Error(msg) - } - } - - await new Promise((resolve) => setTimeout(resolve, delayMs)) - } - - throw new Error(`Server failed to start within ${(maxAttempts * delayMs) / 1000} seconds`) - } finally { - childProcess.removeListener('exit', onExit) - } } // Error management methods for startup failures diff --git a/WebUI/electron/subprocesses/processLifecycle.ts b/WebUI/electron/subprocesses/processLifecycle.ts new file mode 100644 index 000000000..dfcdd0fb5 --- /dev/null +++ b/WebUI/electron/subprocesses/processLifecycle.ts @@ -0,0 +1,210 @@ +import { ChildProcess, exec } from 'node:child_process' +import { promisify } from 'node:util' +import { appLoggerInstance } from '../logging/logger.ts' + +const execAsync = promisify(exec) + +// Minimal logger surface so callers can pass their bound appLogger. Defaults to +// the shared instance when omitted. +type ProcLogger = { + info(message: string, source?: string): void + warn(message: string, source?: string): void + error(message: string, source?: string): void +} + +export type TerminateProcessTreeOptions = { + /** Component name, used as the log source. */ + name: string + /** Extra label distinguishing co-located processes (e.g. 'LLM', 'embedding'). */ + label?: string + /** How long to wait for graceful SIGTERM exit before force killing. */ + gracefulMs?: number + /** How long to wait for the OS to reap the process after force kill. */ + forceMs?: number + appLogger?: ProcLogger +} + +/** + * Stop a child process and, on Windows, its entire descendant tree. + * + * SIGTERM → wait `gracefulMs` → (Windows) `taskkill /T /F`, else SIGKILL → wait + * `forceMs`. Windows `ChildProcess.kill()` only signals the direct PID, so + * descendants (uv subprocesses, embedded Python, OVMS workers) survive and keep + * handles on the service directory and binaries — which makes a subsequent + * reinstall fail to delete the directory (EPERM/EBUSY). `taskkill /T /F` tears + * down the whole tree; we then wait for the OS to reap it and release handles. + */ +export async function terminateProcessTree( + proc: ChildProcess, + opts: TerminateProcessTreeOptions, +): Promise { + const { name, label, gracefulMs = 2000, forceMs = 5000 } = opts + const appLogger = opts.appLogger ?? appLoggerInstance + const tag = label ? `${name} ${label}` : name + + const waitForExit = (ms: number): Promise => + new Promise((resolve) => { + if (proc.exitCode !== null || proc.signalCode !== null) { + resolve(true) + return + } + const timeout = setTimeout(() => resolve(false), ms) + proc.once('exit', () => { + clearTimeout(timeout) + resolve(true) + }) + }) + + // Try graceful shutdown first with SIGTERM. + proc.kill('SIGTERM') + if (await waitForExit(gracefulMs)) return + + appLogger.warn(`Force killing ${tag} process tree`, name) + if (process.platform === 'win32' && proc.pid !== undefined) { + try { + await execAsync(`taskkill /PID ${proc.pid} /T /F`) + } catch (e) { + // taskkill exits non-zero when the process is already gone — not fatal. + appLogger.warn(`taskkill for ${tag} reported: ${e}`, name) + } + } else { + proc.kill('SIGKILL') + } + + if (!(await waitForExit(forceMs))) { + appLogger.warn(`${tag} not confirmed exited after force kill`, name) + } +} + +export type WaitForServerReadyOptions = { + /** Component name, used as the log source. */ + name: string + /** Health probe attempts before giving up. */ + maxAttempts?: number + /** Delay between attempts. */ + delayMs?: number + /** Per-request timeout for the health probe. */ + requestTimeoutMs?: number + /** + * Returns an actionable error message if the server has reported a fatal + * startup error (e.g. ran out of memory), otherwise null. Checked before each + * probe so the wait aborts early with a useful message instead of timing out. + */ + getStartupError?: () => string | null + /** + * When true, capture the process exit code/signal and tail of stderr to build + * a richer crash diagnostic (covers OS-level kills like OOM/SIGSEGV that + * `ChildProcess.killed` does not reflect). + */ + captureExitDiagnostics?: boolean + appLogger?: ProcLogger +} + +/** + * Poll an HTTP health endpoint until the server responds 200, the process dies, + * or `maxAttempts` is exhausted. Throws on every failure path so the caller can + * surface an actionable error. An `'exit'` listener (not just `proc.killed`) is + * used to detect OS-level kills. + */ +export async function waitForServerReadyOrThrow( + healthUrl: string, + proc: ChildProcess, + opts: WaitForServerReadyOptions, +): Promise { + const { + name, + maxAttempts = 120, + delayMs = 1000, + requestTimeoutMs = 1000, + getStartupError, + captureExitDiagnostics = false, + } = opts + const appLogger = opts.appLogger ?? appLoggerInstance + + // process.killed only reflects signals sent by Node.js, so also track real + // exits (OOM, SIGSEGV, etc.) via the 'exit' event. + let processExited = false + let exitCode: number | null = null + let exitSignal: string | null = null + const stderrChunks: string[] = [] + + const onExit = (code: number | null, signal: string | null) => { + processExited = true + exitCode = code + exitSignal = signal + } + proc.on('exit', onExit) + + const onStderr = (data: Buffer) => { + stderrChunks.push(data.toString()) + // Keep only the last 20 chunks of stderr for diagnostics. + if (stderrChunks.length > 20) stderrChunks.shift() + } + if (captureExitDiagnostics) proc.stderr?.on('data', onStderr) + + const isDead = () => processExited || proc.killed + const deathMessage = (fallback: string): string => { + if (!captureExitDiagnostics) return fallback + const reason = exitSignal + ? `killed by signal ${exitSignal}` + : exitCode !== null + ? `exit code ${exitCode}` + : 'exit code null (killed by OS signal, possibly OOM)' + const lastStderr = stderrChunks.join('').trim() + const stderrSuffix = lastStderr ? `\nLast stderr output:\n${lastStderr.slice(-2000)}` : '' + return `${name} process crashed during startup (${reason})${stderrSuffix}` + } + + try { + for (let attempt = 0; attempt < maxAttempts; attempt++) { + // Abort early with an actionable message if the server reported a fatal + // startup error (e.g. not enough memory for the requested context size). + const startupError = getStartupError?.() + if (startupError) { + appLogger.error(startupError, name) + throw new Error(startupError) + } + + if (isDead()) { + const msg = deathMessage(`Process for ${name} exited before server became ready`) + appLogger.warn(`Process for ${name} is not alive, aborting health check: ${msg}`, name) + throw new Error(msg) + } + + let healthy = false + try { + const response = await fetch(healthUrl, { + method: 'GET', + signal: AbortSignal.timeout(requestTimeoutMs), + }) + healthy = response.ok + } catch (_error) { + // Server not up yet — fall through to the liveness check and retry. + } + + if (healthy) { + // Double-check the process is still alive before accepting success. + if (isDead()) { + const msg = deathMessage(`Process for ${name} exited after health check succeeded`) + appLogger.warn(msg, name) + throw new Error(msg) + } + appLogger.info(`Server ready at ${healthUrl}`, name) + return + } + + if (isDead()) { + const msg = deathMessage(`Process for ${name} exited during health check wait`) + appLogger.warn(msg, name) + throw new Error(msg) + } + + await new Promise((resolve) => setTimeout(resolve, delayMs)) + } + + throw new Error(`Server failed to start within ${(maxAttempts * delayMs) / 1000} seconds`) + } finally { + proc.removeListener('exit', onExit) + if (captureExitDiagnostics) proc.stderr?.removeListener('data', onStderr) + } +} diff --git a/WebUI/electron/subprocesses/service.ts b/WebUI/electron/subprocesses/service.ts index 67e6423f3..e902ad4f9 100644 --- a/WebUI/electron/subprocesses/service.ts +++ b/WebUI/electron/subprocesses/service.ts @@ -6,6 +6,7 @@ import path from 'node:path' import { appLoggerInstance } from '../logging/logger.ts' import { packagedResourcesRoot } from '../aipgRoot.ts' import { existingFileOrError, spawnProcessAsync, ProcessError } from './osProcessHelper' +import { terminateProcessTree } from './processLifecycle.ts' import { assert } from 'node:console' import { createHash } from 'crypto' @@ -716,48 +717,7 @@ export abstract class LongLivedPythonApiService implements ApiService { return 'stopped' } - const waitForExit = (ms: number): Promise => - Promise.race([ - new Promise((resolve) => proc.once('exit', () => resolve(true))), - new Promise((resolve) => setTimeout(() => resolve(false), ms)), - ]) - - // Try graceful shutdown first with SIGTERM - proc.kill('SIGTERM') - - // Wait up to 2 seconds for the process to exit gracefully - let exited = await waitForExit(2000) - - if (!exited) { - this.appLogger.warn( - `Backend ${this.name} did not exit gracefully within 2s, force killing`, - this.name, - ) - // On Windows, ChildProcess.kill() only signals the direct child and leaves - // descendant processes (e.g. ComfyUI's python and the uv subprocesses spawned - // by ComfyUI-Manager) running. Those descendants keep handles on the service - // directory — and python's CWD is the service dir itself — which makes a - // subsequent reinstall fail to delete it (EPERM). taskkill /T /F tears down - // the whole tree. - if (process.platform === 'win32' && proc.pid !== undefined) { - try { - await exec(`taskkill /PID ${proc.pid} /T /F`) - } catch (e) { - // taskkill exits non-zero when the process is already gone — not fatal. - this.appLogger.warn(`taskkill for backend ${this.name} reported: ${e}`, this.name) - } - } else { - proc.kill('SIGKILL') - } - // Wait for the OS to actually reap the process and release file handles. - exited = await waitForExit(5000) - if (!exited) { - this.appLogger.warn( - `Backend ${this.name} still not confirmed exited after force kill`, - this.name, - ) - } - } + await terminateProcessTree(proc, { name: this.name, appLogger: this.appLogger }) this.encapsulatedProcess = null this.setStatus('stopped') diff --git a/WebUI/src/assets/js/store/comfyUiMessages.ts b/WebUI/src/assets/js/store/comfyUiMessages.ts new file mode 100644 index 000000000..9beb05094 --- /dev/null +++ b/WebUI/src/assets/js/store/comfyUiMessages.ts @@ -0,0 +1,149 @@ +import { z } from 'zod' + +// Schema for the messages ComfyUI pushes over its WebSocket during execution. +export const ComfyMessageSchema = z.discriminatedUnion('type', [ + z.object({ + type: z.literal('status'), + }), + z.object({ + type: z.literal('execution_start'), + data: z.object({}).passthrough(), + }), + z.object({ + type: z.literal('execution_success'), + data: z.object({}).passthrough(), + }), + z.object({ + type: z.literal('execution_error'), + data: z + .object({ + exception_message: z.string().optional(), + exception_type: z.string().optional(), + node_id: z.union([z.string(), z.number()]).optional(), + node_type: z.string().optional(), + traceback: z.array(z.string()).optional(), + }) + .passthrough(), + }), + z.object({ + type: z.literal('execution_interrupted'), + data: z.object({}).passthrough(), + }), + z.object({ + type: z.literal('execution_cached'), + data: z.object({}).passthrough(), + }), + z.object({ + type: z.literal('progress'), + data: z + .object({ + value: z.number(), + max: z.number(), + }) + .passthrough(), + }), + z.object({ + type: z.literal('executing'), + data: z + .object({ + node: z.string().nullable().optional(), + display_node: z.string().optional(), + }) + .passthrough(), + }), + z.object({ + type: z.literal('executed'), + data: z + .object({ + output: z.union([ + z.object({ + images: z.array( + z.object({ + filename: z.string(), + subfolder: z.string(), + type: z.string(), + }), + ), + animated: z.array(z.boolean()).optional(), + }), + z.object({ + gifs: z.array( + z.object({ + filename: z.string(), + workflow: z.string(), + type: z.string(), + subfolder: z.string(), + format: z.string(), + }), + ), + }), + z.object({ + '3d': z.array( + z.object({ + filename: z.string(), + subfolder: z.string(), + type: z.string(), + }), + ), + }), + ]), + }) + .passthrough(), + }), + z.object({ + type: z.literal('progress_state'), + data: z.object({}).passthrough(), + }), +]) + +export type ComfyExecutionErrorData = { + exception_message?: string + exception_type?: string + node_id?: string | number + node_type?: string + traceback?: string[] +} + +// ComfyUI reports node failures with a full Python exception (often a multi-KB +// state_dict / size-mismatch dump). That is useless and overwhelming as a +// user-facing string, so we map the common, recognizable failures to a short, +// actionable sentence and keep the raw detail for the logs/debug panel only. +export function summarizeComfyExecutionError(data: ComfyExecutionErrorData): string { + const raw = (data.exception_message ?? '').trim() + const lower = raw.toLowerCase() + const type = (data.exception_type ?? '').toLowerCase() + + if ( + lower.includes('size mismatch') || + lower.includes('error(s) in loading state_dict') || + lower.includes('load_state_dict') + ) { + return "The selected model doesn't match this workflow (mismatched weights while loading). Pick a model that fits the preset, or choose a different preset." + } + if ( + lower.includes('out of memory') || + lower.includes('outofmemory') || + lower.includes('failed to allocate') || + (lower.includes('alloc') && lower.includes('memory')) + ) { + return 'Ran out of memory while generating. Try a smaller resolution or batch size.' + } + if ( + type.includes('filenotfound') || + lower.includes('no such file') || + lower.includes('cannot find') || + lower.includes('does not exist') + ) { + return 'A required model or file could not be found. Make sure the needed models are downloaded.' + } + + // Fallback: first meaningful line of the exception, trimmed to a sane length. + const firstLine = + raw + .split('\n') + .map((line) => line.trim()) + .find((line) => line.length > 0) ?? '' + const concise = firstLine.length > 200 ? `${firstLine.slice(0, 200)}…` : firstLine + const prefix = data.node_type ? `${data.node_type}: ` : '' + return concise ? `${prefix}${concise}` : 'The workflow failed during execution.' +} diff --git a/WebUI/src/assets/js/store/comfyUiPresets.ts b/WebUI/src/assets/js/store/comfyUiPresets.ts index dbe169d7b..d902dad73 100644 --- a/WebUI/src/assets/js/store/comfyUiPresets.ts +++ b/WebUI/src/assets/js/store/comfyUiPresets.ts @@ -15,7 +15,6 @@ import { useActivities } from './activities' import { createAppError } from '../errors/appError' import { useBackendServices } from '@/assets/js/store/backendServices.ts' import { usePromptStore } from './promptArea' -import { z } from 'zod' import { imageUrlToDataUri, isImageUrl, mediaUrl } from '@/lib/utils' import { getComfyAuthToken, invalidateComfyAuthToken } from '@/lib/loopbackAuth' import { @@ -23,6 +22,17 @@ import { findKeysByTitle, modifySettingInWorkflow, } from './comfyUiWorkflowHelpers' +import { + ComfyMessageSchema, + summarizeComfyExecutionError, + type ComfyExecutionErrorData, +} from './comfyUiMessages' +import { + bypassNode, + injectOvmsImageUrl, + normalizeModelPathsInWorkflow, + workflowUsesOvmsImage, +} from './comfyUiWorkflowTransforms' /** * Wraps fetch() with the ComfyUI loopback bearer token. The bundled @@ -50,248 +60,6 @@ async function comfyFetch(input: RequestInfo | URL, init?: RequestInit): Promise const WEBSOCKET_OPEN = 1 -const ComfyMessageSchema = z.discriminatedUnion('type', [ - z.object({ - type: z.literal('status'), - }), - z.object({ - type: z.literal('execution_start'), - data: z.object({}).passthrough(), - }), - z.object({ - type: z.literal('execution_success'), - data: z.object({}).passthrough(), - }), - z.object({ - type: z.literal('execution_error'), - data: z - .object({ - exception_message: z.string().optional(), - exception_type: z.string().optional(), - node_id: z.union([z.string(), z.number()]).optional(), - node_type: z.string().optional(), - traceback: z.array(z.string()).optional(), - }) - .passthrough(), - }), - z.object({ - type: z.literal('execution_interrupted'), - data: z.object({}).passthrough(), - }), - z.object({ - type: z.literal('execution_cached'), - data: z.object({}).passthrough(), - }), - z.object({ - type: z.literal('progress'), - data: z - .object({ - value: z.number(), - max: z.number(), - }) - .passthrough(), - }), - z.object({ - type: z.literal('executing'), - data: z - .object({ - node: z.string().nullable().optional(), - display_node: z.string().optional(), - }) - .passthrough(), - }), - z.object({ - type: z.literal('executed'), - data: z - .object({ - output: z.union([ - z.object({ - images: z.array( - z.object({ - filename: z.string(), - subfolder: z.string(), - type: z.string(), - }), - ), - animated: z.array(z.boolean()).optional(), - }), - z.object({ - gifs: z.array( - z.object({ - filename: z.string(), - workflow: z.string(), - type: z.string(), - subfolder: z.string(), - format: z.string(), - }), - ), - }), - z.object({ - '3d': z.array( - z.object({ - filename: z.string(), - subfolder: z.string(), - type: z.string(), - }), - ), - }), - ]), - }) - .passthrough(), - }), - z.object({ - type: z.literal('progress_state'), - data: z.object({}).passthrough(), - }), -]) - -type ComfyExecutionErrorData = { - exception_message?: string - exception_type?: string - node_id?: string | number - node_type?: string - traceback?: string[] -} - -// ComfyUI reports node failures with a full Python exception (often a multi-KB -// state_dict / size-mismatch dump). That is useless and overwhelming as a -// user-facing string, so we map the common, recognizable failures to a short, -// actionable sentence and keep the raw detail for the logs/debug panel only. -function summarizeComfyExecutionError(data: ComfyExecutionErrorData): string { - const raw = (data.exception_message ?? '').trim() - const lower = raw.toLowerCase() - const type = (data.exception_type ?? '').toLowerCase() - - if ( - lower.includes('size mismatch') || - lower.includes('error(s) in loading state_dict') || - lower.includes('load_state_dict') - ) { - return "The selected model doesn't match this workflow (mismatched weights while loading). Pick a model that fits the preset, or choose a different preset." - } - if ( - lower.includes('out of memory') || - lower.includes('outofmemory') || - lower.includes('failed to allocate') || - (lower.includes('alloc') && lower.includes('memory')) - ) { - return 'Ran out of memory while generating. Try a smaller resolution or batch size.' - } - if ( - type.includes('filenotfound') || - lower.includes('no such file') || - lower.includes('cannot find') || - lower.includes('does not exist') - ) { - return 'A required model or file could not be found. Make sure the needed models are downloaded.' - } - - // Fallback: first meaningful line of the exception, trimmed to a sane length. - const firstLine = - raw - .split('\n') - .map((line) => line.trim()) - .find((line) => line.length > 0) ?? '' - const concise = firstLine.length > 200 ? `${firstLine.slice(0, 200)}…` : firstLine - const prefix = data.node_type ? `${data.node_type}: ` : '' - return concise ? `${prefix}${concise}` : 'The workflow failed during execution.' -} - -const OVMS_IMAGE_CLASS_TYPES = ['OpenAICompatibleImageGeneration', 'OpenAICompatibleImageEdit'] - -function workflowUsesOvmsImage(workflow: ComfyUIApiWorkflow): boolean { - return OVMS_IMAGE_CLASS_TYPES.some((ct) => findKeysByClassType(workflow, ct).length > 0) -} - -function injectOvmsImageUrl(workflow: ComfyUIApiWorkflow, url: string): void { - for (const classType of OVMS_IMAGE_CLASS_TYPES) { - for (const key of findKeysByClassType(workflow, classType)) { - const inputs = workflow[key]?.inputs - if (!inputs || typeof inputs !== 'object') continue - if ('base_url' in inputs) { - inputs['base_url'] = url - } - // OVMS registers the served graph under the slash-flattened repo id - // (see `--source_model` in openVINOBackendService.ts), so the model - // sent in the OpenAI-compatible request must use the same form. - if ('model' in inputs && typeof inputs['model'] === 'string') { - inputs['model'] = inputs['model'].split('/').join('---') - } - } - } -} - -/** ComfyUI node input names that hold model/file paths; separator is OS-dependent (see main.ts preset handling). */ -const COMFY_MODEL_PATH_INPUTS = new Set([ - 'ckpt_name', - 'lora_name', - 'text_encoder', - 'vae_name', - 'unet_name', - 'clip_name', - 'model_name', - 'control_net_name', -]) - -function normalizeModelPathsInWorkflow( - workflow: ComfyUIApiWorkflow, - platform: NodeJS.Platform, -): void { - for (const node of Object.values(workflow)) { - const inputs = (node as { inputs?: Record }).inputs - if (!inputs) continue - for (const [inputName, value] of Object.entries(inputs)) { - if (COMFY_MODEL_PATH_INPUTS.has(inputName) && typeof value === 'string') { - inputs[inputName] = modelNameForComfyApi(value, platform) - } - } - } -} - -/** - * Bypass a node by rewiring its outputs to its upstream and removing the node. - * Supported: LoraLoader (output 0 = model from input "model", output 1 = clip from input "clip"). - */ -function bypassNode(workflow: ComfyUIApiWorkflow, nodeId: string): void { - const node = workflow[nodeId] as - | { class_type?: string; inputs?: Record } - | undefined - if (!node?.inputs) return - const classType = node.class_type - let rewire: [number, [string, number]][] - if (classType === 'LoraLoader') { - const model = node.inputs.model as [string, number] | undefined - const clip = node.inputs.clip as [string, number] | undefined - if (!model || !clip) return - rewire = [ - [0, model], - [1, clip], - ] - } else { - return - } - for (const entry of Object.values(workflow)) { - const inputs = (entry as { inputs?: Record }).inputs - if (!inputs) continue - for (const key of Object.keys(inputs)) { - const v = inputs[key] - if ( - Array.isArray(v) && - v.length === 2 && - typeof v[0] === 'string' && - typeof v[1] === 'number' - ) { - if (v[0] === nodeId) { - const slot = v[1] - const upstream = rewire.find(([s]) => s === slot)?.[1] - if (upstream) inputs[key] = upstream - } - } - } - } - delete workflow[nodeId] -} - /** Rewire and remove optional model nodes whose value is None (e.g. LoRA bypass). */ function bypassOptionalModelNodes(workflow: ComfyUIApiWorkflow): void { const imageGeneration = useImageGenerationPresets() diff --git a/WebUI/src/assets/js/store/comfyUiWorkflowTransforms.ts b/WebUI/src/assets/js/store/comfyUiWorkflowTransforms.ts new file mode 100644 index 000000000..f007719a2 --- /dev/null +++ b/WebUI/src/assets/js/store/comfyUiWorkflowTransforms.ts @@ -0,0 +1,98 @@ +import { ComfyUIApiWorkflow } from './presets' +import { modelNameForComfyApi } from './imageGenerationPresets' +import { findKeysByClassType } from './comfyUiWorkflowHelpers' + +const OVMS_IMAGE_CLASS_TYPES = ['OpenAICompatibleImageGeneration', 'OpenAICompatibleImageEdit'] + +export function workflowUsesOvmsImage(workflow: ComfyUIApiWorkflow): boolean { + return OVMS_IMAGE_CLASS_TYPES.some((ct) => findKeysByClassType(workflow, ct).length > 0) +} + +export function injectOvmsImageUrl(workflow: ComfyUIApiWorkflow, url: string): void { + for (const classType of OVMS_IMAGE_CLASS_TYPES) { + for (const key of findKeysByClassType(workflow, classType)) { + const inputs = workflow[key]?.inputs + if (!inputs || typeof inputs !== 'object') continue + if ('base_url' in inputs) { + inputs['base_url'] = url + } + // OVMS registers the served graph under the slash-flattened repo id + // (see `--source_model` in openVINOBackendService.ts), so the model + // sent in the OpenAI-compatible request must use the same form. + if ('model' in inputs && typeof inputs['model'] === 'string') { + inputs['model'] = inputs['model'].split('/').join('---') + } + } + } +} + +/** ComfyUI node input names that hold model/file paths; separator is OS-dependent (see main.ts preset handling). */ +const COMFY_MODEL_PATH_INPUTS = new Set([ + 'ckpt_name', + 'lora_name', + 'text_encoder', + 'vae_name', + 'unet_name', + 'clip_name', + 'model_name', + 'control_net_name', +]) + +export function normalizeModelPathsInWorkflow( + workflow: ComfyUIApiWorkflow, + platform: NodeJS.Platform, +): void { + for (const node of Object.values(workflow)) { + const inputs = (node as { inputs?: Record }).inputs + if (!inputs) continue + for (const [inputName, value] of Object.entries(inputs)) { + if (COMFY_MODEL_PATH_INPUTS.has(inputName) && typeof value === 'string') { + inputs[inputName] = modelNameForComfyApi(value, platform) + } + } + } +} + +/** + * Bypass a node by rewiring its outputs to its upstream and removing the node. + * Supported: LoraLoader (output 0 = model from input "model", output 1 = clip from input "clip"). + */ +export function bypassNode(workflow: ComfyUIApiWorkflow, nodeId: string): void { + const node = workflow[nodeId] as + | { class_type?: string; inputs?: Record } + | undefined + if (!node?.inputs) return + const classType = node.class_type + let rewire: [number, [string, number]][] + if (classType === 'LoraLoader') { + const model = node.inputs.model as [string, number] | undefined + const clip = node.inputs.clip as [string, number] | undefined + if (!model || !clip) return + rewire = [ + [0, model], + [1, clip], + ] + } else { + return + } + for (const entry of Object.values(workflow)) { + const inputs = (entry as { inputs?: Record }).inputs + if (!inputs) continue + for (const key of Object.keys(inputs)) { + const v = inputs[key] + if ( + Array.isArray(v) && + v.length === 2 && + typeof v[0] === 'string' && + typeof v[1] === 'number' + ) { + if (v[0] === nodeId) { + const slot = v[1] + const upstream = rewire.find(([s]) => s === slot)?.[1] + if (upstream) inputs[key] = upstream + } + } + } + } + delete workflow[nodeId] +} diff --git a/WebUI/src/components/InstallationManagement.vue b/WebUI/src/components/InstallationManagement.vue index 519ce5b2c..08680769d 100644 --- a/WebUI/src/components/InstallationManagement.vue +++ b/WebUI/src/components/InstallationManagement.vue @@ -433,7 +433,6 @@ import { mapToDisplayStatus, compareVersions, } from '@/lib/utils.ts' -import * as toast from '@/assets/js/toast.ts' import { useBackendServices } from '@/assets/js/store/backendServices' import LanguageSelector from '@/components/LanguageSelector.vue' import BackendOptions from '@/components/BackendOptions.vue' @@ -444,6 +443,7 @@ import { Checkbox } from '@/components/ui/checkbox' import { Tooltip, TooltipContent, TooltipTrigger, TooltipProvider } from '@/components/ui/tooltip' import type { ErrorDetails } from '../../electron/subprocesses/service' import { useProductMode } from '@/assets/js/store/productMode' +import { useErrors } from '@/assets/js/store/errors' const emits = defineEmits<{ (e: 'close'): void @@ -456,6 +456,7 @@ type ExtendedApiServiceInformation = ApiServiceInformation & { const backendServices = useBackendServices() const productModeStore = useProductMode() +const errors = useErrors() // App version for AI Backend display (fetched directly to avoid timing issues with globalSetup.initSetup) const appVersion = ref('...') @@ -715,10 +716,14 @@ async function installBackend(name: BackendServiceName) { if (setupProgress.success) { await restartBackend(name) } else { - const errorMessage = setupProgress.errorDetails - ? 'Setup failed - Click the info icon for details' - : 'Setup failed' - toast.error(errorMessage) + errors.report('Backend setup failed', { + category: 'setup', + code: 'setup/backend-setup-failed', + userMessage: setupProgress.errorDetails + ? 'Setup failed - Click the info icon for details' + : 'Setup failed', + context: { serviceName: name }, + }) loadingComponents.value.delete(name) } } @@ -727,7 +732,12 @@ async function repairBackend(name: BackendServiceName) { loadingComponents.value.add(name) const stopStatus = await backendServices.stopService(name) if (stopStatus !== 'stopped') { - toast.error('Service failed to stop') + errors.report('Service failed to stop', { + category: 'backend', + code: 'backend/stop-failed', + userMessage: 'Service failed to stop', + context: { serviceName: name }, + }) return } await installBackend(name) @@ -737,7 +747,12 @@ async function restartBackend(name: BackendServiceName) { loadingComponents.value.add(name) const stopStatus = await backendServices.stopService(name) if (stopStatus !== 'stopped') { - toast.error('Service failed to stop') + errors.report('Service failed to stop', { + category: 'backend', + code: 'backend/stop-failed', + userMessage: 'Service failed to stop', + context: { serviceName: name }, + }) loadingComponents.value.delete(name) return } @@ -745,22 +760,30 @@ async function restartBackend(name: BackendServiceName) { try { const startStatus = await backendServices.startService(name) if (startStatus !== 'running') { - // Service failed to start - show detailed error message + // Service failed to start - detail is available via the info icon. const errorDetails = backendServices.getServiceErrorDetails(name) - const errorMessage = errorDetails - ? 'Service failed to start - Click the info icon for details' - : 'Service failed to start' - toast.error(errorMessage) + errors.report('Service failed to start', { + category: 'backend', + code: 'backend/start-failed', + userMessage: errorDetails + ? 'Service failed to start - Click the info icon for details' + : 'Service failed to start', + context: { serviceName: name }, + }) loadingComponents.value.delete(name) return } } catch (error) { - // Exception during startup - show detailed error message + // Exception during startup - detail is available via the info icon. const errorDetails = backendServices.getServiceErrorDetails(name) - const errorMessage = errorDetails - ? 'Service startup failed - Click the info icon for details' - : `Service startup failed: ${error instanceof Error ? error.message : String(error)}` - toast.error(errorMessage) + errors.report(error, { + category: 'backend', + code: 'backend/start-failed', + userMessage: errorDetails + ? 'Service startup failed - Click the info icon for details' + : `Service startup failed: ${error instanceof Error ? error.message : String(error)}`, + context: { serviceName: name }, + }) loadingComponents.value.delete(name) return } diff --git a/WebUI/src/components/McpServerDialog.vue b/WebUI/src/components/McpServerDialog.vue index f49d9bee3..274e0bb00 100644 --- a/WebUI/src/components/McpServerDialog.vue +++ b/WebUI/src/components/McpServerDialog.vue @@ -104,6 +104,7 @@ import { Textarea } from '@/components/ui/textarea' import { Label } from '@/components/ui/label' import { RadioGroup, RadioGroupItem } from '@/components/ui/radio-group' import { useMcp } from '@/assets/js/store/mcp' +import { useErrors } from '@/assets/js/store/errors' import * as toast from '@/assets/js/toast' import type { McpServerConfig } from '../../electron/subprocesses/mcpServers' @@ -120,6 +121,7 @@ const emits = defineEmits<{ }>() const mcp = useMcp() +const errors = useErrors() const transport = ref<'stdio' | 'http'>('stdio') const displayName = ref('') @@ -265,7 +267,14 @@ async function handleSubmit() { error instanceof Error ? error.message : `Failed to ${isAddMode.value ? 'add' : 'update'} server` - toast.error(errorMessage.value) + // The dialog renders errorMessage inline, so report as 'inline' to centralize + // logging/de-duplication without a redundant toast. + errors.report(error, { + category: 'backend', + code: `backend/mcp-${isAddMode.value ? 'add' : 'update'}-failed`, + surface: 'inline', + userMessage: errorMessage.value, + }) } finally { isSubmitting.value = false } @@ -300,7 +309,11 @@ async function handleRemove() { emits('update:open', false) resetForm() } catch (error) { - toast.error(error instanceof Error ? error.message : 'Failed to remove server') + errors.report(error, { + category: 'backend', + code: 'backend/mcp-remove-failed', + userMessage: error instanceof Error ? error.message : 'Failed to remove server', + }) } finally { isSubmitting.value = false } diff --git a/WebUI/src/env.d.ts b/WebUI/src/env.d.ts index 4fb5054f0..b3c8d2129 100644 --- a/WebUI/src/env.d.ts +++ b/WebUI/src/env.d.ts @@ -312,11 +312,9 @@ type electronAPI = { getInitSetting(): Promise updateModelPaths(modelPaths: ModelPaths): Promise restorePathsSettings(): Promise - refreshLLMModles(): Promise loadModels(): Promise zoomIn(): Promise zoomOut(): Promise - getDownloadedLLMs(): Promise getDownloadedGGUFLLMs(): Promise getDownloadedOpenVINOLLMModels(): Promise getDownloadedEmbeddingModels(): Promise diff --git a/WebUI/src/views/PromptArea.vue b/WebUI/src/views/PromptArea.vue index 2bbae3933..cd5d67a61 100644 --- a/WebUI/src/views/PromptArea.vue +++ b/WebUI/src/views/PromptArea.vue @@ -271,6 +271,7 @@ import { import { useOpenAiCompatibleChat } from '@/assets/js/store/openAiCompatibleChat' import { useConversations } from '@/assets/js/store/conversations' import { useActivities } from '@/assets/js/store/activities' +import { useErrors } from '@/assets/js/store/errors' import { useTextInference, type ValidFileExtension, @@ -312,6 +313,7 @@ const openAiCompatibleChat = useOpenAiCompatibleChat() const textInference = useTextInference() const conversations = useConversations() const activities = useActivities() +const errors = useErrors() const textareaRef = ref() const isTextareaFocused = ref(false) const presetsStore = usePresets() @@ -594,7 +596,11 @@ async function handleRecordingClick() { await audioRecorder.startRecording() if (audioRecorder.error) { - toast.error(audioRecorder.error) + errors.report(audioRecorder.error, { + category: 'inference', + code: 'inference/audio-record-failed', + userMessage: audioRecorder.error, + }) } } @@ -704,8 +710,11 @@ async function handleComfyUIImageUpload(imageFiles: File[]) { } } } catch (error) { - console.error('Error processing image:', error) - toast.error('Failed to load image') + errors.report(error, { + category: 'inference', + code: 'inference/image-load-failed', + userMessage: 'Failed to load image', + }) } finally { URL.revokeObjectURL(imageUrl) } @@ -842,8 +851,11 @@ async function addDocumentsToRagList(files: File[]) { await textInference.addDocumentToRagList(newDocument) } catch (error) { - console.error('Error adding document to RAG list:', error) - toast.error(i18nState.RAG_UPLOAD_TYPE_ERROR) + errors.report(error, { + category: 'inference', + code: 'inference/rag-add-failed', + userMessage: i18nState.RAG_UPLOAD_TYPE_ERROR, + }) } } } diff --git a/backend_shared/aipg_loopback_auth.py b/backend_shared/aipg_loopback_auth.py new file mode 100644 index 000000000..fe8c7dff8 --- /dev/null +++ b/backend_shared/aipg_loopback_auth.py @@ -0,0 +1,61 @@ +"""Shared loopback + per-launch-token authentication for AI Playground's local +Python backends (ai-backend, home-agent, ...). + +Each backend binds to 127.0.0.1, but on a shared host any local peer can still +reach the port. We therefore require requests to (a) originate from a loopback +address and (b) carry an `X-AIPG-Auth` header matching the per-launch token the +Electron main process injects via the `AIPG_LOOPBACK_TOKEN` env var. This is the +attack model documented in the CWE-494 report against the local API ports. + +This module is intentionally transport-agnostic (no Flask/aiohttp import) and +stdlib-only, so it can be imported from every backend's isolated virtualenv. It +lives in a sibling `backend_shared/` directory that each backend adds to +`sys.path`; callers build their own HTTP response from the returned verdict. +""" + +import hmac +import os +from typing import Iterable, Optional + +# Remote addresses we treat as loopback. +LOOPBACK_REMOTE_ADDRS = frozenset({"127.0.0.1", "::1"}) + +# Paths reachable without a token. `/healthy` must stay open so the Electron +# service registry can detect the service is up before it has the token. +DEFAULT_AUTH_EXEMPT_PATHS = frozenset({"/healthy"}) + + +def get_loopback_token() -> str: + """The per-launch token provisioned by the Electron main process (or "").""" + return os.environ.get("AIPG_LOOPBACK_TOKEN", "") + + +def evaluate_loopback_auth( + remote_addr: Optional[str], + method: str, + path: str, + provided_token: str, + *, + expected_token: str, + exempt_paths: Iterable[str] = DEFAULT_AUTH_EXEMPT_PATHS, +) -> Optional[tuple[int, str]]: + """Decide whether a request is allowed. + + Returns ``None`` when the request should be allowed to proceed, or a + ``(status_code, error_message)`` tuple describing the rejection. CORS + preflight (``OPTIONS``) requests pass the loopback check and then return + ``None`` — the browser strips custom headers from preflight, so the caller + is responsible for answering it (e.g. a 204) without requiring the token. + The real request that follows is authenticated normally. + """ + if remote_addr not in LOOPBACK_REMOTE_ADDRS: + return 403, "loopback only" + if method == "OPTIONS": + return None + if path in exempt_paths: + return None + if not expected_token: + return 503, "service not provisioned" + if not provided_token or not hmac.compare_digest(provided_token, expected_token): + return 401, "unauthorized" + return None diff --git a/home-agent/web_api.py b/home-agent/web_api.py index d7eb353ed..0ff2dd877 100644 --- a/home-agent/web_api.py +++ b/home-agent/web_api.py @@ -8,10 +8,10 @@ """ import argparse -import hmac import logging import os import re +import sys import threading from flask import Flask, jsonify, request @@ -21,6 +21,13 @@ from channels import registry from channels.types import ALL_CHANNEL_KINDS +# Shared loopback-auth lives in a sibling backend_shared/ directory so the same +# logic is used by every local Python backend (see backend_shared/). +sys.path.insert( + 0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "backend_shared") +) +from aipg_loopback_auth import evaluate_loopback_auth, get_loopback_token # noqa: E402 + app = Flask(__name__) CORS(app) @@ -28,26 +35,21 @@ # Flask binds to 127.0.0.1 but on a shared host any local peer can still reach # our port. Require an `X-AIPG-Auth` header matching the per-launch token the # Electron main process injected via env. Mirrors the `ai-backend` pattern. -_LOOPBACK_AUTH_TOKEN = os.environ.get("AIPG_LOOPBACK_TOKEN", "") -_LOOPBACK_REMOTE_ADDRS = frozenset({"127.0.0.1", "::1"}) -# `/healthy` must remain reachable so the service registry can probe readiness -# before it has obtained the token. -_AUTH_EXEMPT_PATHS = frozenset({"/healthy"}) +_LOOPBACK_AUTH_TOKEN = get_loopback_token() @app.before_request def _enforce_loopback_and_auth(): - if request.remote_addr not in _LOOPBACK_REMOTE_ADDRS: - return jsonify({"error": "loopback only"}), 403 - if request.method == "OPTIONS": - return None - if request.path in _AUTH_EXEMPT_PATHS: - return None - if not _LOOPBACK_AUTH_TOKEN: - return jsonify({"error": "service not provisioned"}), 503 - provided = request.headers.get("X-AIPG-Auth", "") - if not provided or not hmac.compare_digest(provided, _LOOPBACK_AUTH_TOKEN): - return jsonify({"error": "unauthorized"}), 401 + rejection = evaluate_loopback_auth( + request.remote_addr, + request.method, + request.path, + request.headers.get("X-AIPG-Auth", ""), + expected_token=_LOOPBACK_AUTH_TOKEN, + ) + if rejection is not None: + status, message = rejection + return jsonify({"error": message}), status return None @@ -113,7 +115,9 @@ def _redacting_factory(*args, **kwargs): # type: ignore[no-untyped-def] if isinstance(record.args, tuple): record.args = tuple(_redact_token(a, patterns) for a in record.args) elif isinstance(record.args, dict): - record.args = {k: _redact_token(v, patterns) for k, v in record.args.items()} + record.args = { + k: _redact_token(v, patterns) for k, v in record.args.items() + } return record logging.setLogRecordFactory(_redacting_factory) @@ -139,6 +143,7 @@ def _redacting_factory(*args, **kwargs): # type: ignore[no-untyped-def] # ── Health ──────────────────────────────────────────────────────────────────── + @app.get("/healthy") def healthy(): return jsonify({"status": "ok"}) @@ -146,6 +151,7 @@ def healthy(): # ── Upstream control ────────────────────────────────────────────────────────── + @app.post("/set-upstream") def set_upstream(): global _upstream_url @@ -162,6 +168,7 @@ def set_upstream(): # All channel-specific routes funnel through these handlers, which look up # the concrete channel from `registry.CHANNELS` and delegate. + def _get_channel_or_404(kind: str): ch = registry.get(kind) if ch is None: @@ -235,6 +242,7 @@ def channel_send(kind: str, action: str): # ── Chat completions proxy ──────────────────────────────────────────────────── + @app.post("/v1/chat/completions") def chat_completions(): upstream = request.headers.get("X-Upstream-Url") diff --git a/service/model_download_adpater.py b/service/model_download_adpater.py index 2e2d5b499..d4a7f2a46 100644 --- a/service/model_download_adpater.py +++ b/service/model_download_adpater.py @@ -14,7 +14,7 @@ class Model_Downloader_Adapter: msg_queue: Queue finish: bool - singal: threading.Event + signal: threading.Event file_downloader: FileDownloader hf_downloader: HFPlaygroundDownloader has_error: bool @@ -24,7 +24,7 @@ def __init__(self, hf_token=None): self.msg_queue = Queue(-1) self.finish = False self.user_stop = False - self.singal = threading.Event() + self.signal = threading.Event() self.file_downloader = FileDownloader() self.file_downloader.on_download_progress = ( self.download_model_progress_callback @@ -40,7 +40,7 @@ def __init__(self, hf_token=None): def put_msg(self, data): self.msg_queue.put_nowait(data) - self.singal.set() + self.signal.set() def download_model_progress_callback( self, repo_id: str, download_size: int, total_size: int, speed: int @@ -101,7 +101,10 @@ def error_callback(self, ex: Exception): def download(self, model_download_list: List[DownloadModelData]): self.has_error = False - threading.Thread(target=self.__start_download, kwargs={"model_download_list": model_download_list}).start() + threading.Thread( + target=self.__start_download, + kwargs={"model_download_list": model_download_list}, + ).start() return self.generator() def __start_download(self, model_download_list: List[DownloadModelData]): @@ -120,7 +123,10 @@ def __start_download(self, model_download_list: List[DownloadModelData]): # Copy faceswap/facerestore models to ComfyUI directory after download completes # This must happen here (not in callback) to ensure file is fully written - if item.backend == "comfyui" and item.type in ("faceswap", "facerestore"): + if item.backend == "comfyui" and item.type in ( + "faceswap", + "facerestore", + ): utils.copy_faceswap_facerestore_to_comfyui( item.type, item.repo_id, @@ -132,7 +138,7 @@ def __start_download(self, model_download_list: List[DownloadModelData]): self.error_callback(ex) finally: self.finish = True - self.singal.set() + self.signal.set() if _adapter is self: _adapter = None @@ -153,8 +159,8 @@ def generator(self): except Empty: break if not self.finish: - self.singal.clear() - self.singal.wait() + self.signal.clear() + self.signal.wait() else: break diff --git a/service/web_api.py b/service/web_api.py index cbab8629d..c086e4e88 100644 --- a/service/web_api.py +++ b/service/web_api.py @@ -49,12 +49,24 @@ def get_added_dll_directories(): except ImportError: pass - import hmac import os import threading from flask import jsonify, request, Response, stream_with_context from apiflask import APIFlask + # Shared loopback-auth lives in a sibling backend_shared/ directory so the + # same logic is used by every local Python backend (see backend_shared/). + sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "backend_shared" + ), + ) + from aipg_loopback_auth import ( # noqa: E402 + evaluate_loopback_auth, + get_loopback_token, + ) + import model_download_adpater import utils from model_downloader import HFPlaygroundDownloader @@ -71,12 +83,7 @@ def get_added_dll_directories(): # would be reachable by any other process on the box (including low-IL # processes and host-networked containers), which is the attack model # documented in the CWE-494 report against /api/comfyUi/loadCustomNodes. - _LOOPBACK_AUTH_TOKEN = os.environ.get("AIPG_LOOPBACK_TOKEN", "") - _LOOPBACK_REMOTE_ADDRS = frozenset({"127.0.0.1", "::1"}) - # Endpoints that are allowed without a bearer token. /healthy must remain - # reachable so the Electron service registry can detect when the service - # is up before it has obtained the token. - _AUTH_EXEMPT_PATHS = frozenset({"/healthy"}) + _LOOPBACK_AUTH_TOKEN = get_loopback_token() # Loopback hostnames whose Origin we are willing to echo back as # Access-Control-Allow-Origin. The renderer may load via 127.0.0.1 or # localhost depending on platform/devtools; production Electron loads @@ -119,27 +126,29 @@ def _origin_is_loopback(origin: str) -> bool: @app.before_request def _enforce_loopback_and_auth(): - if request.remote_addr not in _LOOPBACK_REMOTE_ADDRS: - logging.warning( - f"rejecting non-loopback request from {request.remote_addr} to {request.path}" - ) - return jsonify({"error": "loopback only"}), 403 + # Use a dedicated X-AIPG-Auth header so the existing + # `Authorization: Bearer ` semantics for /api/downloadModel + # remain intact. + rejection = evaluate_loopback_auth( + request.remote_addr, + request.method, + request.path, + request.headers.get("X-AIPG-Auth", ""), + expected_token=_LOOPBACK_AUTH_TOKEN, + ) + if rejection is not None: + status, message = rejection + if status == 403: + logging.warning( + f"rejecting non-loopback request from {request.remote_addr} to {request.path}" + ) + return jsonify({"error": message}), status # CORS preflight requests do NOT carry the X-AIPG-Auth header by # design (the browser strips custom headers from preflight). Reply # 204 here and let _attach_cors_headers below add the actual CORS # headers; the real request that follows will be authenticated. if request.method == "OPTIONS": return Response(status=204) - if request.path in _AUTH_EXEMPT_PATHS: - return None - if not _LOOPBACK_AUTH_TOKEN: - return jsonify({"error": "service not provisioned"}), 503 - # Use a dedicated X-AIPG-Auth header so the existing - # `Authorization: Bearer ` semantics for /api/downloadModel - # remain intact. - provided = request.headers.get("X-AIPG-Auth", "") - if not provided or not hmac.compare_digest(provided, _LOOPBACK_AUTH_TOKEN): - return jsonify({"error": "unauthorized"}), 401 return None @app.after_request