Refuse to start Comfy while AITraining has the GPU.

Co-authored-by: Cursor <cursoragent@cursor.com>

scripts/comfy-host-agent.mjs

@@ -12,18 +12,43 @@ 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)
+const trainingUrl = String(process.env.TRAINING_CONTROL_URL || 'http://127.0.0.1:8200').replace(/\/$/, '')
 
 let lastWorkAt = Date.now()
 let lastQueueRunning = 0
 let stoppedByAgent = false
 let lastHealthyPort = 0
 let lastProcessUp = false
+let lastTraining = { busy: false, jobId: null, status: null, name: null, message: '' }
 
 function markWork() {
   lastWorkAt = Date.now()
   stoppedByAgent = false
 }
 
+async function trainingLock() {
+  try {
+    const res = await fetch(`${trainingUrl}/status`, { signal: AbortSignal.timeout(2500) })
+    const body = await res.json().catch(() => null)
+    const job = body?.job
+    const status = String(job?.status || '')
+    const busy = Boolean(job && ['preparing', 'queued', 'running', 'stopping'].includes(status))
+    lastTraining = {
+      busy,
+      jobId: job?.id || null,
+      status: busy ? status : null,
+      name: busy ? (job?.outputName || null) : null,
+      message: busy
+        ? `AITraining is ${status}${job?.outputName ? ` (${job.outputName})` : ''}. Stop that job before poking Comfy.`
+        : ''
+    }
+    return lastTraining
+  } catch {
+    lastTraining = { busy: false, jobId: null, status: null, name: null, message: '' }
+    return lastTraining
+  }
+}
+
 function json(res, status, body) {
   const payload = JSON.stringify(body)
   res.writeHead(status, { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(payload) })
@@ -603,10 +628,20 @@ const server = http.createServer(async (req, res) => {
       lastActivityAt: new Date(lastWorkAt).toISOString(),
       idleMs,
       port: healthyPort || null,
-      proxyPort: healthyPort ? proxyPort : null
+      proxyPort: healthyPort ? proxyPort : null,
+      training: lastTraining
     })
   }
   if (req.method === 'POST' && url.pathname === '/start') {
+    const training = await trainingLock()
+    if (training.busy) {
+      return json(res, 409, {
+        ok: false,
+        error: 'train-busy',
+        message: training.message,
+        training
+      })
+    }
     const healthyPort = await syncProxy()
     if (healthyPort) {
       markWork()
@@ -663,6 +698,8 @@ server.listen(port, '0.0.0.0', async () => {
     idleMs,
     candidates: candidatePorts()
   }))
+  void trainingLock()
+  setInterval(() => { void trainingLock() }, 3000)
   let lastIdleCheck = 0
   setInterval(() => {
     syncProxy().then(async (nextPort) => {

server/utils/comfyLifecycle.ts

@@ -109,6 +109,8 @@ async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) {
       stopped?: boolean
       killed?: unknown
       error?: string
+      message?: string
+      training?: { busy?: boolean; jobId?: string | null; status?: string | null; name?: string | null; message?: string }
       running?: number
       pending?: number
       status?: number
@@ -118,6 +120,8 @@ async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) {
         ok: false as const,
         status: res.status,
         error: body?.error,
+        message: body?.message,
+        training: body?.training,
         running: body?.running,
         pending: body?.pending,
         killed: body?.killed
@@ -133,6 +137,15 @@ export async function fetchComfyControlStatus() {
   return await controlRequest('/status')
 }
 
+function throwIfTrainingBusy(remote: { training?: { busy?: boolean; message?: string }; error?: string; message?: string } | null) {
+  const busy = Boolean(remote?.training?.busy || remote?.error === 'train-busy')
+  if (!busy) return
+  throw createError({
+    statusCode: 409,
+    statusMessage: remote?.training?.message || remote?.message || 'AITraining is using the GPU. Stop that job before poking Comfy.'
+  })
+}
+
 export async function requestComfyWake() {
   const { controlUrl } = settings()
   if (!controlUrl) {
@@ -148,10 +161,11 @@ export async function requestComfyWake() {
       statusMessage: 'Beast host agent is unreachable. Confirm it is running on the desktop.'
     })
   }
+  throwIfTrainingBusy(remote)
   if (remote.ok === false) {
     throw createError({
       statusCode: Number(remote.status) || 502,
-      statusMessage: 'Beast host agent refused to start Comfy.'
+      statusMessage: String(remote.message || remote.error || 'Beast host agent refused to start Comfy.')
     })
   }
   return {
@@ -250,6 +264,7 @@ async function startComfy() {
   lastStartAt = now
 
   const remote = await controlRequest('/start', 'POST')
+  throwIfTrainingBusy(remote)
   if (remote?.ok || remote?.started) {
     log('start-remote', { pid: remote.pid || null })
     return { started: true, reason: 'remote' as const }
@@ -322,6 +337,8 @@ function hostWithPort(host: string, port: number) {
 async function ensureUnlocked(onStatus: StatusFn, skipBusyWait = false) {
   const cfg = settings()
   const host = cfg.host
+  const lock = await controlRequest('/status')
+  throwIfTrainingBusy(lock)
   onStatus({ state: 'offline', message: 'Checking ComfyUI...', host })
 
   const health = await checkComfyHttp(cfg.healthTimeoutMs)