aigen/server/api/generate/active.get.ts

import { fetchLiveQueue } from '~/server/utils/comfy'
import { jobSnapshot, listJobs, type Job } from '~/server/utils/jobs'
import { libraryOwnerKey } from '~/server/utils/library'
import { listPendingJobs, resolvePendingJob } from '~/server/utils/pending'
import { jobIsLocallySubmitting } from '~/server/utils/studioQueue'
import { ensurePendingWatch } from '~/server/utils/watch'

const SUBMIT_WINDOW_MS = 45 * 1000

function livePayload(job: Job) {
  return {
    jobId: job.id,
    kind: job.kind || 'video',
    hideThumbnail: job.hideThumbnail === true,
    folderLocked: job.library?.folderLocked === true,
    ...jobSnapshot(job)
  }
}

function jobLooksLive(job: Job, comfyBusy: boolean) {
  if (job.library?.stopAfterCurrent) return false
  if (jobIsLocallySubmitting(job)) return true
  if (job.status === 'running' && job.promptId) return comfyBusy
  if (job.status === 'queued' && Date.now() - job.startedAt < SUBMIT_WINDOW_MS) return true
  return false
}

export default defineEventHandler(async (event) => {
  const owner = libraryOwnerKey(event)
  const queue = await fetchLiveQueue()
  const comfyBusy = queue ? queue.running > 0 || queue.pending > 0 : true
  const live = listJobs()
    .filter(job => job.library?.ownerKey === owner && jobLooksLive(job, comfyBusy))
    .sort((a, b) => b.startedAt - a.startedAt)
  const videoLive = live.find(job => job.kind !== 'edit' && job.kind !== 'music')
  const editLive = live.find(job => job.kind === 'edit')
  const musicLive = live.find(job => job.kind === 'music')

  let video = videoLive ? livePayload(videoLive) : null
  const edit = editLive ? livePayload(editLive) : null
  const music = musicLive ? livePayload(musicLive) : null

  if (!video) {
    const pending = listPendingJobs()
      .filter(job => (
        job.ownerKey === owner
        && !job.stopAfterCurrent
        && Date.now() - job.startedAt < SUBMIT_WINDOW_MS * 4
      ))
      .sort((a, b) => b.startedAt - a.startedAt)[0]
    if (pending && comfyBusy) {
      const resolved = await resolvePendingJob(pending).catch(() => null)
      if (resolved && resolved.status !== 'complete' && resolved.status !== 'error') {
        video = {
          jobId: pending.jobId,
          kind: 'video' as const,
          hideThumbnail: pending.hideThumbnail === true,
          folderLocked: pending.folderLocked === true,
          ...resolved
        }
      } else if (!resolved) {
        const job = ensurePendingWatch(pending)
        video = {
          jobId: pending.jobId,
          kind: 'video',
          hideThumbnail: pending.hideThumbnail === true,
          folderLocked: pending.folderLocked === true,
          ...jobSnapshot(job)
        }
      }
    }
  }

  return {
    jobId: video?.jobId || edit?.jobId || null,
    video,
    edit,
    music
  }
})