aigen/scripts/comfy-host-agent.mjs

import http from 'node:http'
import net from 'node:net'
import { execFile, spawn } from 'node:child_process'
import { promisify } from 'node:util'
import { readdirSync, existsSync, rmSync } from 'node:fs'
import { basename, join, resolve, relative, isAbsolute } from 'node:path'

const execFileAsync = promisify(execFile)
const port = Number(process.env.COMFY_CONTROL_PORT || 8199)
const proxyPort = Number(process.env.COMFY_PROXY_PORT || 8198)
const token = process.env.COMFY_CONTROL_TOKEN || ''
const defaultHttp = (process.env.COMFY_HOST || 'http://127.0.0.1:8188').replace(/\/$/, '')

function json(res, status, body) {
  const payload = JSON.stringify(body)
  res.writeHead(status, { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(payload) })
  res.end(payload)
}

function authorized(req) {
  if (!token) return true
  const header = String(req.headers.authorization || '')
  return header === `Bearer ${token}`
}

function candidatePorts() {
  const ports = new Set([8188, 8189])
  try {
    const url = new URL(defaultHttp.includes('://') ? defaultHttp : `http://${defaultHttp}`)
    if (url.port) ports.add(Number(url.port))
  } catch { /* ignore */ }
  try {
    const dir = join(process.env.APPDATA || '', 'Comfy Desktop', 'port-locks')
    for (const name of readdirSync(dir)) {
      const match = name.match(/^port-(\d+)\.json$/)
      if (match) ports.add(Number(match[1]))
    }
  } catch { /* ignore */ }
  return [...ports]
}

async function probe(portNum) {
  try {
    const res = await fetch(`http://127.0.0.1:${portNum}/system_stats`, { signal: AbortSignal.timeout(2500) })
    return res.ok
  } catch {
    return false
  }
}

async function findHealthyPort() {
  for (const next of candidatePorts()) {
    if (await probe(next)) return next
  }
  return 0
}

let proxyServer = null
let proxyTarget = 0

function ensureProxy(targetPort) {
  if (!targetPort) return
  if (proxyServer && proxyTarget === targetPort) return
  proxyTarget = targetPort
  if (proxyServer) {
    try { proxyServer.close() } catch { /* ignore */ }
    proxyServer = null
  }
  proxyServer = net.createServer((client) => {
    const upstream = net.connect(proxyTarget, '127.0.0.1')
    const fail = () => {
      try { client.destroy() } catch { /* ignore */ }
      try { upstream.destroy() } catch { /* ignore */ }
    }
    client.on('error', fail)
    upstream.on('error', fail)
    client.pipe(upstream)
    upstream.pipe(client)
  })
  proxyServer.on('error', (error) => {
    console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
  })
  proxyServer.listen(proxyPort, '0.0.0.0', () => {
    console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy', listen: proxyPort, target: proxyTarget }))
  })
}

async function processUp() {
  try {
    const { stdout } = await execFileAsync('tasklist', ['/FO', 'CSV', '/NH'], { windowsHide: true, timeout: 4000 })
    if (/Comfy Desktop|ComfyUI/i.test(stdout)) return true
  } catch { /* ignore */ }
  try {
    const { stdout } = await execFileAsync('powershell.exe', [
      '-NoProfile', '-Command',
      "Get-CimInstance Win32_Process | Where-Object { $_.CommandLine -match 'ComfyUI\\\\main.py|ComfyUI/main.py' } | Select-Object -First 1 -ExpandProperty ProcessId"
    ], { windowsHide: true, timeout: 6000 })
    return Boolean(String(stdout).trim())
  } catch {
    return false
  }
}

function launchCommand() {
  if (process.env.COMFY_LAUNCH_CMD) return process.env.COMFY_LAUNCH_CMD
  const local = process.env.LOCALAPPDATA || ''
  const desktop = `${local}\\Programs\\Comfy Desktop\\Comfy Desktop.exe`
  return `start "" "${desktop}"`
}

function startComfy() {
  const command = launchCommand()
  const cwd = process.env.COMFY_LAUNCH_CWD || undefined
  const child = spawn(command, { shell: true, detached: true, stdio: 'ignore', windowsHide: true, cwd })
  child.unref()
  return { started: true, command, cwd: cwd || null }
}

function sharedRoots() {
  const local = process.env.LOCALAPPDATA || ''
  return {
    input: process.env.COMFY_INPUT_DIR || join(local, 'Comfy-Desktop', 'ComfyUI-Shared', 'input'),
    output: process.env.COMFY_OUTPUT_DIR || join(local, 'Comfy-Desktop', 'ComfyUI-Shared', 'output')
  }
}

function inputRoots() {
  const local = process.env.LOCALAPPDATA || ''
  const roots = [
    process.env.COMFY_INPUT_DIR || join(local, 'Comfy-Desktop', 'ComfyUI-Shared', 'input')
  ]
  const installs = join(local, 'Comfy-Desktop', 'ComfyUI-Installs')
  try {
    for (const dir of readdirSync(installs, { withFileTypes: true })) {
      if (!dir.isDirectory()) continue
      roots.push(join(installs, dir.name, 'ComfyUI', 'input'))
    }
  } catch { /* ignore */ }
  return roots
}

function safeFile(root, subfolder, filename) {
  const name = basename(String(filename || ''))
  if (!name || name === '.' || name === '..') return null
  const rootAbs = resolve(root)
  const target = resolve(rootAbs, String(subfolder || '').replace(/\.\./g, ''), name)
  const rel = relative(rootAbs, target)
  if (!rel || rel.startsWith('..') || isAbsolute(rel)) return null
  return target
}

function removeFile(path) {
  if (!path || !existsSync(path)) return false
  rmSync(path, { force: true })
  return true
}

function readJson(req) {
  return new Promise((resolve) => {
    const chunks = []
    req.on('data', (chunk) => chunks.push(chunk))
    req.on('end', () => {
      try {
        resolve(JSON.parse(Buffer.concat(chunks).toString('utf8') || '{}'))
      } catch {
        resolve({})
      }
    })
    req.on('error', () => resolve({}))
  })
}

function purgeDesktopFiles(body) {
  const roots = sharedRoots()
  const deleted = []
  const imageName = String(body?.imageName || body?.filename || '')
  const imageSub = String(body?.imageSubfolder || '')
  if (imageName) {
    for (const root of inputRoots()) {
      if (!existsSync(root)) continue
      const nested = safeFile(root, imageSub, imageName)
      if (removeFile(nested)) deleted.push(nested)
      if (imageSub) {
        const flat = safeFile(root, '', imageName)
        if (removeFile(flat)) deleted.push(flat)
      }
    }
  }
  const video = body?.video || null
  const videoName = String(video?.filename || '')
  if (videoName) {
    const outputPath = safeFile(roots.output, video.subfolder || 'video', videoName)
    if (removeFile(outputPath)) deleted.push(outputPath)
  }
  console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'purge', deleted: deleted.length, files: deleted.map((path) => basename(path)), imageName }))
  return { ok: true, deleted }
}

const server = http.createServer(async (req, res) => {
  if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
  const url = new URL(req.url || '/', 'http://localhost')
  if (req.method === 'GET' && url.pathname === '/status') {
    const [healthyPort, process] = await Promise.all([findHealthyPort(), processUp()])
    if (healthyPort) ensureProxy(healthyPort)
    return json(res, 200, {
      ok: true,
      http: Boolean(healthyPort),
      process,
      port: healthyPort || null,
      proxyPort: healthyPort ? proxyPort : null
    })
  }
  if (req.method === 'POST' && url.pathname === '/start') {
    const healthyPort = await findHealthyPort()
    if (healthyPort) {
      ensureProxy(healthyPort)
      return json(res, 200, { ok: true, started: false, already: true, port: healthyPort, proxyPort })
    }
    const launched = startComfy()
    return json(res, 200, { ok: true, ...launched })
  }
  if (req.method === 'POST' && url.pathname === '/purge') {
    const body = await readJson(req)
    return json(res, 200, purgeDesktopFiles(body))
  }
  json(res, 404, { ok: false, error: 'not found' })
})

server.listen(port, '0.0.0.0', async () => {
  const healthyPort = await findHealthyPort()
  if (healthyPort) ensureProxy(healthyPort)
  console.log(JSON.stringify({
    src: 'comfy-host-agent',
    event: 'listen',
    port,
    proxyPort,
    comfyPort: healthyPort || null,
    candidates: candidatePorts()
  }))
})