Stop idle Comfy after 30 minutes and add a Wake GPU poke.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-29 17:59:18 -05:00
parent cdc69d2bd6
commit ad1f4f905a
8 changed files with 283 additions and 26 deletions
@@ -10,6 +10,16 @@ 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(/\/$/, '')
+const idleMs = Math.max(60_000, Number(process.env.COMFY_IDLE_MS || 30 * 60 * 1000) || 30 * 60 * 1000)
+
+let lastWorkAt = Date.now()
+let lastQueueRunning = 0
+let stoppedByAgent = false
+
+function markWork() {
+  lastWorkAt = Date.now()
+  stoppedByAgent = false
+}

 function json(res, status, body) {
  const payload = JSON.stringify(body)
@@ -110,7 +120,11 @@ function ensureProxy(targetPort) {
    }
    client.on('error', fail)
    upstream.on('error', fail)
-    client.pipe(upstream)
+    client.once('data', (chunk) => {
+      if (isWorkHttp(chunk)) markWork()
+      upstream.write(chunk)
+      client.pipe(upstream)
+    })
    upstream.pipe(client)
  })
  proxyServer.on('error', (error) => {
@@ -151,6 +165,7 @@ function launchCommand() {
 }

 function startComfy() {
+  markWork()
  const command = launchCommand()
  const cwd = process.env.COMFY_LAUNCH_CWD || undefined
  const child = spawn(command, { shell: true, detached: true, stdio: 'ignore', windowsHide: true, cwd })
@@ -158,6 +173,96 @@ function startComfy() {
  return { started: true, command, cwd: cwd || null }
 }

+async function fetchLocalQueue(portNum) {
+  if (!portNum) return { running: 0, pending: 0, ok: false }
+  try {
+    const res = await fetch(`http://127.0.0.1:${portNum}/queue`, { signal: AbortSignal.timeout(2500) })
+    if (!res.ok) return { running: 0, pending: 0, ok: false }
+    const payload = await res.json()
+    return {
+      running: Array.isArray(payload?.queue_running) ? payload.queue_running.length : 0,
+      pending: Array.isArray(payload?.queue_pending) ? payload.queue_pending.length : 0,
+      ok: true
+    }
+  } catch {
+    return { running: 0, pending: 0, ok: false }
+  }
+}
+
+function isWorkHttp(chunk) {
+  const head = chunk.toString('latin1', 0, Math.min(chunk.length, 256))
+  const line = head.split('\r\n')[0] || ''
+  const match = line.match(/^(GET|HEAD|POST|PUT|PATCH|DELETE)\s+(\S+)/i)
+  if (!match) return false
+  const method = match[1].toUpperCase()
+  let path = String(match[2] || '').split('?')[0] || ''
+  try { path = decodeURIComponent(path) } catch { /* keep raw */ }
+  if (method === 'GET' || method === 'HEAD') return false
+  if (/^\/(system_stats|queue|object_info|internal)\b/i.test(path)) return false
+  return true
+}
+
+async function taskkillImage(image) {
+  try {
+    await execFileAsync('taskkill', ['/IM', image, '/F', '/T'], { windowsHide: true, timeout: 8000 })
+    return true
+  } catch {
+    return false
+  }
+}
+
+async function stopComfyProcesses() {
+  const killed = []
+  if (await taskkillImage('Comfy Desktop.exe')) killed.push('Comfy Desktop.exe')
+  if (await taskkillImage('ComfyUI.exe')) killed.push('ComfyUI.exe')
+  try {
+    const { stdout } = await execFileAsync('powershell.exe', [
+      '-NoProfile', '-Command',
+      "Get-CimInstance Win32_Process | Where-Object { $_.CommandLine -match 'ComfyUI\\\\main.py|ComfyUI/main.py' } | ForEach-Object { Stop-Process -Id $_.ProcessId -Force -ErrorAction SilentlyContinue; $_.ProcessId }"
+    ], { windowsHide: true, timeout: 8000 })
+    if (String(stdout).trim()) killed.push('main.py')
+  } catch { /* ignore */ }
+  return killed
+}
+
+async function noteQueue(portNum) {
+  const queue = await fetchLocalQueue(portNum)
+  if (!queue.ok) return queue
+  if (queue.running > 0 || queue.pending > 0) {
+    markWork()
+  } else if (lastQueueRunning > 0 && queue.running === 0) {
+    lastWorkAt = Date.now()
+  }
+  lastQueueRunning = queue.running
+  return queue
+}
+
+async function maybeIdleStop(healthyPort) {
+  if (healthyPort) {
+    stoppedByAgent = false
+    const queue = await noteQueue(healthyPort)
+    if (!queue.ok) return
+    if (queue.running > 0 || queue.pending > 0) return
+  } else if (await processUp()) {
+    stoppedByAgent = false
+  } else {
+    stoppedByAgent = true
+    return
+  }
+  if (Date.now() - lastWorkAt < idleMs) return
+  console.log(JSON.stringify({
+    src: 'comfy-host-agent',
+    event: 'idle-stop',
+    idleMs,
+    lastActivityAt: new Date(lastWorkAt).toISOString(),
+    http: Boolean(healthyPort)
+  }))
+  const killed = await stopComfyProcesses()
+  stoppedByAgent = true
+  lastQueueRunning = 0
+  console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'stopped', reason: 'idle', killed }))
+}
+
 function sharedRoots() {
  const local = process.env.LOCALAPPDATA || ''
  return {
@@ -274,10 +379,16 @@ const server = http.createServer(async (req, res) => {
  const url = new URL(req.url || '/', 'http://localhost')
  if (req.method === 'GET' && url.pathname === '/status') {
    const [healthyPort, process] = await Promise.all([syncProxy(), processUp()])
+    if (healthyPort) await noteQueue(healthyPort)
+    const processUpNow = Boolean(process || healthyPort)
    return json(res, 200, {
      ok: true,
      http: Boolean(healthyPort),
-      process,
+      process: processUpNow,
+      processUp: processUpNow,
+      asleep: Boolean(stoppedByAgent || (!healthyPort && !process)),
+      lastActivityAt: new Date(lastWorkAt).toISOString(),
+      idleMs,
      port: healthyPort || null,
      proxyPort: healthyPort ? proxyPort : null
    })
@@ -285,10 +396,28 @@ const server = http.createServer(async (req, res) => {
  if (req.method === 'POST' && url.pathname === '/start') {
    const healthyPort = await syncProxy()
    if (healthyPort) {
-      return json(res, 200, { ok: true, started: false, already: true, port: healthyPort, proxyPort })
+      markWork()
+      return json(res, 200, { ok: true, started: false, already: true, asleep: false, port: healthyPort, proxyPort })
    }
    const launched = startComfy()
-    return json(res, 200, { ok: true, ...launched })
+    return json(res, 200, { ok: true, asleep: false, ...launched })
+  }
+  if (req.method === 'POST' && url.pathname === '/stop') {
+    const healthyPort = await syncProxy()
+    const queue = healthyPort ? await fetchLocalQueue(healthyPort) : { running: 0, pending: 0, ok: false }
+    if (queue.ok && (queue.running > 0 || queue.pending > 0)) {
+      return json(res, 409, {
+        ok: false,
+        error: 'queue-busy',
+        running: queue.running,
+        pending: queue.pending
+      })
+    }
+    const killed = await stopComfyProcesses()
+    stoppedByAgent = true
+    lastQueueRunning = 0
+    console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'stopped', reason: 'stop', killed }))
+    return json(res, 200, { ok: true, stopped: true, asleep: true, killed })
  }
  if (req.method === 'POST' && url.pathname === '/purge') {
    const body = await readJson(req)
@@ -299,15 +428,23 @@ const server = http.createServer(async (req, res) => {

 server.listen(port, '0.0.0.0', async () => {
  const healthyPort = await syncProxy()
+  if (healthyPort) await noteQueue(healthyPort)
  console.log(JSON.stringify({
    src: 'comfy-host-agent',
    event: 'listen',
    port,
    proxyPort,
    comfyPort: healthyPort || null,
+    idleMs,
    candidates: candidatePorts()
  }))
+  let lastIdleCheck = 0
  setInterval(() => {
-    syncProxy().catch(() => {})
+    syncProxy().then(async (nextPort) => {
+      const now = Date.now()
+      if (now - lastIdleCheck < 10_000) return
+      lastIdleCheck = now
+      await maybeIdleStop(nextPort)
+    }).catch(() => {})
  }, 3000)
 })