aigen/server/utils/musicChain.ts

import { createJob, emitJob, type Job } from '~/server/utils/jobs'
import { extractAudio, fetchHistory, fetchHistoryAll, findHistoryAudio, queuePrompt } from '~/server/utils/comfy'
import { comfyWsUrl } from '~/server/utils/comfy'
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { downloadComfyAudio, saveTrack } from '~/server/utils/library'
import { buildMusicWorkflow, musicFilenamePrefix, assertMusicEngineNodes } from '~/server/utils/musicWorkflow'
import { NODE_LABELS } from '~/server/utils/workflow'
import { musicExtension } from '~/utils/music'
import type { MusicEngine } from '~/utils/music'

export type MusicJobParams = {
  ownerKey: string
  folderId: string
  name: string
  tags: string
  lyrics: string
  duration: number
  steps: number
  seed: number
  cfg: number
  lyricsStrength: number
  instrumental: boolean
  folderLocked?: boolean
  engine?: MusicEngine
  samplerName?: string
  scheduler?: string
}

function sleep(ms: number) {
  return new Promise(resolve => setTimeout(resolve, ms))
}

function classifyError(message: string) {
  const lower = message.toLowerCase()
  if (lower.includes('out of memory') || (lower.includes('cuda') && lower.includes('alloc')) || lower.includes('vram')) {
    return 'ComfyUI VRAM allocation failed. Try fewer steps or a shorter duration.'
  }
  if (lower.includes('econnrefused') || lower.includes('unreachable') || lower.includes('fetch failed')) {
    return 'ComfyUI host connection dropped. Confirm the desktop instance is running.'
  }
  if (lower.includes('timeout')) {
    return 'Network timeout talking to ComfyUI. The job may still be running on the desktop.'
  }
  return message
}

async function settleMusicStudio(job: Job) {
  try {
    const { onLiveVideoSettled } = await import('~/server/utils/studioQueue')
    await onLiveVideoSettled(job)
  } catch {
    /* queue kick is best-effort */
  }
}

function nodeLabel(node: string) {
  if (!node) return ''
  return NODE_LABELS[node] || ''
}

async function waitForSocket(job: Job, ms = 4000) {
  const started = Date.now()
  while (Date.now() - started < ms) {
    if (job.socketReady) return
    await sleep(100)
  }
}

function watchMusicJob(job: Job): Promise<void> {
  job.socketReady = false
  job.saving = false
  const ws = new WebSocket(comfyWsUrl(job.clientId))
  let settled = false
  let finishing = false
  let pollTimer: ReturnType<typeof setInterval> | null = null
  let timeout: ReturnType<typeof setTimeout> | null = null

  return new Promise((resolve) => {
    let localProgress = 0
    const startedAt = Date.now()
    let lastActivity = Date.now()
    const IDLE_MS = 30 * 60 * 1000
    const ABSOLUTE_MS = 3 * 60 * 60 * 1000

    const markActivity = () => { lastActivity = Date.now() }

    const emitLocal = (event: Parameters<typeof emitJob>[1], samplePct?: number) => {
      markActivity()
      if (typeof event.progress === 'number') localProgress = event.progress
      if (event.message && typeof samplePct === 'number') {
        event.message = `${event.message} (${samplePct}%)`
      }
      emitJob(job, event)
    }

    const cleanup = () => {
      if (timeout) clearTimeout(timeout)
      if (pollTimer) clearInterval(pollTimer)
      timeout = null
      pollTimer = null
      try { ws.close() } catch { /* ignore */ }
    }

    const armIdle = () => {
      if (timeout) clearTimeout(timeout)
      if (settled) return
      const absLeft = ABSOLUTE_MS - (Date.now() - startedAt)
      if (absLeft <= 0) {
        void fail('Timed out waiting for ComfyUI (3 hours).')
        return
      }
      const idleLeft = IDLE_MS - (Date.now() - lastActivity)
      timeout = setTimeout(() => {
        if (settled || finishing) return
        if (Date.now() - lastActivity >= IDLE_MS) {
          void fail('Timed out waiting for ComfyUI progress (30 minutes with no updates).')
          return
        }
        armIdle()
      }, Math.max(1000, Math.min(idleLeft, absLeft)))
    }

    const fail = async (error: string) => {
      if (settled || finishing) return
      settled = true
      cleanup()
      job.status = job.status === 'cancelled' ? 'cancelled' : 'error'
      job.error = classifyError(error)
      emitJob(job, { type: 'error', error: job.error, message: job.error })
      await settleMusicStudio(job)
      resolve()
    }

    const adoptPromptId = (id: unknown) => {
      const promptId = String(id || '').trim()
      if (promptId && !job.promptId) job.promptId = promptId
    }

    const loadJobAudio = async () => {
      if (job.promptId) {
        const scoped = await fetchHistory(job.promptId)
        const fromPrompt = extractAudio(scoped, job.promptId)
        if (fromPrompt) return fromPrompt
      }
      const found = findHistoryAudio(await fetchHistoryAll(), {
        promptId: job.promptId,
        tags: job.library?.tags,
        since: job.startedAt - 60_000
      })
      if (found) {
        job.promptId = found.promptId
        return found.audio
      }
      return null
    }

    const succeed = async () => {
      if (settled || finishing) return false
      finishing = true
      job.saving = true
      try {
        const audio = await loadJobAudio()
        if (!audio) {
          finishing = false
          job.saving = false
          return false
        }
        if (settled) return false
        settled = true
        cleanup()
        try {
          job.audio = audio
          emitLocal({ type: 'status', message: 'Saving to library...', progress: 98 })
          if (job.library) {
            const buffer = await downloadComfyAudio(audio)
            const ext = musicExtension(audio.filename)
            const track = await saveTrack({
              ownerKey: job.library.ownerKey,
              folderId: job.library.folderId,
              name: job.library.name,
              tags: job.library.tags || job.library.prompt,
              lyrics: job.library.lyrics || '',
              duration: job.library.duration || 0,
              steps: job.library.steps,
              seed: job.library.seed,
              cfg: job.library.cfg || 0,
              instrumental: job.library.instrumental === true,
              audio: buffer,
              ext,
              engine: job.library.engine,
              samplerName: job.library.samplerName,
              scheduler: job.library.scheduler,
              lyricsStrength: job.library.lyricsStrength,
              comfyFilename: audio.filename
            })
            job.trackId = track.id
            job.library.audioExt = ext
          }
          job.status = 'complete'
          emitJob(job, {
            type: 'complete',
            message: job.library?.folderLocked
              ? 'Saved to the locked folder. Unlock it to listen.'
              : 'Track ready',
            progress: 100,
            filename: job.library?.folderLocked ? undefined : audio.filename,
            subfolder: job.library?.folderLocked ? undefined : audio.subfolder,
            mediaType: job.library?.folderLocked ? undefined : audio.type,
            trackId: job.trackId,
            folderLocked: job.library?.folderLocked
          })
        } catch (saveError) {
          const message = saveError instanceof Error ? saveError.message : String(saveError)
          job.status = 'error'
          job.error = `Audio generated but library save failed: ${message}`
          emitJob(job, { type: 'error', error: job.error, message: job.error })
        }
        await settleMusicStudio(job)
        job.saving = false
        resolve()
        return true
      } catch (error) {
        finishing = false
        job.saving = false
        throw error
      }
    }

    const pollHistory = async () => {
      if (settled || finishing) return
      try {
        const history = await fetchHistoryAll()
        const found = findHistoryAudio(history, {
          promptId: job.promptId,
          tags: job.library?.tags,
          since: job.startedAt - 60_000
        })
        if (found) {
          job.promptId = found.promptId
          await succeed()
          return
        }
        const entry = (job.promptId ? history?.[job.promptId] : undefined) as {
          status?: { status_str?: string; completed?: boolean }
        } | undefined
        const status = entry?.status?.status_str
        if (status === 'error') await fail('ComfyUI reported an execution error')
        else if (status === 'interrupted') await fail('Job interrupted.')
      } catch {
        // history can miss while the graph is still running
      }
    }

    armIdle()
    pollTimer = setInterval(() => { void pollHistory() }, 4000)

    ws.addEventListener('open', () => {
      job.socketReady = true
      emitLocal({ type: 'status', message: 'Connected to ComfyUI', progress: Math.max(localProgress, 4) })
    })
    ws.addEventListener('error', () => {
      job.socketReady = false
    })
    ws.addEventListener('message', (event) => {
      let payload: { type?: string; data?: Record<string, any> }
      try {
        payload = JSON.parse(String(event.data))
      } catch {
        return
      }
      const type = String(payload.type || '')
      const data = payload.data || {}
      adoptPromptId(data.prompt_id)
      if (type === 'status') markActivity()
      if (type === 'progress') {
        const value = Number(data.value || 0)
        const max = Number(data.max || 0)
        const pct = max > 0 ? Math.min(95, Math.round((value / max) * 90) + 8) : Math.max(localProgress, 12)
        const node = String(data.node || '')
        emitLocal({
          type: 'progress',
          message: nodeLabel(node) || 'Sampling',
          progress: pct,
          step: value,
          maxStep: max,
          node
        }, max > 0 ? Math.round((value / max) * 100) : undefined)
      }
      if (type === 'executing') {
        const node = data.node == null ? '' : String(data.node)
        if (!node) void succeed()
        else if (node) {
          emitLocal({
            type: 'status',
            message: nodeLabel(node) || 'Running node',
            progress: Math.max(localProgress, 10),
            node
          })
        }
      }
      if (type === 'executed') {
        void succeed()
      }
      if (type === 'execution_error') {
        const message = String(data.exception_message || data.message || 'ComfyUI node execution failed')
        void fail(message)
      }
      if (type === 'execution_interrupted') {
        job.status = 'cancelled'
        void fail('Job interrupted.')
      }
    })
  })
}

export async function startMusicJob(params: MusicJobParams) {
  const job = createJob('music')
  job.library = {
    ownerKey: params.ownerKey,
    folderId: params.folderId,
    hideThumbnail: false,
    folderLocked: params.folderLocked === true,
    name: params.name,
    prompt: params.tags,
    tags: params.tags,
    lyrics: params.lyrics,
    instrumental: params.instrumental,
    aspect: 'audio',
    width: 0,
    height: 0,
    steps: params.steps,
    turbo: false,
    seed: params.seed,
    cfg: params.cfg,
    duration: params.duration,
    sound: true,
    engine: params.engine || 'ace-step',
    samplerName: params.samplerName || 'euler',
    scheduler: params.scheduler || 'simple',
    lyricsStrength: params.lyricsStrength
  }
  job.maxStep = params.steps
  job.status = 'queued'
  emitJob(job, { type: 'status', message: 'Checking ComfyUI...', progress: 1 })

  void (async () => {
    try {
      const done = watchMusicJob(job)
      await ensureComfyReady((status) => {
        emitJob(job, {
          type: status.state === 'busy' ? 'busy' : 'status',
          message: status.message,
          progress: status.state === 'online' ? 3 : 1,
          busy: status.state === 'busy',
          queueRunning: status.queueRunning,
          queuePending: status.queuePending
        })
      })
      job.status = 'running'
      const engineLabel = params.engine === 'yue' ? 'YuE' : params.engine === 'ace-step-1.5' ? 'ACE-Step 1.5' : 'ACE-Step'
      emitJob(job, { type: 'status', message: `Queueing ${engineLabel}...`, progress: 6 })
      await waitForSocket(job, 4000)
      await assertMusicEngineNodes(params.engine)
      const graph = buildMusicWorkflow({
        tags: params.tags,
        lyrics: params.instrumental ? '' : params.lyrics,
        duration: params.duration,
        steps: params.steps,
        seed: params.seed,
        cfg: params.cfg,
        lyricsStrength: params.instrumental ? 0 : params.lyricsStrength,
        samplerName: params.samplerName,
        scheduler: params.scheduler,
        filenamePrefix: musicFilenamePrefix(),
        engine: params.engine,
        vocalMultiplier: params.instrumental ? 0.2 : 1
      })
      try {
        const queued = await queuePrompt(graph, job.clientId, { timeoutMs: 25_000 })
        if (queued?.prompt_id) job.promptId = queued.prompt_id
        emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
      } catch (queueError: unknown) {
        const statusCode = Number((queueError as { statusCode?: number; status?: number })?.statusCode
          || (queueError as { status?: number })?.status || 0)
        if (statusCode >= 400) throw queueError
        emitJob(job, {
          type: 'status',
          message: 'Waiting for ComfyUI to pick up the job...',
          progress: 8
        })
      }
      await done
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error)
      if (job.status !== 'cancelled' && job.status !== 'complete') {
        job.status = 'error'
        job.error = classifyError(message)
        emitJob(job, { type: 'error', error: job.error, message: job.error })
        await settleMusicStudio(job)
      }
    }
  })()

  return job
}