diff --git a/locales/en.json b/locales/en.json index feda5029d..5cb177b87 100644 --- a/locales/en.json +++ b/locales/en.json @@ -481,6 +481,7 @@ "launchStart": "Starting ComfyUI…", "repair": "Repairing installation…", "torchRepair": "Restoring GPU PyTorch…", + "agentRequirements": "Installing agent requirements…", "securityScan": "Running security scan…", "mountLibraries": "Loading libraries…", "gpu": "Initializing GPU…", diff --git a/locales/zh.json b/locales/zh.json index 14c114ca7..ab96a2412 100644 --- a/locales/zh.json +++ b/locales/zh.json @@ -481,6 +481,7 @@ "launchStart": "正在启动 ComfyUI…", "repair": "正在修复安装…", "torchRepair": "正在恢复 GPU PyTorch…", + "agentRequirements": "正在安装智能体依赖…", "securityScan": "正在进行安全扫描…", "mountLibraries": "正在加载库…", "gpu": "正在初始化 GPU…", diff --git a/src/main/lib/agentRequirementsLaunch.test.ts b/src/main/lib/agentRequirementsLaunch.test.ts new file mode 100644 index 000000000..1dee434d6 --- /dev/null +++ b/src/main/lib/agentRequirementsLaunch.test.ts @@ -0,0 +1,390 @@ +import fs from 'fs' +import os from 'os' +import path from 'path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const mirrors = { pypiMirror: 'https://mirror.example/simple/', useChineseMirrors: false } + +vi.mock('../settings', () => ({ + getMirrorConfig: () => mirrors +})) + +vi.mock('./pip', () => ({ + installFilteredRequirementsDetailed: vi.fn(async () => ({ code: 0, output: '' })) +})) + +import { installAgentRequirements, planAgentRequirementsInstall } from './agentRequirementsLaunch' +import { installFilteredRequirementsDetailed } from './pip' +import { getUvPath, getVenvPythonPath, getLegacyVenvUvPath } from './pythonEnv' +import type { InstallationRecord } from '../installations' + +const mockInstall = vi.mocked(installFilteredRequirementsDetailed) + +let installDir = '' + +/** Create a file (and its parents) the way the real layout has it on disk. */ +function touch(target: string): void { + fs.mkdirSync(path.dirname(target), { recursive: true }) + fs.writeFileSync(target, '') +} + +/** A standalone install: `standalone-env` uv + the managed `ComfyUI/.venv`. */ +function managedInstall(): InstallationRecord { + touch(getUvPath(installDir)) + touch(getVenvPythonPath(installDir)) + return { installPath: installDir } as unknown as InstallationRecord +} + +function writeAgentRequirements(): string { + const reqPath = path.join(installDir, 'ComfyUI', 'agent_requirements.txt') + fs.mkdirSync(path.dirname(reqPath), { recursive: true }) + fs.writeFileSync(reqPath, 'comfyui-agent==1.0.0\n') + return reqPath +} + +describe('planAgentRequirementsInstall', () => { + beforeEach(() => { + vi.clearAllMocks() + installDir = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-reqs-')) + fs.mkdirSync(path.join(installDir, 'ComfyUI'), { recursive: true }) + }) + + afterEach(() => { + fs.rmSync(installDir, { recursive: true, force: true }) + }) + + it('plans the install when the user typed the flag by hand', () => { + const reqPath = writeAgentRequirements() + const inst = managedInstall() + + expect( + planAgentRequirementsInstall(inst, ['-s', 'main.py', '--enable-agent', '--listen']) + ).toEqual({ + reqPath, + uvPath: getUvPath(installDir), + pythonPath: getVenvPythonPath(installDir), + installPath: installDir + }) + }) + + it('plans the install when a beta grant added the flag', () => { + // A grant reaches the final args ahead of the user's own, which is the only + // difference from the hand-typed case - the args are all this reads. + const reqPath = writeAgentRequirements() + const inst = managedInstall() + + const plan = planAgentRequirementsInstall(inst, [ + '-s', + 'main.py', + '--feature-flag', + 'show_signin_button=true', + '--enable-agent', + '--listen' + ]) + + expect(plan?.reqPath).toBe(reqPath) + }) + + it('plans nothing when the flag is absent', () => { + writeAgentRequirements() + const inst = managedInstall() + + expect(planAgentRequirementsInstall(inst, ['-s', 'main.py', '--listen'])).toBeNull() + }) + + it('plans nothing when core ships no agent requirements file', () => { + const inst = managedInstall() + + expect(planAgentRequirementsInstall(inst, ['--enable-agent'])).toBeNull() + }) + + it('plans nothing for an install with no Desktop-managed Python environment', () => { + // Portable, git and build installs: the requirements file may well be there, + // but there is no uv/venv pair to install it into. + writeAgentRequirements() + const inst = { installPath: installDir } as unknown as InstallationRecord + + expect(planAgentRequirementsInstall(inst, ['--enable-agent'])).toBeNull() + }) + + it('plans nothing when uv is missing from an otherwise managed install', () => { + writeAgentRequirements() + touch(getVenvPythonPath(installDir)) + const inst = { installPath: installDir } as unknown as InstallationRecord + + expect(planAgentRequirementsInstall(inst, ['--enable-agent'])).toBeNull() + }) + + it('targets the legacy venv for an adopted install', () => { + const reqPath = writeAgentRequirements() + const adoptedBaseDir = path.join(installDir, 'legacy') + const adoptedPythonPath = path.join(adoptedBaseDir, '.venv', 'python-for-test') + touch(getLegacyVenvUvPath(adoptedBaseDir)) + touch(adoptedPythonPath) + const inst = { + installPath: installDir, + adopted: true, + adoptedBaseDir, + adoptedPythonPath + } as unknown as InstallationRecord + + expect(planAgentRequirementsInstall(inst, ['--enable-agent'])).toEqual({ + reqPath, + uvPath: getLegacyVenvUvPath(adoptedBaseDir), + pythonPath: adoptedPythonPath, + installPath: installDir + }) + }) +}) + +describe('installAgentRequirements', () => { + const plan = { + reqPath: '/inst/ComfyUI/agent_requirements.txt', + uvPath: '/inst/standalone-env/bin/uv', + pythonPath: '/inst/ComfyUI/.venv/bin/python3', + installPath: '/inst' + } + + beforeEach(() => { + vi.clearAllMocks() + }) + + it('installs the planned file through the shared uv helper', async () => { + const sendOutput = vi.fn() + + await installAgentRequirements(plan, sendOutput) + + expect(mockInstall).toHaveBeenCalledWith( + plan.reqPath, + plan.uvPath, + plan.pythonPath, + plan.installPath, + '.launch-agent-reqs.txt', + sendOutput, + expect.any(AbortSignal), + mirrors + ) + expect(sendOutput.mock.calls.join('')).toContain('Installing agent requirements') + }) + + it('reports a failed install and resolves so the launch continues', async () => { + mockInstall.mockResolvedValueOnce({ code: 2, output: 'No solution found\n' }) + const sendOutput = vi.fn() + + await expect(installAgentRequirements(plan, sendOutput)).resolves.toBeUndefined() + + const reported = sendOutput.mock.calls.join('') + expect(reported).toContain('exited with code 2') + expect(reported).toContain('No solution found') + }) + + it('reports a thrown install and resolves so the launch continues', async () => { + mockInstall.mockRejectedValueOnce(new Error('EACCES: permission denied')) + const sendOutput = vi.fn() + + await expect(installAgentRequirements(plan, sendOutput)).resolves.toBeUndefined() + + expect(sendOutput.mock.calls.join('')).toContain('EACCES: permission denied') + }) + + it('abandons an install that outlives the ceiling and lets the launch continue', async () => { + // The whole point of the bound: a stalled uv must not hold the user at the + // launcher. Core starts with the flag and disables the agent itself. + vi.useFakeTimers() + try { + let uvSignal: AbortSignal | undefined + mockInstall.mockImplementationOnce( + async (...args: Parameters) => { + uvSignal = args[6] + // Resolve only once something aborts uv, the way the real helper does. + return new Promise((resolve) => { + uvSignal!.addEventListener('abort', () => resolve({ code: 1, output: '' }), { + once: true + }) + }) + } + ) + const sendOutput = vi.fn() + + const pending = installAgentRequirements(plan, sendOutput) + await vi.advanceTimersByTimeAsync(120_000) + await expect(pending).resolves.toBeUndefined() + + expect(uvSignal?.aborted).toBe(true) + expect(sendOutput.mock.calls.join('')).toContain('starting ComfyUI without it') + } finally { + vi.useRealTimers() + } + }) + + it('stops waiting for a uv that never exits after being killed', async () => { + // The ceiling only asks uv to stop: killProcTree sends SIGTERM and does not + // wait, and the helper settles on the child's exit, so a uv that ignores the + // signal would hold the launch open past the bound meant to prevent exactly + // that. The wait has to end on its own. + vi.useFakeTimers() + try { + mockInstall.mockImplementationOnce( + // Never settles, however it is signalled. + () => new Promise(() => {}) + ) + const sendOutput = vi.fn() + + const pending = installAgentRequirements(plan, sendOutput) + await vi.advanceTimersByTimeAsync(120_000) + await vi.advanceTimersByTimeAsync(10_000) + + await expect(pending).resolves.toBeUndefined() + expect(sendOutput.mock.calls.join('')).toContain('uv did not stop') + } finally { + vi.useRealTimers() + } + }) + + it('keeps waiting while uv is still within the grace period', async () => { + // The grace must not cut short a uv that is on its way out, or the warning + // would fire on every ordinary cancellation. + vi.useFakeTimers() + try { + mockInstall.mockImplementationOnce( + async (...args: Parameters) => { + const uvSignal = args[6]! + return new Promise((resolve) => { + uvSignal.addEventListener( + 'abort', + () => setTimeout(() => resolve({ code: 1, output: '' }), 2_000), + { once: true } + ) + }) + } + ) + const sendOutput = vi.fn() + + const pending = installAgentRequirements(plan, sendOutput) + await vi.advanceTimersByTimeAsync(120_000) + await vi.advanceTimersByTimeAsync(2_000) + await pending + + const reported = sendOutput.mock.calls.join('') + expect(reported).toContain('starting ComfyUI without it') + expect(reported).not.toContain('uv did not stop') + } finally { + vi.useRealTimers() + } + }) + + it('reports success for an install that finishes inside the grace period', async () => { + // The ceiling can fire while uv is already on its way out with a zero exit. + // Reporting that from the timer rather than the exit code would tell the + // user the agent was skipped on a launch that actually installed it. + vi.useFakeTimers() + try { + mockInstall.mockImplementationOnce( + async (...args: Parameters) => { + const uvSignal = args[6]! + return new Promise((resolve) => { + uvSignal.addEventListener( + 'abort', + () => setTimeout(() => resolve({ code: 0, output: '' }), 1_000), + { once: true } + ) + }) + } + ) + const sendOutput = vi.fn() + + const pending = installAgentRequirements(plan, sendOutput) + await vi.advanceTimersByTimeAsync(120_000) + await vi.advanceTimersByTimeAsync(1_000) + await pending + + const reported = sendOutput.mock.calls.join('') + expect(reported).not.toContain('without it') + expect(reported).not.toContain('uv did not stop') + expect(reported).not.toContain('exited with code') + } finally { + vi.useRealTimers() + } + }) + + it('does not cancel the launch when the ceiling fires', async () => { + // The deadline aborts a controller this module owns, never the launch's own + // signal - the launch must proceed, not report itself cancelled. + vi.useFakeTimers() + try { + const launchAbort = new AbortController() + mockInstall.mockImplementationOnce( + async (...args: Parameters) => { + const uvSignal = args[6]! + return new Promise((resolve) => { + uvSignal.addEventListener('abort', () => resolve({ code: 1, output: '' }), { + once: true + }) + }) + } + ) + + const pending = installAgentRequirements(plan, vi.fn(), launchAbort.signal) + await vi.advanceTimersByTimeAsync(120_000) + await pending + + expect(launchAbort.signal.aborted).toBe(false) + } finally { + vi.useRealTimers() + } + }) + + it('leaves the deadline behind no timer once the install finishes', async () => { + vi.useFakeTimers() + try { + await installAgentRequirements(plan, vi.fn()) + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) + + it('kills uv when the launch itself is cancelled', async () => { + const launchAbort = new AbortController() + let uvSignal: AbortSignal | undefined + mockInstall.mockImplementationOnce( + async (...args: Parameters) => { + uvSignal = args[6] + launchAbort.abort() + return { code: 1, output: '' } + } + ) + + await installAgentRequirements(plan, vi.fn(), launchAbort.signal) + + expect(uvSignal?.aborted).toBe(true) + }) + + it('stays quiet about a thrown install when the launch was cancelled', async () => { + const launchAbort = new AbortController() + mockInstall.mockImplementationOnce(async () => { + launchAbort.abort() + throw new Error('EIO') + }) + const sendOutput = vi.fn() + + await installAgentRequirements(plan, sendOutput, launchAbort.signal) + + expect(sendOutput.mock.calls.join('')).not.toContain('EIO') + }) + + it('stays quiet about the exit code when the launch was cancelled mid-install', async () => { + // Cancelling kills uv, so its non-zero exit IS the cancellation; reporting it + // would put a spurious failure in the output of a launch the user stopped. + const abort = new AbortController() + mockInstall.mockImplementationOnce(async () => { + abort.abort() + return { code: 1, output: '' } + }) + const sendOutput = vi.fn() + + await installAgentRequirements(plan, sendOutput, abort.signal) + + expect(sendOutput.mock.calls.join('')).not.toContain('exited with code') + }) +}) diff --git a/src/main/lib/agentRequirementsLaunch.ts b/src/main/lib/agentRequirementsLaunch.ts new file mode 100644 index 000000000..7ce8b8e87 --- /dev/null +++ b/src/main/lib/agentRequirementsLaunch.ts @@ -0,0 +1,176 @@ +import fs from 'fs' +import path from 'path' +import * as settings from '../settings' +import { installFilteredRequirementsDetailed } from './pip' +import type { UvPipResult } from './pip' +import { withOutputTail } from './logged-process' +import { getActivePythonPath, getActiveUvPath } from './pythonEnv' +import type { InstallationRecord } from '../installations' + +/** The Core flag that starts ComfyUI with the agent enabled. */ +const ENABLE_AGENT_ARG = '--enable-agent' + +/** Requirements file Core ships beside `main.py` for the agent. */ +const AGENT_REQUIREMENTS = 'agent_requirements.txt' + +/** + * Ceiling on how long a launch waits for the install. + * + * The launch blocks on this so the agent is usable on the run that enables it, + * but core starts perfectly well without the packages, so holding a user at the + * launcher indefinitely to acquire them inverts that. Past the ceiling uv is + * killed and the launch goes ahead with the flag still set: core prints its own + * install hint and disables the agent. The matching bound one step earlier is + * the args-schema probe's 15s, which fails open the same way. + * + * The ceiling is generous rather than tight because uv streams nothing while a + * wheel downloads, so a slow transfer is indistinguishable from a stall. A link + * too slow to finish inside it never gets the agent from this path. + */ +const INSTALL_TIMEOUT_MS = 120_000 + +/** + * Grace between asking uv to stop and the launch giving up on it. + * + * Killing is best-effort and not awaited: `killProcTree` sends SIGTERM to the + * process group on POSIX and swallows a failed `taskkill` on Windows, while the + * install settles only on the child's own exit. A uv that does not take the + * signal would therefore hold the launch open past the very ceiling that exists + * to stop that, so the wait is bounded here too and the launch proceeds either + * way. An abandoned uv may still be writing to the environment, which is worth + * one warning line and is strictly better than never starting. + */ +const KILL_GRACE_MS = 10_000 + +/** Which way the bounded wait ended: uv exited, it threw, or the launch stopped + * waiting for a uv that would not stop. */ +type InstallOutcome = + | { kind: 'settled'; result: UvPipResult } + | { kind: 'failed'; error: unknown } + | { kind: 'abandoned' } + +export interface AgentRequirementsInstall { + reqPath: string + uvPath: string + pythonPath: string + installPath: string +} + +/** + * Decide whether this launch installs the agent's Python requirements, and + * resolve everything the install needs. + * + * `args` must be the FINAL spawn args. `--enable-agent` reaches them either + * from the user's own launch args or from a Core beta grant, so the args are + * the one place that knows whether the agent is really starting. A core whose + * schema does not know the flag has normally had it filtered out by then, with + * one exception: when schema discovery itself failed the launch keeps the raw + * args, so a hand-typed flag survives to a core that cannot parse it. That + * launch fails on argparse either way; the cost is one wasted install. + * + * The environment test is the one every `manager_requirements.txt` site makes, + * which is what confines this to the Desktop-managed installs (standalone and + * adopted). Portable, git and build installs drive a Python environment this + * helper has no uv binary for; Core's own `pip install -r` hint covers them. + */ +export function planAgentRequirementsInstall( + installation: InstallationRecord, + args: readonly string[] +): AgentRequirementsInstall | null { + if (!args.includes(ENABLE_AGENT_ARG)) return null + const reqPath = path.join(installation.installPath, 'ComfyUI', AGENT_REQUIREMENTS) + if (!fs.existsSync(reqPath)) return null + const uvPath = getActiveUvPath(installation) + const pythonPath = getActivePythonPath(installation) + if (!pythonPath || !fs.existsSync(uvPath)) return null + return { reqPath, uvPath, pythonPath, installPath: installation.installPath } +} + +/** + * Install the planned requirements, streaming uv's output into the launch. + * + * Bounded and fail-open. Never throws and reports nothing back: neither a + * failure nor a timeout here may stop the launch. ComfyUI still starts with the + * flag, prints its own install hint and disables the agent itself, which beats + * refusing to start. + * + * `signal` is the launch's own, and stays untouched: uv is driven through a + * controller owned here, so the deadline ends the wait without cancelling the + * launch. Because killing uv is best-effort, the wait is bounded twice over: + * the ceiling asks it to stop, and `KILL_GRACE_MS` later the launch stops + * waiting whether or not it did. Nothing downstream depends on which happened. + */ +export async function installAgentRequirements( + plan: AgentRequirementsInstall, + sendOutput: (text: string) => void, + signal?: AbortSignal +): Promise { + sendOutput('\nInstalling agent requirements…\n') + const uvAbort = new AbortController() + const onLaunchAbort = (): void => uvAbort.abort() + let timedOut = false + let graceTimer: ReturnType | undefined + let abandon = (): void => {} + const abandoned = new Promise((resolve) => { + abandon = () => resolve({ kind: 'abandoned' }) + }) + // Armed on whichever side raised the abort, so neither the ceiling nor a user + // cancel can be held open by a uv that never takes the signal. + uvAbort.signal.addEventListener( + 'abort', + () => { + graceTimer = setTimeout(abandon, KILL_GRACE_MS) + }, + { once: true } + ) + if (signal?.aborted) uvAbort.abort() + else signal?.addEventListener('abort', onLaunchAbort, { once: true }) + const deadline = setTimeout(() => { + timedOut = true + uvAbort.abort() + }, INSTALL_TIMEOUT_MS) + + // Settled into a value rather than awaited directly: losing the race leaves + // this pending, and a later rejection with nothing awaiting it would surface + // as an unhandled rejection. + const install: Promise = installFilteredRequirementsDetailed( + plan.reqPath, + plan.uvPath, + plan.pythonPath, + plan.installPath, + '.launch-agent-reqs.txt', + sendOutput, + uvAbort.signal, + settings.getMirrorConfig() + ).then( + (result) => ({ kind: 'settled', result }), + (error: unknown) => ({ kind: 'failed', error }) + ) + + try { + const outcome = await Promise.race([install, abandoned]) + // A cancelled launch kills uv mid-install, so whatever it reports is the + // cancellation rather than a failure worth showing. + if (signal?.aborted) return + if (outcome.kind === 'abandoned') { + sendOutput( + `\n⚠ agent requirements install exceeded ${INSTALL_TIMEOUT_MS / 1000}s and uv did not stop; starting ComfyUI anyway\n` + ) + } else if (outcome.kind === 'failed') { + sendOutput(`⚠ ${AGENT_REQUIREMENTS} failed: ${(outcome.error as Error).message}\n`) + } else if (outcome.result.code !== 0) { + // The exit code decides, not the timer that was racing it: an install + // that finished inside the grace period succeeded, however close to the + // ceiling it landed, and must not be reported as skipped. + sendOutput( + timedOut + ? `\n⚠ agent requirements install exceeded ${INSTALL_TIMEOUT_MS / 1000}s; starting ComfyUI without it\n` + : `\n${withOutputTail(`⚠ agent requirements install exited with code ${outcome.result.code}`, outcome.result.output)}\n` + ) + } + } finally { + clearTimeout(deadline) + if (graceTimer !== undefined) clearTimeout(graceTimer) + signal?.removeEventListener('abort', onLaunchAbort) + } +} diff --git a/src/main/lib/ipc/sessionActions/launch.test.ts b/src/main/lib/ipc/sessionActions/launch.test.ts index a1b20285c..953481c8b 100644 --- a/src/main/lib/ipc/sessionActions/launch.test.ts +++ b/src/main/lib/ipc/sessionActions/launch.test.ts @@ -116,6 +116,26 @@ vi.mock('../../coreBetaGrants', async (importOriginal) => { return { ...actual, getCoreBetaGrantsAsync: async () => launchHarness.grants } }) +/** Stands in for the uv subprocess the agent-requirements step drives, so a launch + * under test installs nothing and a test can choose the outcome. */ +const pipHarness = vi.hoisted(() => ({ + calls: [] as unknown[][], + result: { code: 0, output: '' }, + duringInstall: null as null | (() => void) +})) + +vi.mock('../../pip', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + installFilteredRequirementsDetailed: async (...args: unknown[]) => { + pipHarness.calls.push(args) + pipHarness.duringInstall?.() + return pipHarness.result + } + } +}) + vi.mock('../../hardwareTap', async (importOriginal) => { const actual = await importOriginal() return { @@ -162,6 +182,9 @@ import type * as SharedModule from '../shared' import type * as ComfyArgsModule from '../../comfy-args' import type * as CoreBetaGrantsModule from '../../coreBetaGrants' import type * as HardwareTapModule from '../../hardwareTap' +import type * as PipModule from '../../pip' +import { getUvPath, getVenvPythonPath } from '../../pythonEnv' +import { getLogDir } from '../../logRotation' const installOf = (sourceId: string) => ({ sourceId }) as InstallationRecord @@ -1477,3 +1500,209 @@ describe('emitCoreBetaTelemetry', () => { ]) }) }) + +describe('agent requirements at launch', () => { + const AGENT_GRANT: CoreBetaGrant = { arg: '--enable-agent', minCoreVersion: '0.3.80' } + let installDir = '' + let sent: string[] = [] + let progress: { phase: string; steps?: { phase: string }[] }[] = [] + let spawnArgs: string[] = [] + let spawned = 0 + + /** Create a file (and its parents) the way the real install layout has it. */ + const touch = (target: string): void => { + fs.mkdirSync(path.dirname(target), { recursive: true }) + fs.writeFileSync(target, '') + } + + const agentReqPath = (): string => path.join(installDir, 'ComfyUI', 'agent_requirements.txt') + + const harnessInstall = (): InstallationRecord => + ({ + id: 'agent-reqs-inst', + name: 'Agent Harness', + sourceId: 'harness-source', + installPath: installDir, + version: '0.3.81', + comfyVersion: { + commit: '61e5e3b5a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4', + baseTag: 'v0.3.81', + commitsAhead: 0, + baseTagVerified: true + } + }) as unknown as InstallationRecord + + const ctxFor = (installationId: string): ActionContext => ({ + event: { + sender: { + isDestroyed: () => false, + send: (channel: string, payload: { text?: string; phase?: string }) => { + if (typeof payload?.text === 'string') sent.push(payload.text) + if (channel === 'install-progress' && typeof payload?.phase === 'string') { + progress.push(payload as { phase: string; steps?: { phase: string }[] }) + } + } + } + } as unknown as Electron.IpcMainInvokeEvent, + installationId, + inst: harnessInstall(), + actionData: {} + }) + + beforeEach(() => { + installDir = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-reqs-launch-')) + fs.mkdirSync(path.join(installDir, 'ComfyUI'), { recursive: true }) + // The managed pair the shared uv helper targets, plus the file Core ships. + touch(getUvPath(installDir)) + touch(getVenvPythonPath(installDir)) + fs.writeFileSync(agentReqPath(), 'comfyui-agent==1.0.0\n') + sent = [] + progress = [] + spawnArgs = [] + spawned = 0 + pipHarness.calls = [] + pipHarness.result = { code: 0, output: '' } + pipHarness.duringInstall = null + launchHarness.schemaThrows = false + launchHarness.registryThrows = false + launchHarness.betaEnabled = true + launchHarness.betaEnabledThrows = false + launchHarness.schemaNames = ['enable-agent', 'listen', 'feature-flag'] + launchHarness.grants = [] + launchHarness.duringResourceAcquire = null + launchHarness.spawn = (_cmd: unknown, args: unknown) => { + spawned += 1 + spawnArgs = args as string[] + const proc = new EventEmitter() as FakeChild + proc.stdout = new EventEmitter() + proc.stderr = new EventEmitter() + proc.pid = 4243 + proc.killed = false + proc.kill = () => true + return proc + } + launchHarness.launchCommand = { + cmd: process.execPath, + args: ['-s', path.join(installDir, 'ComfyUI', 'main.py'), '--enable-agent', '--listen'], + cwd: installDir, + skipPortWait: true + } + vi.spyOn(telemetry, 'emit').mockImplementation((() => {}) as unknown as typeof telemetry.emit) + }) + + afterEach(async () => { + // The launch log stream opens asynchronously after the spawn: removing the + // install dir first fails that open with an ENOENT nothing is listening for. + // The file appears exactly when the open lands, so waiting on it is the + // teardown's own gate rather than a guess at how long the open takes. + if (spawned > 0) { + const logFile = path.join(getLogDir(installDir), 'comfyui.log') + await vi.waitFor(() => expect(fs.existsSync(logFile)).toBe(true)) + } + vi.restoreAllMocks() + fs.rmSync(installDir, { recursive: true, force: true }) + }) + + it('installs the requirements when the user typed the flag by hand', async () => { + const res = await handleLaunch(ctxFor('agent-reqs-hand-typed')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toHaveLength(1) + expect(pipHarness.calls[0]![0]).toBe(agentReqPath()) + expect(pipHarness.calls[0]![1]).toBe(getUvPath(installDir)) + expect(pipHarness.calls[0]![2]).toBe(getVenvPythonPath(installDir)) + expect(spawnArgs).toContain('--enable-agent') + }) + + it('installs the requirements when a beta grant added the flag', async () => { + // The grant is the only source of the flag here: the install's own args have none. + launchHarness.grants = [AGENT_GRANT] + launchHarness.launchCommand = { + cmd: process.execPath, + args: ['-s', path.join(installDir, 'ComfyUI', 'main.py'), '--listen'], + cwd: installDir, + skipPortWait: true + } + + const res = await handleLaunch(ctxFor('agent-reqs-granted')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toHaveLength(1) + expect(spawnArgs).toContain('--enable-agent') + }) + + it('installs nothing when the flag never reaches the final args', async () => { + launchHarness.launchCommand = { + cmd: process.execPath, + args: ['-s', path.join(installDir, 'ComfyUI', 'main.py'), '--listen'], + cwd: installDir, + skipPortWait: true + } + + const res = await handleLaunch(ctxFor('agent-reqs-flag-absent')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toEqual([]) + }) + + it('installs nothing when the running core cannot parse the flag', async () => { + // A core predating `--enable-agent` has it filtered out of the final args, + // so there is nothing to install for and nothing reaches the spawn either. + launchHarness.schemaNames = ['listen', 'feature-flag'] + + const res = await handleLaunch(ctxFor('agent-reqs-unsupported-core')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toEqual([]) + expect(spawnArgs).not.toContain('--enable-agent') + }) + + it('installs nothing when core ships no agent requirements file', async () => { + fs.rmSync(agentReqPath()) + + const res = await handleLaunch(ctxFor('agent-reqs-file-absent')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toEqual([]) + expect(spawnArgs).toContain('--enable-agent') + }) + + it('installs nothing for an install with no Desktop-managed Python environment', async () => { + fs.rmSync(getUvPath(installDir)) + + const res = await handleLaunch(ctxFor('agent-reqs-unmanaged-env')) + + expect(res.ok).toBe(true) + expect(pipHarness.calls).toEqual([]) + expect(spawnArgs).toContain('--enable-agent') + }) + + it('launches with the flag still set when the install fails', async () => { + pipHarness.result = { code: 1, output: 'No solution found\n' } + + const res = await handleLaunch(ctxFor('agent-reqs-install-failed')) + + expect(res.ok).toBe(true) + expect(spawnArgs).toContain('--enable-agent') + expect(sent.join('')).toContain('exited with code 1') + }) + + it('publishes the install as a step the renderer can show', async () => { + // The renderer drops progress for a phase absent from the steps payload, so + // the payload is what makes the step visible at all. + await handleLaunch(ctxFor('agent-reqs-step-published')) + + const lastSteps = progress.filter((p) => p.phase === 'steps').at(-1) + expect(lastSteps?.steps?.map((s) => s.phase)).toContain('agentRequirements') + expect(progress.some((p) => p.phase === 'agentRequirements')).toBe(true) + }) + + it('cancels the launch without spawning when it is aborted mid-install', async () => { + pipHarness.duringInstall = () => _operationAborts.get('agent-reqs-cancelled')?.abort() + + const res = await handleLaunch(ctxFor('agent-reqs-cancelled')) + + expect(res).toEqual({ ok: false, cancelled: true }) + expect(spawned).toBe(0) + }) +}) diff --git a/src/main/lib/ipc/sessionActions/launch.ts b/src/main/lib/ipc/sessionActions/launch.ts index 164dda849..aa3f93d3f 100644 --- a/src/main/lib/ipc/sessionActions/launch.ts +++ b/src/main/lib/ipc/sessionActions/launch.ts @@ -65,7 +65,7 @@ import { createAssetsTap } from '../../assetsTap' import { createExecutionTap } from '../../executionTap' import { createHardwareTap } from '../../hardwareTap' import { createLaunchProgressTracker } from '../../launchProgress' -import { buildLaunchPhases } from '../../launchPhases' +import { buildLaunchPhases, AGENT_REQUIREMENTS_PHASE } from '../../launchPhases' import { getTemplateDownloadState, summarizeTemplateState, @@ -90,6 +90,10 @@ import { } from '../../bootPhaseBuffer' import { appendLog } from '../../logsBroadcast' import { reconcileManagerConfigForLaunch } from '../../managerConfigLaunch' +import { + installAgentRequirements, + planAgentRequirementsInstall +} from '../../agentRequirementsLaunch' import { recoverInterruptedComfyOp } from '../../opMarker' import { waitLaunchSpawnHold } from '../../e2eOverrides' import { migrateEnvLayout } from '../../../sources/standalone/install' @@ -1015,6 +1019,28 @@ async function runLaunch( return { ok: false, message: i18n.t('errors.managerConfigWriteFailed') } } + // The agent flag is final here (only path args are appended after this), so + // this is the first point that knows the agent is actually starting - whether + // the user typed the flag or a beta grant added it, and whether the running + // core can parse it at all. The package is tens of megabytes, so it gets its + // own launch step. Bounded and fail-open: a failure or a timeout is reported + // in the launch output and the flag is kept, leaving core to print its install + // hint and disable the agent itself. + // + // The step is added through `addLatePhase` rather than `preLaunchPhases` + // because a torch repair may already have armed the tracker, freezing that list. + const agentRequirements = planAgentRequirementsInstall(inst, launchCmd.args ?? []) + if (agentRequirements) { + const tracker = await armLaunchTracker() + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + await installAgentRequirements( + agentRequirements, + makeSendOutput(sender, installationId), + abort.signal + ) + if (abort.signal.aborted) return { ok: false, cancelled: true } + } + const { preLaunchExtras, manageModelFolders, modelDirsForLaunch, modelSyncOptions } = applyStorageLaunchArgs(inst, installationId, launchCmd) diff --git a/src/main/lib/launchPhases.ts b/src/main/lib/launchPhases.ts index ca03c2cb4..da86cd6fa 100644 --- a/src/main/lib/launchPhases.ts +++ b/src/main/lib/launchPhases.ts @@ -99,6 +99,17 @@ const PRE_LAUNCH_PHASES: Record = { torchRepair: { phase: 'torchRepair', match: NEVER, weight: 0.1, streaming: true } } +/** The agent requirements install. Not one of `PRE_LAUNCH_PHASES`: whether it + * runs is only known once the launch args are final, which is after a torch + * repair may already have armed the tracker and frozen the phase list. It is + * handed to `addLatePhase` instead of being injected here. */ +export const AGENT_REQUIREMENTS_PHASE: LaunchPhaseDef = { + phase: 'agentRequirements', + match: NEVER, + weight: 0.1, + streaming: true +} + /** Starter-template model download, shown as the LAST launch step. Synthetic + * streaming: its bytes downloaded in the background since install-begin, and a * 500 ms reader in `handleLaunch` feeds the rich substatus from the shared diff --git a/src/main/lib/launchProgress.test.ts b/src/main/lib/launchProgress.test.ts index 5c17bfea1..ffc80d141 100644 --- a/src/main/lib/launchProgress.test.ts +++ b/src/main/lib/launchProgress.test.ts @@ -2,7 +2,7 @@ import { describe, it, expect } from 'vitest' import fs from 'fs' import path from 'path' import { createLaunchProgressTracker } from './launchProgress' -import { DEFAULT_LAUNCH_PHASES, buildLaunchPhases } from './launchPhases' +import { DEFAULT_LAUNCH_PHASES, buildLaunchPhases, AGENT_REQUIREMENTS_PHASE } from './launchPhases' const FIXTURE = fs.readFileSync( path.join(__dirname, '__fixtures__', 'launch', 'first-run.log'), @@ -259,3 +259,100 @@ describe('buildLaunchPhases — extensibility', () => { expect(order).toContain('mountLibraries') }) }) + +describe('addLatePhase for work discovered after the tracker was armed', () => { + /** Arm a tracker over `phases` and return it with its emit log. */ + function armed(phases = buildLaunchPhases({}, { preLaunchPhases: ['torchRepair'] })): { + tracker: ReturnType + emits: Emit[] + } { + const emits: Emit[] = [] + const tracker = createLaunchProgressTracker({ + phases, + sendProgress: (phase, detail) => emits.push({ phase, ...detail }) + }) + tracker.start() + return { tracker, emits } + } + + it('re-emits the steps payload so the renderer knows the new phase', () => { + // Without the re-emit the renderer drops the phase's progress outright: + // progressStore returns early when the phase is absent from `steps`. + const { tracker, emits } = armed() + const initial = emits + .filter((e) => e.phase === 'steps') + .at(-1)! + .steps!.map((s) => s.phase) + expect(initial).not.toContain('agentRequirements') + + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + const stepPayloads = emits.filter((e) => e.phase === 'steps') + expect(stepPayloads).toHaveLength(2) + expect(stepPayloads[1]!.steps?.map((s) => s.phase)).toContain('agentRequirements') + }) + + it('enters the new phase, so it is the active step while the work runs', () => { + const { tracker, emits } = armed() + const before = emits.length + + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + const after = emits.slice(before).filter((e) => e.phase !== 'steps') + expect(after.map((e) => e.phase)).toEqual(['agentRequirements']) + expect(after[0]!.percent).toBe(-1) + }) + + it('inserts directly after the active phase so the bar cannot regress', () => { + // torchRepair is active (phase 0). The new phase must land at 1, ahead of + // it and ahead of nothing already completed. + const { tracker, emits } = armed() + + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + const published = emits + .filter((e) => e.phase === 'steps') + .at(-1)! + .steps!.map((s) => s.phase) + expect(published.slice(0, 2)).toEqual(['torchRepair', 'agentRequirements']) + expect(published).toEqual([ + 'torchRepair', + 'agentRequirements', + ...DEFAULT_LAUNCH_PHASES.map((p) => p.phase) + ]) + }) + + it('still advances into the real boot phases afterwards', () => { + const { tracker, emits } = armed() + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + tracker.ingest('Total VRAM 24576 MB\n') + + expect(phaseOrder(emits).at(-1)).toBe('gpu') + }) + + it('works when nothing else injected a phase (the common launch)', () => { + const { tracker, emits } = armed(buildLaunchPhases({})) + + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + const published = emits + .filter((e) => e.phase === 'steps') + .at(-1)! + .steps!.map((s) => s.phase) + // launchStart is active at index 0, so the install lands right after it. + expect(published.slice(0, 2)).toEqual(['launchStart', 'agentRequirements']) + expect(phaseOrder(emits).at(-1)).toBe('agentRequirements') + }) + + it("does not mutate the caller's phase array", () => { + const phases = buildLaunchPhases({}) + const { tracker } = armed(phases) + const lengthBefore = phases.length + + tracker.addLatePhase(AGENT_REQUIREMENTS_PHASE) + + expect(phases).toHaveLength(lengthBefore) + expect(phases.map((p) => p.phase)).not.toContain('agentRequirements') + }) +}) diff --git a/src/main/lib/launchProgress.ts b/src/main/lib/launchProgress.ts index 273f048df..0c6bb2c02 100644 --- a/src/main/lib/launchProgress.ts +++ b/src/main/lib/launchProgress.ts @@ -73,6 +73,12 @@ export interface LaunchProgressTracker { start: () => void /** Feed a stdout/stderr chunk. Safe to call with partial lines. */ ingest: (chunk: string) => void + /** Add a phase for work that only became known after the tracker was armed, + * and enter it. The steps payload is re-emitted, without which the renderer + * drops the phase's progress entirely (it ignores a phase absent from the + * payload). Inserted directly after the active phase, so the bar advances + * into a new slot rather than regressing. */ + addLatePhase: (def: LaunchPhaseDef) => void /** Restart per-attempt `onPhaseEnter` observation for a boot retry. The UI * phase index stays monotonic (progress never regresses), but the retried * process re-logs its boot from the top - without this reset those re-hit @@ -95,7 +101,10 @@ export function createLaunchProgressTracker(opts: { * break progress, so it is swallowed. */ onPhaseEnter?: (phase: string) => void }): LaunchProgressTracker { - const { phases, sendProgress, onPhaseEnter } = opts + const { sendProgress, onPhaseEnter } = opts + // Copied, not aliased: `addLatePhase` splices into this list, and the caller's + // array must not change under it. + const phases = [...opts.phases] const nodeCount = opts.nodeCount && opts.nodeCount > 0 ? opts.nodeCount : 0 // Index of the currently-active phase; -1 until the first milestone. @@ -252,6 +261,13 @@ export function createLaunchProgressTracker(opts: { const lines = pending.split(/\r?\n/) pending = lines.pop() ?? '' for (const line of lines) handleLine(line) + }, + addLatePhase(def: LaunchPhaseDef): void { + const idx = activeIdx + 1 + phases.splice(idx, 0, { ...def }) + stepsSent = false + emitSteps() + enterPhase(idx) } } }